Only clear object iid when automatically set

Introduce an additional tracking variable in AtomicInternalId, so that
the values will only be cleared from model objects who had it set
automatically by the ensure_ hook. Manually set values will persist on
the model in the instance where the save rollsback.

app/models/concerns/atomic_internal_id.rb

@@ -81,6 +81,8 @@ module AtomicInternalId
             internal_id_scope_usage,
             init)
           write_attribute(column, value)
+
+          @internal_id_set_manually = false
         end
 
         value
@@ -112,6 +114,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
 
@@ -132,6 +135,8 @@ module AtomicInternalId
       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)

spec/models/concerns/atomic_internal_id_spec.rb

@@ -32,81 +32,6 @@ 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
 
@@ -162,6 +87,95 @@ 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 '.with_project_iid_supply' do
     let(:iid) { 100 }
 

spec/support/shared_examples/models/atomic_internal_id_shared_examples.rb

@@ -50,16 +50,22 @@ RSpec.shared_examples 'AtomicInternalId' do
 
     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!
+        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).not_to be_nil
-
-            raise ActiveRecord::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(read_internal_id).to be_nil
+            expect_iid_to_be_set_and_rollback
+
+            expect(read_internal_id).not_to be_nil
+          end
         end
       end
 
@@ -70,15 +76,19 @@ RSpec.shared_examples 'AtomicInternalId' do
 
           expect(original_id).not_to be_nil
 
-          ActiveRecord::Base.transaction(requires_new: true) do
-            instance.save!
+          expect_iid_to_be_set_and_rollback
 
-            expect(read_internal_id).not_to be_nil
+          expect(read_internal_id).to eq(original_id)
+        end
+      end
 
-            raise ActiveRecord::Rollback
-          end
+      def expect_iid_to_be_set_and_rollback
+        ActiveRecord::Base.transaction(requires_new: true) do
+          instance.save!
 
-          expect(read_internal_id).to eq(original_id)
+          expect(read_internal_id).not_to be_nil
+
+          raise ActiveRecord::Rollback
         end
       end
     end