Use before_{create,update} callbacks to set iids

Convert AtomicInternalId to set iid values in before_create or
before_update callbacks, rather than before_validation.

app/models/ci/pipeline.rb

@@ -45,7 +45,7 @@ module Ci
     belongs_to :external_pull_request
     belongs_to :ci_ref, class_name: 'Ci::Ref', foreign_key: :ci_ref_id, inverse_of: :pipelines
 
-    has_internal_id :iid, scope: :project, presence: false,
+    has_internal_id :iid, scope: :project,
       track_if: -> { !importing? },
       ensure_if: -> { !importing? },
       init: ->(pipeline, scope) do

app/models/concerns/atomic_internal_id.rb

@@ -28,18 +28,22 @@ module AtomicInternalId
 
   class_methods do
     def has_internal_id( # rubocop:disable Naming/PredicateName
-      column, scope:, init: :not_given, ensure_if: nil, track_if: nil,
-      presence: true, backfill: false, hook_names: :create)
+      column, scope:, init: :not_given, ensure_if: nil, track_if: nil, hook_names: :create)
       raise "has_internal_id init must not be nil if given." if init.nil?
       raise "has_internal_id needs to be defined on association." unless self.reflect_on_association(scope)
 
       init = infer_init(scope) if init == :not_given
-      before_validation :"track_#{scope}_#{column}!", on: hook_names, if: track_if
-      before_validation :"ensure_#{scope}_#{column}!", on: hook_names, if: ensure_if
-      validates column, presence: presence
+      callback_names = Array.wrap(hook_names).map { |hook_name| :"before_#{hook_name}" }
+      callback_names.each do |callback_name|
+        # rubocop:disable GitlabSecurity/PublicSend
+        public_send(callback_name, :"track_#{scope}_#{column}!", if: track_if)
+        public_send(callback_name, :"ensure_#{scope}_#{column}!", if: ensure_if)
+        # rubocop:enable GitlabSecurity/PublicSend
+      end
+      after_rollback :"clear_#{scope}_#{column}!", on: hook_names, if: ensure_if
 
       define_singleton_internal_id_methods(scope, column, init)
-      define_instance_internal_id_methods(scope, column, init, backfill)
+      define_instance_internal_id_methods(scope, column, init)
     end
 
     private
@@ -62,10 +66,8 @@ module AtomicInternalId
     #   - track_{scope}_{column}!
     #   - reset_{scope}_{column}
     #   - {column}=
-    def define_instance_internal_id_methods(scope, column, init, backfill)
+    def define_instance_internal_id_methods(scope, column, init)
       define_method("ensure_#{scope}_#{column}!") do
-        return if backfill && self.class.where(column => nil).exists?
-
         scope_value = internal_id_read_scope(scope)
         value = read_attribute(column)
         return value unless scope_value
@@ -128,6 +130,12 @@ module AtomicInternalId
 
         read_attribute(column)
       end
+
+      define_method("clear_#{scope}_#{column}!") do
+        return unless public_send(:"#{column}_previously_changed?") # rubocop:disable GitlabSecurity/PublicSend
+
+        write_attribute(column, nil)
+      end
     end
 
     # Defines class methods:

app/models/design_management/design.rb

@@ -27,7 +27,7 @@ module DesignManagement
 
     has_many :events, as: :target, dependent: :delete_all # rubocop:disable Cop/ActiveRecordDependent
 
-    has_internal_id :iid, scope: :project, presence: true,
+    has_internal_id :iid, scope: :project,
       hook_names: %i[create update], # Deal with old records
       track_if: -> { !importing? }
 

app/models/operations/feature_flags/user_list.rb

@@ -13,7 +13,7 @@ module Operations
       has_many :strategy_user_lists
       has_many :strategies, through: :strategy_user_lists
 
-      has_internal_id :iid, scope: :project, presence: true
+      has_internal_id :iid, scope: :project
 
       validates :project, presence: true
       validates :name,

