Specify if `iid` should be generated during Import

Refactor `InternalId` generation.
Remove a hack to handle ci_pipelines.iid for imports.

app/models/ci/pipeline.rb

@@ -25,7 +25,7 @@ module Ci
     belongs_to :merge_request, class_name: 'MergeRequest'
     belongs_to :external_pull_request
 
-    has_internal_id :iid, scope: :project, presence: false, init: ->(s) do
+    has_internal_id :iid, scope: :project, presence: false, ensure_if: -> { !importing? }, init: ->(s) do
       s&.project&.all_pipelines&.maximum(:iid) || s&.project&.all_pipelines&.count
     end
 

app/models/concerns/atomic_internal_id.rb

@@ -27,53 +27,74 @@ module AtomicInternalId
   extend ActiveSupport::Concern
 
   class_methods do
-    def has_internal_id(column, scope:, init:, presence: true) # rubocop:disable Naming/PredicateName
+    def has_internal_id(column, scope:, init:, ensure_if: nil, presence: true) # rubocop:disable Naming/PredicateName
       # We require init here to retain the ability to recalculate in the absence of a
       # InternaLId record (we may delete records in `internal_ids` for example).
       raise "has_internal_id requires a init block, none given." unless init
+      raise "has_internal_id needs to be defined on association." unless self.reflect_on_association(scope)
 
-      before_validation :"ensure_#{scope}_#{column}!", on: :create
+      before_validation :"track_#{scope}_#{column}!", on: :create
+      before_validation :"ensure_#{scope}_#{column}!", on: :create, if: ensure_if
       validates column, presence: presence
 
       define_method("ensure_#{scope}_#{column}!") do
-        scope_value = association(scope).reader
+        scope_value = internal_id_read_scope(scope)
         value = read_attribute(column)
-
         return value unless scope_value
 
-        scope_attrs = { scope_value.class.table_name.singularize.to_sym => scope_value }
-        usage = self.class.table_name.to_sym
-
-        if value.present? && (@iid_needs_tracking || Feature.enabled?(:iid_always_track, default_enabled: true))
-          # The value was set externally, e.g. by the user
-          # We update the InternalId record to keep track of the greatest value.
-          InternalId.track_greatest(self, scope_attrs, usage, value, init)
-
-          @iid_needs_tracking = false
-        elsif !value.present?
+        if value.nil?
           # We don't have a value yet and use a InternalId record to generate
           # the next value.
-          value = InternalId.generate_next(self, scope_attrs, usage, init)
+          value = InternalId.generate_next(
+            self,
+            internal_id_scope_attrs(scope),
+            internal_id_scope_usage,
+            init)
           write_attribute(column, value)
         end
 
         value
       end
 
+      define_method("track_#{scope}_#{column}!") do
+        if !@internal_id_needs_tracking
+          return unless Feature.enabled?(:iid_always_track, default_enabled: true)
+        end
+
+        @internal_id_needs_tracking = false
+
+        scope_value = internal_id_read_scope(scope)
+        value = read_attribute(column)
+        return unless scope_value
+
+        if value.present?
+          # The value was set externally, e.g. by the user
+          # We update the InternalId record to keep track of the greatest value.
+          InternalId.track_greatest(
+            self,
+            internal_id_scope_attrs(scope),
+            internal_id_scope_usage,
+            value,
+            init)
+        end
+      end
+
       define_method("#{column}=") do |value|
         super(value).tap do |v|
           # Indicate the iid was set from externally
-          @iid_needs_tracking = true
+          @internal_id_needs_tracking = true
         end
       end
 
       define_method("reset_#{scope}_#{column}") do
         if value = read_attribute(column)
-          scope_value = association(scope).reader
-          scope_attrs = { scope_value.class.table_name.singularize.to_sym => scope_value }
-          usage = self.class.table_name.to_sym
+          did_reset = InternalId.reset(
+            self,
+            internal_id_scope_attrs(scope),
+            internal_id_scope_usage,
+            value)
 
-          if InternalId.reset(self, scope_attrs, usage, value)
+          if did_reset
             write_attribute(column, nil)
           end
         end
