Merge branch '292225-atomic-internal-id-edge-cases' into 'master'

Set iids on before_save callbacks

See merge request gitlab-org/gitlab!50789

app/models/concerns/atomic_internal_id.rb

@@ -26,20 +26,31 @@
 module AtomicInternalId
   extend ActiveSupport::Concern
 
+  MissingValueError = Class.new(StandardError)
+
   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, presence: true, 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
+
+      if presence
+        before_create :"validate_#{column}_exists!"
+        before_update :"validate_#{column}_exists!"
+      end
 
       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 +73,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
@@ -79,6 +88,8 @@ module AtomicInternalId
             internal_id_scope_usage,
             init)
           write_attribute(column, value)
+
+          @internal_id_set_manually = false
         end
 
         value
@@ -110,6 +121,7 @@ module AtomicInternalId
         super(value).tap do |v|
           # Indicate the iid was set from externally
           @internal_id_needs_tracking = true
+          @internal_id_set_manually = true
         end
       end
 
@@ -128,6 +140,20 @@ module AtomicInternalId
 
         read_attribute(column)
       end
+
+      define_method("clear_#{scope}_#{column}!") do
+        return if @internal_id_set_manually
+
+        return unless public_send(:"#{column}_previously_changed?") # rubocop:disable GitlabSecurity/PublicSend
+
+        write_attribute(column, nil)
+      end
+
+      define_method("validate_#{column}_exists!") do
+        value = read_attribute(column)
+
+        raise MissingValueError, "#{column} was unexpectedly blank!" if value.blank?
+      end
     end
 
     # Defines class methods:

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

@@ -11,13 +11,17 @@ 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) }
       it { is_expected.to validate_length_of(:description).is_at_most(::Issuable::DESCRIPTION_LENGTH_MAX).on(:create) }
 
-      it_behaves_like 'validates description length with custom validation'
+      it_behaves_like 'validates description length with custom validation' do
+        before do
+          allow(InternalId).to receive(:generate_next).and_call_original
+        end
+      end
+
       it_behaves_like 'truncates the description to its allowed maximum length on import'
     end
   end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -1874,7 +1874,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         has_internal_id :iid,
           scope: :project,
           init: ->(s, _scope) { s&.project&.issues&.maximum(:iid) },
-          backfill: true,
           presence: false
       end
     end
@@ -1928,258 +1927,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/concerns/atomic_internal_id_spec.rb

@@ -87,6 +87,158 @@ RSpec.describe AtomicInternalId do
     end
   end
 
+  describe '#clear_scope_iid!' do
+    context 'when no ensure_if condition is given' do
+      it 'clears automatically set IIDs' do
+        expect(milestone).to receive(:clear_project_iid!).and_call_original
+
+        expect_iid_to_be_set_and_rollback(milestone)
+
+        expect(milestone.iid).to be_nil
+      end
+
+      it 'does not clear manually set IIDS' do
+        milestone.iid = external_iid
+
+        expect(milestone).to receive(:clear_project_iid!).and_call_original
+
+        expect_iid_to_be_set_and_rollback(milestone)
+
+        expect(milestone.iid).to eq(external_iid)
+      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 true' do
+        it 'clears automatically set IIDs' do
+          expect(instance).to receive(:clear_project_iid!).and_call_original
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to be_nil
+        end
+
+        it 'does not clear manually set IIDs' do
+          instance.iid = external_iid
+
+          expect(instance).to receive(:clear_project_iid!).and_call_original
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to eq(external_iid)
+        end
+      end
+
+      context 'when the ensure_if condition evaluates to false' do
+        before do
+          instance.importing = true
+        end
+
+        it 'does not clear IIDs' do
+          instance.iid = external_iid
+
+          expect(instance).not_to receive(:clear_project_iid!)
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to eq(external_iid)
+        end
+      end
+    end
+
+    def expect_iid_to_be_set_and_rollback(instance)
+      ActiveRecord::Base.transaction(requires_new: true) do
+        instance.save!
+
+        expect(instance.iid).not_to be_nil
+
+        raise ActiveRecord::Rollback
+      end
+    end
+  end
+
+  describe '#validate_scope_iid_exists!' do
+    let(:test_class) do
+      Class.new(ApplicationRecord) do
+        include AtomicInternalId
+        include Importable
+
+        self.table_name = :milestones
+
+        belongs_to :project
+
+        def self.name
+          'TestClass'
+        end
+      end
+    end
+
+    let(:instance) { test_class.new(milestone.attributes) }
+
+    before do
+      test_class.has_internal_id :iid, scope: :project, presence: presence, ensure_if: -> { !importing }
+
+      instance.importing = true
+    end
+
+    context 'when the presence flag is set' do
+      let(:presence) { true }
+
+      it 'raises an error for blank iids on create' do
+        expect do
+          instance.save!
+        end.to raise_error(described_class::MissingValueError, 'iid was unexpectedly blank!')
+      end
+
+      it 'raises an error for blank iids on update' do
+        instance.iid = 100
+        instance.save!
+
+        instance.iid = nil
+
+        expect do
+          instance.save!
+        end.to raise_error(described_class::MissingValueError, 'iid was unexpectedly blank!')
+      end
+    end
+
+    context 'when the presence flag is not set' do
+      let(:presence) { false }
+
+      it 'does not raise an error for blank iids on create' do
+        expect { instance.save! }.not_to raise_error
+      end
+
+      it 'does not raise an error for blank iids on update' do
+        instance.iid = 100
+        instance.save!
+
+        instance.iid = nil
+
+        expect { instance.save! }.not_to raise_error
+      end
+    end
+  end
+
   describe '.with_project_iid_supply' do
     let(:iid) { 100 }
 

