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,22 +32,80 @@ RSpec.describe AtomicInternalId do
           milestone.save!
         end
       end
+    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
+  describe '#track_project_iid!' do
+    subject { milestone.track_project_iid! }
 
-            ActiveRecord::Base.transaction(requires_new: true) do
-              milestone.save!
+    it 'tracks the present value' do
+      milestone.iid = external_iid
 
-              expect(milestone.iid).to eq(external_iid)
+      expect(InternalId).to receive(:track_greatest).once.with(milestone, scope_attrs, usage, external_iid, anything)
+      expect(InternalId).not_to receive(:generate_next)
 
-              raise ActiveRecord::Rollback
+      subject
+    end
+
+    context 'when value is set by ensure_project_iid!' do
+      it 'does not track the value' do
+        expect(InternalId).not_to receive(:track_greatest)
+
+        milestone.ensure_project_iid!
+        subject
+      end
+
+      it 'tracks the iid for the scope that is actually present' do
+        milestone.iid = external_iid
+
+        expect(InternalId).to receive(:track_greatest).once.with(milestone, scope_attrs, usage, external_iid, anything)
+        expect(InternalId).not_to receive(:generate_next)
+
+        # group scope is not present here, the milestone does not have a group
+        milestone.track_group_iid!
+        subject
+      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
+
+  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
@@ -70,95 +128,51 @@ RSpec.describe AtomicInternalId do
 
       let(:instance) { test_class.new(milestone.attributes) }
 
-          context 'when the ensure_if condition evaluates to false' do
-            it 'clears the instance IID' do
+      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
 
-              ActiveRecord::Base.transaction(requires_new: true) do
-                instance.save!
-
-                expect(instance.iid).not_to be_nil
-
-                raise ActiveRecord::Rollback
-              end
+          expect_iid_to_be_set_and_rollback(instance)
 
           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!)
+        it 'does not clear manually set IIDs' do
+          instance.iid = external_iid
 
-              ActiveRecord::Base.transaction(requires_new: true) do
-                instance.save!
+          expect(instance).to receive(:clear_project_iid!).and_call_original
 
-                expect(instance.iid).not_to be_nil
+          expect_iid_to_be_set_and_rollback(instance)
 
-                raise ActiveRecord::Rollback
-              end
-
-              expect(instance.iid).not_to be_nil
-            end
-          end
-        end
-      end
+          expect(instance.iid).to eq(external_iid)
         end
       end
 
-  describe '#track_project_iid!' do
-    subject { milestone.track_project_iid! }
-
-    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)
-      expect(InternalId).not_to receive(:generate_next)
-
-      subject
+      context 'when the ensure_if condition evaluates to false' do
+        before do
+          instance.importing = true
         end
 
-    context 'when value is set by ensure_project_iid!' do
-      it 'does not track the value' do
-        expect(InternalId).not_to receive(:track_greatest)
-
-        milestone.ensure_project_iid!
-        subject
-      end
+        it 'does not clear IIDs' do
+          instance.iid = external_iid
 
-      it 'tracks the iid for the scope that is actually present' do
-        milestone.iid = external_iid
+          expect(instance).not_to receive(:clear_project_iid!)
 
-        expect(InternalId).to receive(:track_greatest).once.with(milestone, scope_attrs, usage, external_iid, anything)
-        expect(InternalId).not_to receive(:generate_next)
+          expect_iid_to_be_set_and_rollback(instance)
 
-        # group scope is not present here, the milestone does not have a group
-        milestone.track_group_iid!
-        subject
+          expect(instance.iid).to eq(external_iid)
         end
       end
     end
 
-  describe '#ensure_project_iid!' do
-    subject { milestone.ensure_project_iid! }
+    def expect_iid_to_be_set_and_rollback(instance)
+      ActiveRecord::Base.transaction(requires_new: true) do
+        instance.save!
 
-    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(instance.iid).not_to be_nil
 
-      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
+        raise ActiveRecord::Rollback
       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
 

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
+        context 'when the internal id is automatically set' do
           it 'clears it on the instance' do
-          ActiveRecord::Base.transaction(requires_new: true) do
-            instance.save!
+            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
 
-          expect(read_internal_id).to be_nil
+        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
 
@@ -70,6 +76,13 @@ RSpec.shared_examples 'AtomicInternalId' do
 
           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!
 
@@ -77,9 +90,6 @@ RSpec.shared_examples 'AtomicInternalId' do
 
           raise ActiveRecord::Rollback
         end
-
-          expect(read_internal_id).to eq(original_id)
-        end
       end
     end