Merge branch 'master' into 'backstage/gb/use-persisted-stages-to-improve-pipelines-table'

# Conflicts:
#   db/schema.rb

2 files changed, 51 insertions, 0 deletions

diff --git a/app/models/concerns/avatarable.rb b/app/models/concerns/avatarable.rb
index 13246a774e3..095897b08e3 100644
--- a/app/models/concerns/avatarable.rb
+++ b/app/models/concerns/avatarable.rb
@@ -4,11 +4,14 @@ module Avatarable
   included do
     prepend ShadowMethods
     include ObjectStorage::BackgroundMove
+    include Gitlab::Utils::StrongMemoize
 
     validate :avatar_type, if: ->(user) { user.avatar.present? && user.avatar_changed? }
     validates :avatar, file_size: { maximum: 200.kilobytes.to_i }
 
     mount_uploader :avatar, AvatarUploader
+
+    after_initialize :add_avatar_to_batch
   end
 
   module ShadowMethods
@@ -18,6 +21,17 @@ module Avatarable
 
       avatar_path(only_path: args.fetch(:only_path, true)) || super
     end
+
+    def retrieve_upload(identifier, paths)
+      upload = retrieve_upload_from_batch(identifier)
+
+      # This fallback is needed when deleting an upload, because we may have
+      # already been removed from the DB. We have to check an explicit `#nil?`
+      # because it's a BatchLoader instance.
+      upload = super if upload.nil?
+
+      upload
+    end
   end
 
   def avatar_type
@@ -52,4 +66,37 @@ module Avatarable
 
     url_base + avatar.local_url
   end
+
+  # Path that is persisted in the tracking Upload model. Used to fetch the
+  # upload from the model.
+  def upload_paths(identifier)
+    avatar_mounter.blank_uploader.store_dirs.map { |store, path| File.join(path, identifier) }
+  end
+
+  private
+
+  def retrieve_upload_from_batch(identifier)
+    BatchLoader.for(identifier: identifier, model: self).batch(key: self.class) do |upload_params, loader, args|
+      model_class = args[:key]
+      paths = upload_params.flat_map do |params|
+        params[:model].upload_paths(params[:identifier])
+      end
+
+      Upload.where(uploader: AvatarUploader, path: paths).find_each do |upload|
+        model = model_class.instantiate('id' => upload.model_id)
+
+        loader.call({ model: model, identifier: File.basename(upload.path) }, upload)
+      end
+    end
+  end
+
+  def add_avatar_to_batch
+    return unless avatar_mounter
+
+    avatar_mounter.read_identifiers.each { |identifier| retrieve_upload_from_batch(identifier) }
+  end
+
+  def avatar_mounter
+    strong_memoize(:avatar_mounter) { _mounter(:avatar) }
+  end
 end
diff --git a/app/models/concerns/with_uploads.rb b/app/models/concerns/with_uploads.rb
index e7cfffb775b..4245d083a49 100644
--- a/app/models/concerns/with_uploads.rb
+++ b/app/models/concerns/with_uploads.rb
@@ -36,4 +36,8 @@ module WithUploads
       upload.destroy
     end
   end
+
+  def retrieve_upload(_identifier, paths)
+    uploads.find_by(path: paths)
+  end
 end