spec/models/concerns/issuable_spec.rb

@@ -45,13 +45,17 @@ 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) }
       it { is_expected.to validate_length_of(:description).is_at_most(described_class::DESCRIPTION_LENGTH_MAX).on(:create) }
 
-      it_behaves_like 'validates description length with custom validation'
+      it_behaves_like 'validates description length with custom validation' do
+        before do
+          allow(InternalId).to receive(:generate_next).and_call_original
+        end
+      end
+
       it_behaves_like 'truncates the description to its allowed maximum length on import'
     end
   end

spec/support/shared_examples/models/atomic_internal_id_shared_examples.rb

@@ -9,19 +9,30 @@ RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
     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")
+        context 'when creating an object' do
+          before do
+            allow_any_instance_of(described_class).to receive(:"ensure_#{scope}_#{internal_id_attribute}!")
+          end
+
+          it 'raises an error if the internal id is blank' do
+            expect { instance.save! }.to raise_error(AtomicInternalId::MissingValueError)
+          end
+        end
+
+        context 'when updating an object' do
+          it 'raises an error if the internal id is blank' do
+            instance.save!
+
+            write_internal_id(nil)
+            allow(instance).to receive(:"ensure_#{scope}_#{internal_id_attribute}!")
+
+            expect { instance.save! }.to raise_error(AtomicInternalId::MissingValueError)
+          end
         end
       end
 
@@ -30,8 +41,27 @@ RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
           skip if validate_presence
         end
 
-        it 'does not validate presence' do
-          expect(instance.errors[internal_id_attribute]).to be_empty
+        context 'when creating an object' do
+          before do
+            allow_any_instance_of(described_class).to receive(:"ensure_#{scope}_#{internal_id_attribute}!")
+          end
+
+          it 'does not raise an error if the internal id is blank' do
+            expect(read_internal_id).to be_nil
+
+            expect { instance.save! }.not_to raise_error
+          end
+        end
+
+        context 'when updating an object' do
+          it 'does not raise an error if the internal id is blank' do
+            instance.save!
+
+            write_internal_id(nil)
+            allow(instance).to receive(:"ensure_#{scope}_#{internal_id_attribute}!")
+
+            expect { instance.save! }.not_to raise_error
+          end
         end
       end
     end
@@ -76,6 +106,51 @@ 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
+        context 'when the internal id is automatically set' do
+          it 'clears it on the instance' do
+            expect_iid_to_be_set_and_rollback
+
+            expect(read_internal_id).to be_nil
+          end
+        end
+
+        context 'when the internal id is manually set' do
+          it 'does not clear it on the instance' do
+            write_internal_id(100)
+
+            expect_iid_to_be_set_and_rollback
+
+            expect(read_internal_id).not_to be_nil
+          end
+        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
+
+          expect_iid_to_be_set_and_rollback
+
+          expect(read_internal_id).to eq(original_id)
+        end
+      end
+
+      def expect_iid_to_be_set_and_rollback
+        ActiveRecord::Base.transaction(requires_new: true) do
+          instance.save!
+
+          expect(read_internal_id).not_to be_nil
+
+          raise ActiveRecord::Rollback
+        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" }