@@ -82,4 +103,18 @@ module AtomicInternalId
       end
     end
   end
+
+  def internal_id_scope_attrs(scope)
+    scope_value = internal_id_read_scope(scope)
+
+    { scope_value.class.table_name.singularize.to_sym => scope_value } if scope_value
+  end
+
+  def internal_id_scope_usage
+    self.class.table_name.to_sym
+  end
+
+  def internal_id_read_scope(scope)
+    association(scope).reader
+  end
 end

lib/gitlab/import_export/project_tree_restorer.rb

@@ -184,18 +184,9 @@ module Gitlab
         return if tree_hash[relation_key].blank?
 
         tree_array = [tree_hash[relation_key]].flatten
-        null_iid_pipelines = []
 
         # Avoid keeping a possible heavy object in memory once we are done with it
-        while relation_item = (tree_array.shift || null_iid_pipelines.shift)
-          if nil_iid_pipeline?(relation_key, relation_item) && tree_array.any?
-            # Move pipelines with NULL IIDs to the end
-            # so they don't clash with existing IIDs.
-            null_iid_pipelines << relation_item
-
-            next
-          end
-
+        while relation_item = tree_array.shift
           remove_feature_dependent_sub_relations(relation_item)
 
           # The transaction at this level is less speedy than one single transaction
@@ -259,10 +250,6 @@ module Gitlab
       def excluded_keys_for_relation(relation)
         reader.attributes_finder.find_excluded_keys(relation)
       end
-
-      def nil_iid_pipeline?(relation_key, relation_item)
-        relation_key == 'ci_pipelines' && relation_item['iid'].nil?
-      end
     end
   end
 end

spec/models/concerns/atomic_internal_id_spec.rb

@@ -9,16 +9,10 @@ describe AtomicInternalId do
   let(:scope_attrs) { { project: milestone.project } }
   let(:usage) { :milestones }
 
-  describe '#ensure_project_iid!' do
-    subject { milestone.ensure_project_iid! }
-
-    it 'generates a new value if non is present' do
-      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
+  describe '#track_project_iid!' do
+    subject { milestone.track_project_iid! }
 
-      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
-    end
-
-    it 'tracks the present value if not generated by InternalId.generate_next' do
+    it 'tracks the present value' do
       milestone.iid = external_iid
 
       expect(InternalId).to receive(:track_greatest).once.with(milestone, scope_attrs, usage, external_iid, anything)
@@ -27,27 +21,51 @@ describe AtomicInternalId do
       subject
     end
 
-    it 'generates a new value if first set with iid= but later set to nil' do
-      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
-
-      milestone.iid = external_iid
-      milestone.iid = nil
+    context 'when value is set by ensure_project_iid!' do
+      context 'with iid_always_track true' do
+        before do
+          stub_feature_flags(iid_always_track: false)
+        end
 
-      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
-    end
+        it 'does not track the value' do  
+          expect(InternalId).not_to receive(:track_greatest)
 
-    context 'with iid_always_track disabled' do
-      before do
-        stub_feature_flags(iid_always_track: false)
+          milestone.ensure_project_iid!
+          subject
+        end
       end
 
-      it 'does not track the present value if generated by InternalId.generate_next' do
-        milestone.ensure_project_iid!
+      context 'with iid_always_track enabled' do
+        before do
+          stub_feature_flags(iid_always_track: true)
+        end
 
-        expect(InternalId).not_to receive(:track_greatest)
+        it 'does not track the value' do
+          expect(InternalId).to receive(:track_greatest)
 
-        subject
+          milestone.ensure_project_iid!
+          subject
+        end
       end
     end
   end
+
+  describe '#ensure_project_iid!' do
+    subject { milestone.ensure_project_iid! }
+
+    it 'generates a new value if non is present' do
+      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
+
+      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
+    end
+
+    it 'generates a new value if first set with iid= but later set to nil' do
+      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
+
+      milestone.iid = external_iid
+      milestone.iid = nil
+
+      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
+    end
+  end
 end