Merge branch 'georgekoltsov/poc-import-update-object-persistance' into 'master'

Update Import process to persist object subrelations in separate db transactions

See merge request gitlab-org/gitlab!79963

config/feature_flags/development/import_relation_object_persistence.yml

+---
+name: import_relation_object_persistence
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/79963
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/354497
+milestone: '14.9'
+type: development
+group: group::import
+default_enabled: false

lib/gitlab/import_export/base/relation_object_saver.rb

+# frozen_string_literal: true
+
+# RelationObjectSaver allows for an alternative approach to persisting
+# objects during Project/Group Import which persists object's
+# nested collection subrelations separately, in batches.
+#
+# Instead of the regular `relation_object.save!` that opens one db
+# transaction for the object itself and all of its subrelations we
+# separate collection subrelations from the object and save them
+# in batches in smaller more frequent db transactions.
+module Gitlab
+  module ImportExport
+    module Base
+      class RelationObjectSaver
+        include Gitlab::Utils::StrongMemoize
+
+        BATCH_SIZE = 100
+        MIN_RECORDS_SIZE = 5
+
+        # @param relation_object [Object] Object of a project/group, e.g. an issue
+        # @param relation_key [String] Name of the object association to group/project, e.g. :issues
+        # @param relation_definition [Hash] Object subrelations as defined in import_export.yml
+        # @param importable [Project|Group] Project or group where relation object is getting saved to
+        #
+        # @example
+        #   Gitlab::ImportExport::Base::RelationObjectSaver.new(
+        #     relation_key: 'merge_requests',
+        #     relation_object: #<MergeRequest id: root/mrs!1, notes: [#<Note id: nil, note: 'test', ...>, #<Note id: nil, noteL 'another note'>]>,
+        #     relation_definition: {"metrics"=>{}, "award_emoji"=>{}, "notes"=>{"author"=>{}, ... }}
+        #     importable: @importable
+        #   ).execute
+        def initialize(relation_object:, relation_key:, relation_definition:, importable:)
+          @relation_object = relation_object
+          @relation_key = relation_key
+          @relation_definition = relation_definition
+          @importable = importable
+          @invalid_subrelations = []
+        end
+
+        def execute
+          move_subrelations
+
+          relation_object.save!
+
+          save_subrelations
+        ensure
+          log_invalid_subrelations
+        end
+
+        private
+
+        attr_reader :relation_object, :relation_key, :relation_definition,
+                    :importable, :collection_subrelations, :invalid_subrelations
+
+        # rubocop:disable GitlabSecurity/PublicSend
+        def save_subrelations
+          collection_subrelations.each_pair do |relation_name, records|
+            records.each_slice(BATCH_SIZE) do |batch|
+              valid_records, invalid_records = batch.partition { |record| record.valid? }
+
+              invalid_subrelations << invalid_records
+              relation_object.public_send(relation_name) << valid_records
+            end
+          end
+        end
+
+        def move_subrelations
+          strong_memoize(:collection_subrelations) do
+            relation_definition.each_key.each_with_object({}) do |definition, collection_subrelations|
+              subrelation = relation_object.public_send(definition)
+              association = relation_object.class.reflect_on_association(definition)
+
+              if association&.collection? && subrelation.size > MIN_RECORDS_SIZE
+                collection_subrelations[definition] = subrelation.records
+
+                subrelation.clear
+              end
+            end
+          end
+        end
+        # rubocop:enable GitlabSecurity/PublicSend
+
+        def log_invalid_subrelations
+          invalid_subrelations.flatten.each do |record|
+            Gitlab::Import::Logger.info(
+              message: '[Project/Group Import] Invalid subrelation',
+              importable_column_name => importable.id,
+              relation_key: relation_key,
+              error_messages: record.errors.full_messages.to_sentence
+            )
+
+            ImportFailure.create(
+              source: 'RelationObjectSaver#save!',
+              relation_key: relation_key,
+              exception_class: 'RecordInvalid',
+              exception_message: record.errors.full_messages.to_sentence,
+              correlation_id_value: Labkit::Correlation::CorrelationId.current_or_new_id,
+              importable_column_name => importable.id
+            )
+          end
+        end
+
+        def importable_column_name
+          @column_name ||= importable.class.reflect_on_association(:import_failures).foreign_key.to_sym
+        end
+      end
+    end
+  end
+end