ee/spec/models/concerns/ee/issuable_spec.rb

@@ -11,7 +11,6 @@ RSpec.describe EE::Issuable do
         allow(InternalId).to receive(:generate_next).and_return(nil)
       end
 
-      it { is_expected.to validate_presence_of(:iid) }
       it { is_expected.to validate_presence_of(:author) }
       it { is_expected.to validate_presence_of(:title) }
       it { is_expected.to validate_length_of(:title).is_at_most(::Issuable::TITLE_LENGTH_MAX) }

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -1873,9 +1873,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
         has_internal_id :iid,
           scope: :project,
-          init: ->(s, _scope) { s&.project&.issues&.maximum(:iid) },
-          backfill: true,
-          presence: false
+          init: ->(s, _scope) { s&.project&.issues&.maximum(:iid) }
       end
     end
 
@@ -1928,258 +1926,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       expect(issue_b.iid).to eq(3)
     end
 
-    context 'when the new code creates a row post deploy but before the migration runs' do
-      it 'does not change the row iid' do
-        project = setup
-        issue = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue.reload.iid).to eq(1)
-      end
-
-      it 'backfills iids for rows already in the database' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-      end
-
-      it 'backfills iids across multiple projects' do
-        project_a = setup
-        project_b = setup
-        issue_a = issues.create!(project_id: project_a.id)
-        issue_b = issues.create!(project_id: project_b.id)
-        issue_c = Issue.create!(project_id: project_a.id)
-        issue_d = Issue.create!(project_id: project_b.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(1)
-        expect(issue_c.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(2)
-      end
-
-      it 'generates iids properly for models created after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        issue_d = Issue.create!(project_id: project.id)
-        issue_e = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-        expect(issue_d.iid).to eq(4)
-        expect(issue_e.iid).to eq(5)
-      end
-
-      it 'backfills iids and properly generates iids for new models across multiple projects' do
-        project_a = setup
-        project_b = setup
-        issue_a = issues.create!(project_id: project_a.id)
-        issue_b = issues.create!(project_id: project_b.id)
-        issue_c = Issue.create!(project_id: project_a.id)
-        issue_d = Issue.create!(project_id: project_b.id)
-
-        model.backfill_iids('issues')
-
-        issue_e = Issue.create!(project_id: project_a.id)
-        issue_f = Issue.create!(project_id: project_b.id)
-        issue_g = Issue.create!(project_id: project_a.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(1)
-        expect(issue_c.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(2)
-        expect(issue_e.iid).to eq(3)
-        expect(issue_f.iid).to eq(3)
-        expect(issue_g.iid).to eq(4)
-      end
-    end
-
-    context 'when the new code creates a model and then old code creates a model post deploy but before the migration runs' do
-      it 'backfills iids' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = Issue.create!(project_id: project.id)
-        issue_c = issues.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-      end
-
-      it 'generates an iid for a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_d = issues.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        issue_e = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-        expect(issue_d.reload.iid).to eq(4)
-        expect(issue_e.iid).to eq(5)
-      end
-    end
-
-    context 'when the new code and old code alternate creating models post deploy but before the migration runs' do
-      it 'backfills iids' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = Issue.create!(project_id: project.id)
-        issue_c = issues.create!(project_id: project.id)
-        issue_d = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-        expect(issue_d.reload.iid).to eq(4)
-      end
-
-      it 'generates an iid for a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_d = issues.create!(project_id: project.id)
-        issue_e = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        issue_f = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_c.reload.iid).to eq(3)
-        expect(issue_d.reload.iid).to eq(4)
-        expect(issue_e.reload.iid).to eq(5)
-        expect(issue_f.iid).to eq(6)
-      end
-    end
-
-    context 'when the new code creates and deletes a model post deploy but before the migration runs' do
-      it 'backfills iids for rows already in the database' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-      end
-
-      it 'successfully creates a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-
-        model.backfill_iids('issues')
-
-        issue_d = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_d.iid).to eq(3)
-      end
-    end
-
-    context 'when the new code creates and deletes a model and old code creates a model post deploy but before the migration runs' do
-      it 'backfills iids' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-        issue_d = issues.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(3)
-      end
-
-      it 'successfully creates a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-        issue_d = issues.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        issue_e = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(3)
-        expect(issue_e.iid).to eq(4)
-      end
-    end
-
-    context 'when the new code creates and deletes a model and then creates another model post deploy but before the migration runs' do
-      it 'successfully generates an iid for a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-        issue_d = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(3)
-      end
-
-      it 'successfully generates an iid for a new model after the migration' do
-        project = setup
-        issue_a = issues.create!(project_id: project.id)
-        issue_b = issues.create!(project_id: project.id)
-        issue_c = Issue.create!(project_id: project.id)
-        issue_c.delete
-        issue_d = Issue.create!(project_id: project.id)
-
-        model.backfill_iids('issues')
-
-        issue_e = Issue.create!(project_id: project.id)
-
-        expect(issue_a.reload.iid).to eq(1)
-        expect(issue_b.reload.iid).to eq(2)
-        expect(issue_d.reload.iid).to eq(3)
-        expect(issue_e.iid).to eq(4)
-      end
-    end
-
     context 'when the first model is created for a project after the migration' do
       it 'generates an iid' do
         project_a = setup

spec/models/ci/pipeline_spec.rb

@@ -592,7 +592,7 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
   end
 
   describe 'modules' do
-    it_behaves_like 'AtomicInternalId', validate_presence: false do
+    it_behaves_like 'AtomicInternalId' do
       let(:internal_id_attribute) { :iid }
       let(:instance) { build(:ci_pipeline) }
       let(:scope) { :project }

spec/models/concerns/atomic_internal_id_spec.rb

@@ -32,6 +32,81 @@ RSpec.describe AtomicInternalId do
           milestone.save!
         end
       end
+
+      context 'when the save is rolled back' do
+        context 'when no ensure_if condition is given' do
+          it 'clears the instance IID' do
+            expect(milestone).to receive(:clear_project_iid!).and_call_original
+
+            ActiveRecord::Base.transaction(requires_new: true) do
+              milestone.save!
+
+              expect(milestone.iid).to eq(external_iid)
+
+              raise ActiveRecord::Rollback
+            end
+
+            expect(milestone.iid).to be_nil
+          end
+        end
+
+        context 'when an ensure_if condition is given' do
+          let(:test_class) do
+            Class.new(ApplicationRecord) do
+              include AtomicInternalId
+              include Importable
+
+              self.table_name = :milestones
+
+              belongs_to :project
+
+              has_internal_id :iid, scope: :project, track_if: -> { !importing }, ensure_if: -> { !importing }
+
+              def self.name
+                'TestClass'
+              end
+            end
+          end
+
+          let(:instance) { test_class.new(milestone.attributes) }
+
+          context 'when the ensure_if condition evaluates to false' do
+            it 'clears the instance IID' do
+              expect(instance).to receive(:clear_project_iid!).and_call_original
+
+              ActiveRecord::Base.transaction(requires_new: true) do
+                instance.save!
+
+                expect(instance.iid).not_to be_nil
+
+                raise ActiveRecord::Rollback
+              end
+
+              expect(instance.iid).to be_nil
+            end
+          end
+
+          context 'when the ensure_if condition evaluates to true' do
+            before do
+              instance.importing = true
+            end
+
+            it 'does not clear the instance IID' do
+              expect(instance).not_to receive(:clear_project_iid!)
+
+              ActiveRecord::Base.transaction(requires_new: true) do
+                instance.save!
+
+                expect(instance.iid).not_to be_nil
+
+                raise ActiveRecord::Rollback
+              end
+
+              expect(instance.iid).not_to be_nil
+            end
+          end
+        end
+      end
     end
   end
 

spec/models/concerns/issuable_spec.rb

@@ -45,7 +45,6 @@ RSpec.describe Issuable do
       end
 
       it { is_expected.to validate_presence_of(:project) }
-      it { is_expected.to validate_presence_of(:iid) }
       it { is_expected.to validate_presence_of(:author) }
       it { is_expected.to validate_presence_of(:title) }
       it { is_expected.to validate_length_of(:title).is_at_most(described_class::TITLE_LENGTH_MAX) }

spec/models/design_management/design_spec.rb

@@ -11,7 +11,7 @@ RSpec.describe DesignManagement::Design do
   let_it_be(:design3) { create(:design, :with_versions, issue: issue, versions_count: 1) }
   let_it_be(:deleted_design) { create(:design, :with_versions, deleted: true) }
 
-  it_behaves_like 'AtomicInternalId', validate_presence: true do
+  it_behaves_like 'AtomicInternalId' do
     let(:internal_id_attribute) { :iid }
     let(:instance) { build(:design, issue: issue) }
     let(:scope) { :project }

spec/models/operations/feature_flag_spec.rb

@@ -64,7 +64,7 @@ RSpec.describe Operations::FeatureFlag do
       end
     end
 
-    it_behaves_like 'AtomicInternalId', validate_presence: true do
+    it_behaves_like 'AtomicInternalId' do
       let(:internal_id_attribute) { :iid }
       let(:instance) { build(:operations_feature_flag) }
       let(:scope) { :project }

spec/support/shared_examples/models/atomic_internal_id_shared_examples.rb

 # frozen_string_literal: true
 
-RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
+RSpec.shared_examples 'AtomicInternalId' do
   describe '.has_internal_id' do
     describe 'Module inclusion' do
       subject { described_class }
@@ -8,34 +8,6 @@ RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
       it { is_expected.to include_module(AtomicInternalId) }
     end
 
-    describe 'Validation' do
-      before do
-        allow_any_instance_of(described_class).to receive(:"ensure_#{scope}_#{internal_id_attribute}!")
-
-        instance.valid?
-      end
-
-      context 'when presence validation is required' do
-        before do
-          skip unless validate_presence
-        end
-
-        it 'validates presence' do
-          expect(instance.errors[internal_id_attribute]).to include("can't be blank")
-        end
-      end
-
-      context 'when presence validation is not required' do
-        before do
-          skip if validate_presence
-        end
-
-        it 'does not validate presence' do
-          expect(instance.errors[internal_id_attribute]).to be_empty
-        end
-      end
-    end
-
     describe 'Creating an instance' do
       subject { instance.save! }
 
@@ -76,6 +48,41 @@ RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
       end
     end
 
+    describe 'unsetting the instance internal id on rollback' do
+      context 'when the internal id has been changed' do
+        it 'clears it on the instance' do
+          ActiveRecord::Base.transaction(requires_new: true) do
+            instance.save!
+
+            expect(read_internal_id).not_to be_nil
+
+            raise ActiveRecord::Rollback
+          end
+
+          expect(read_internal_id).to be_nil
+        end
+      end
+
+      context 'when the internal id has not been changed' do
+        it 'preserves the value on the instance' do
+          instance.save!
+          original_id = read_internal_id
+
+          expect(original_id).not_to be_nil
+
+          ActiveRecord::Base.transaction(requires_new: true) do
+            instance.save!
+
+            expect(read_internal_id).not_to be_nil
+
+            raise ActiveRecord::Rollback
+          end
+
+          expect(read_internal_id).to eq(original_id)
+        end
+      end
+    end
+
     describe 'supply of internal ids' do
       let(:scope_value) { scope_attrs.each_value.first }
       let(:method_name) { :"with_#{scope}_#{internal_id_attribute}_supply" }