LFS export records repository_type data

A project can have the same `LfsObject` linked with up to three
`LfsObjectsProject` records. Each of these records would be for a
different repository, recorded in the `repository_type` property. The
different repositories at time of writing are "project", "wiki", and
"design". See https://gitlab.com/gitlab-org/gitlab-ee/merge_requests/13894

This change exports the list of `repository_type`s as a JSON mapping of
oid => repository_types, which are imported to recreate the correct
`LfsObjectsProject` records.

https://gitlab.com/gitlab-org/gitlab-ee/issues/11090

2 files changed, 91 insertions, 7 deletions

diff --git a/lib/gitlab/import_export/lfs_restorer.rb b/lib/gitlab/import_export/lfs_restorer.rb
index 345c7880e30..1de8a5bf9ec 100644
--- a/lib/gitlab/import_export/lfs_restorer.rb
+++ b/lib/gitlab/import_export/lfs_restorer.rb
@@ -3,6 +3,10 @@
 module Gitlab
   module ImportExport
     class LfsRestorer
+      include Gitlab::Utils::StrongMemoize
+
+      attr_accessor :project, :shared
+
       def initialize(project:, shared:)
         @project = project
         @shared = shared
@@ -17,7 +21,7 @@ module Gitlab
 
         true
       rescue => e
-        @shared.error(e)
+        shared.error(e)
         false
       end
 
@@ -29,16 +33,57 @@ module Gitlab
 
         lfs_object = LfsObject.find_or_initialize_by(oid: oid, size: size)
         lfs_object.file = File.open(path) unless lfs_object.file&.exists?
+        lfs_object.save! if lfs_object.changed?
 
-        @project.all_lfs_objects << lfs_object
+        repository_types(oid).each do |repository_type|
+          LfsObjectsProject.create!(
+            project: project,
+            lfs_object: lfs_object,
+            repository_type: repository_type
+          )
+        end
+      end
+
+      def repository_types(oid)
+        # We allow support for imports created before the `lfs-objects.json`
+        # file was generated. In this case, the restorer will link an LFS object
+        # with a single `lfs_objects_projects` relation.
+        #
+        # This allows us backwards-compatibility without version bumping.
+        # See https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/30830#note_192608870
+        return ['project'] unless has_lfs_json?
+
+        lfs_json[oid]
       end
 
       def lfs_file_paths
         @lfs_file_paths ||= Dir.glob("#{lfs_storage_path}/*")
       end
 
+      def has_lfs_json?
+        strong_memoize(:has_lfs_json) do
+          File.exist?(lfs_json_path)
+        end
+      end
+
+      def lfs_json
+        return {} unless has_lfs_json?
+
+        @lfs_json ||=
+          begin
+            json = IO.read(lfs_json_path)
+            ActiveSupport::JSON.decode(json)
+          rescue
+            raise Gitlab::ImportExport::Error.new('Incorrect JSON format')
+          end
+      end
+
       def lfs_storage_path
-        File.join(@shared.export_path, 'lfs-objects')
+        File.join(shared.export_path, ImportExport.lfs_objects_storage)
+      end
+
+      def lfs_json_path
+        File.join(shared.export_path, ImportExport.lfs_objects_filename)
       end
     end
   end
diff --git a/lib/gitlab/import_export/lfs_saver.rb b/lib/gitlab/import_export/lfs_saver.rb
index 954f6f00078..18c590e1ca9 100644
--- a/lib/gitlab/import_export/lfs_saver.rb
+++ b/lib/gitlab/import_export/lfs_saver.rb
@@ -5,25 +5,40 @@ module Gitlab
     class LfsSaver
       include Gitlab::ImportExport::CommandLineUtil
 
+      attr_accessor :lfs_json, :project, :shared
+
+      BATCH_SIZE = 100
+
       def initialize(project:, shared:)
         @project = project
         @shared = shared
+        @lfs_json = {}
       end
 
       def save
-        @project.all_lfs_objects.each do |lfs_object|
-          save_lfs_object(lfs_object)
+        project.all_lfs_objects.find_in_batches(batch_size: BATCH_SIZE) do |batch|
+          batch.each do |lfs_object|
+            save_lfs_object(lfs_object)
+          end
+
+          append_lfs_json_for_batch(batch) if write_lfs_json_enabled?
         end
 
+        write_lfs_json if write_lfs_json_enabled?
+
         true
       rescue => e
-        @shared.error(e)
+        shared.error(e)
 
         false
       end
 
       private
 
+      def write_lfs_json_enabled?
+        ::Feature.enabled?(:export_lfs_objects_projects, default_enabled: true)
+      end
+
       def save_lfs_object(lfs_object)
         if lfs_object.local_store?
           copy_file_for_lfs_object(lfs_object)
@@ -45,12 +60,36 @@ module Gitlab
         copy_files(lfs_object.file.path, destination_path_for_object(lfs_object))
       end
 
+      def append_lfs_json_for_batch(lfs_objects_batch)
+        lfs_objects_projects = LfsObjectsProject
+                                .select('lfs_objects.oid, array_agg(distinct lfs_objects_projects.repository_type) as repository_types')
+                                .joins(:lfs_object)
+                                .where(project: project, lfs_object: lfs_objects_batch)
+                                .group('lfs_objects.oid')
+
+        lfs_objects_projects.each do |group|
+          oid = group.oid
+
+          lfs_json[oid] ||= []
+          lfs_json[oid] += group.repository_types
+        end
+      end
+
+      def write_lfs_json
+        mkdir_p(shared.export_path)
+        File.write(lfs_json_path, lfs_json.to_json)
+      end
+
       def destination_path_for_object(lfs_object)
         File.join(lfs_export_path, lfs_object.oid)
       end
 
       def lfs_export_path
-        File.join(@shared.export_path, 'lfs-objects')
+        File.join(shared.export_path, ImportExport.lfs_objects_storage)
+      end
+
+      def lfs_json_path
+        File.join(shared.export_path, ImportExport.lfs_objects_filename)
       end
     end
   end