lib/gitlab/import_export/group/relation_tree_restorer.rb

@@ -79,10 +79,7 @@ module Gitlab
 
           relation_object.assign_attributes(importable_class_sym => @importable)
 
-          import_failure_service.with_retry(action: 'relation_object.save!', relation_key: relation_key, relation_index: relation_index) do
-            relation_object.save!
-            log_relation_creation(@importable, relation_key, relation_object)
-          end
+          save_relation_object(relation_object, relation_key, relation_definition, relation_index)
         rescue StandardError => e
           import_failure_service.log_import_failure(
             source: 'process_relation_item!',
@@ -91,6 +88,23 @@ module Gitlab
             exception: e)
         end
 
+        def save_relation_object(relation_object, relation_key, relation_definition, relation_index)
+          if Feature.enabled?(:import_relation_object_persistence, default_enabled: :yaml) && relation_object.new_record?
+            Gitlab::ImportExport::Base::RelationObjectSaver.new(
+              relation_object: relation_object,
+              relation_key: relation_key,
+              relation_definition: relation_definition,
+              importable: @importable
+            ).execute
+          else
+            import_failure_service.with_retry(action: 'relation_object.save!', relation_key: relation_key, relation_index: relation_index) do
+              relation_object.save!
+            end
+          end
+
+          log_relation_creation(@importable, relation_key, relation_object)
+        end
+
         def import_failure_service
           @import_failure_service ||= ImportFailureService.new(@importable)
         end

spec/lib/gitlab/import_export/base/relation_object_saver_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::ImportExport::Base::RelationObjectSaver do
+  let(:project) { create(:project) }
+  let(:relation_object) { build(:issue, project: project) }
+  let(:relation_definition) { {} }
+  let(:importable) { project }
+  let(:relation_key) { 'issues' }
+
+  subject(:saver) do
+    described_class.new(
+      relation_object: relation_object,
+      relation_key: relation_key,
+      relation_definition: relation_definition,
+      importable: importable
+    )
+  end
+
+  describe '#save' do
+    before do
+      expect(relation_object).to receive(:save!).and_call_original
+    end
+
+    it 'saves relation object' do
+      expect { saver.execute }.to change(project.issues, :count).by(1)
+    end
+
+    context 'when subrelation is present' do
+      let(:notes) { build_list(:note, 6, project: project, importing: true) }
+      let(:relation_object) { build(:issue, project: project, notes: notes) }
+      let(:relation_definition) { { 'notes' => {} } }
+
+      it 'saves relation object with subrelations' do
+        expect(relation_object.notes).to receive(:<<).and_call_original
+
+        saver.execute
+
+        issue = project.issues.last
+        expect(issue.notes.count).to eq(6)
+      end
+    end
+
+    context 'when subrelation is not a collection' do
+      let(:sentry_issue) { build(:sentry_issue, importing: true) }
+      let(:relation_object) { build(:issue, project: project, sentry_issue: sentry_issue) }
+      let(:relation_definition) { { 'sentry_issue' => {} } }
+
+      it 'saves subrelation as part of the relation object itself' do
+        expect(relation_object.notes).not_to receive(:<<)
+
+        saver.execute
+
+        issue = project.issues.last
+        expect(issue.sentry_issue.persisted?).to eq(true)
+      end
+    end
+
+    context 'when subrelation collection count is small' do
+      let(:notes) { build_list(:note, 2, project: project, importing: true) }
+      let(:relation_object) { build(:issue, project: project, notes: notes) }
+      let(:relation_definition) { { 'notes' => {} } }
+
+      it 'saves subrelation as part of the relation object itself' do
+        expect(relation_object.notes).not_to receive(:<<)
+
+        saver.execute
+
+        issue = project.issues.last
+        expect(issue.notes.count).to eq(2)
+      end
+    end
+
+    context 'when some subrelations are invalid' do
+      let(:notes) { build_list(:note, 5, project: project, importing: true) }
+      let(:invalid_note) { build(:note) }
+      let(:relation_object) { build(:issue, project: project, notes: notes + [invalid_note]) }
+      let(:relation_definition) { { 'notes' => {} } }
+
+      it 'saves valid subrelations and logs invalid subrelation' do
+        expect(relation_object.notes).to receive(:<<).and_call_original
+        expect(Gitlab::Import::Logger)
+          .to receive(:info)
+          .with(
+            message: '[Project/Group Import] Invalid subrelation',
+            project_id: project.id,
+            relation_key: 'issues',
+            error_messages: "Noteable can't be blank and Project does not match noteable project"
+          )
+
+        saver.execute
+
+        issue = project.issues.last
+        import_failure = project.import_failures.last
+
+        expect(issue.notes.count).to eq(5)
+        expect(import_failure.source).to eq('RelationObjectSaver#save!')
+        expect(import_failure.exception_message).to eq("Noteable can't be blank and Project does not match noteable project")
+      end
+
+      context 'when importable is group' do
+        let(:relation_key) { 'labels' }
+        let(:relation_definition) { { 'priorities' => {} } }
+        let(:importable) { create(:group) }
+        let(:valid_priorities) { build_list(:label_priority, 5, importing: true) }
+        let(:invalid_priority) { build(:label_priority, priority: -1) }
+        let(:relation_object) { build(:group_label, group: importable, title: 'test', priorities: valid_priorities + [invalid_priority]) }
+
+        it 'logs invalid subrelation for a group' do
+          expect(Gitlab::Import::Logger)
+            .to receive(:info)
+            .with(
+              message: '[Project/Group Import] Invalid subrelation',
+              group_id: importable.id,
+              relation_key: 'labels',
+              error_messages: 'Priority must be greater than or equal to 0'
+            )
+
+          saver.execute
+
+          label = importable.labels.last
+          import_failure = importable.import_failures.last
+
+          expect(label.priorities.count).to eq(5)
+          expect(import_failure.source).to eq('RelationObjectSaver#save!')
+          expect(import_failure.exception_message).to eq('Priority must be greater than or equal to 0')
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/import_export/group/tree_restorer_spec.rb

@@ -5,6 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::ImportExport::Group::TreeRestorer do
   include ImportExport::CommonUtil
 
+  shared_examples 'group restoration' do
     describe 'restore group tree' do
       before_all do
         # Using an admin for import, so we can check assignment of existing members
@@ -201,4 +202,21 @@ RSpec.describe Gitlab::ImportExport::Group::TreeRestorer do
       include_examples 'with visibility level', :private, [0, 0, 0]
       include_examples 'with visibility level', :internal, [10, 10, 0]
     end
+  end
+
+  context 'when import_relation_object_persistence feature flag is enabled' do
+    before do
+      stub_feature_flags(import_relation_object_persistence: true)
+    end
+
+    include_examples 'group restoration'
+  end
+
+  context 'when import_relation_object_persistence feature flag is disabled' do
+    before do
+      stub_feature_flags(import_relation_object_persistence: false)
+    end
+
+    include_examples 'group restoration'
+  end
 end

spec/lib/gitlab/import_export/project/tree_restorer_spec.rb

@@ -1058,6 +1058,27 @@ RSpec.describe Gitlab::ImportExport::Project::TreeRestorer do
     end
   end
 
+  context 'when import_relation_object_persistence feature flag is enabled' do
+    before do
+      stub_feature_flags(import_relation_object_persistence: true)
+    end
+
+    context 'enable ndjson import' do
+      it_behaves_like 'project tree restorer work properly', :legacy_reader, true
+
+      it_behaves_like 'project tree restorer work properly', :ndjson_reader, true
+    end
+
+    context 'disable ndjson import' do
+      it_behaves_like 'project tree restorer work properly', :legacy_reader, false
+    end
+  end
+
+  context 'when import_relation_object_persistence feature flag is disabled' do
+    before do
+      stub_feature_flags(import_relation_object_persistence: false)
+    end
+
     context 'enable ndjson import' do
       it_behaves_like 'project tree restorer work properly', :legacy_reader, true
 
@@ -1067,4 +1088,5 @@ RSpec.describe Gitlab::ImportExport::Project::TreeRestorer do
     context 'disable ndjson import' do
       it_behaves_like 'project tree restorer work properly', :legacy_reader, false
     end
+  end
 end