Raise error if iid is not set as expected

With the change to using before_save hook to set the iid on an object,
normal validations can't be used to verify the presence of the iid.

In normal situations, the iid should always be set by either
ensure_scope_iid hook or application logic, but since we're lacking
database constraints on some tables we don't want to silently insert bad
data. To prevent unforeseen issues cause by bugs, add an additional
before_save hook to raise an exception if the iid is not set as it
should be.

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,
+    has_internal_id :iid, scope: :project, presence: false,
       track_if: -> { !importing? },
       ensure_if: -> { !importing? },
       init: ->(pipeline, scope) do

app/models/concerns/atomic_internal_id.rb

@@ -26,9 +26,11 @@
 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, 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)
 
@@ -42,6 +44,11 @@ module AtomicInternalId
       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)
     end
@@ -141,6 +148,12 @@ module AtomicInternalId
 
         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:

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,
+    has_internal_id :iid, scope: :project, presence: true,
       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
+      has_internal_id :iid, scope: :project, presence: true
 
       validates :project, presence: true
       validates :name,

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

@@ -16,7 +16,12 @@ RSpec.describe EE::Issuable do
       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

@@ -1873,7 +1873,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
         has_internal_id :iid,
           scope: :project,
-          init: ->(s, _scope) { s&.project&.issues&.maximum(:iid) }
+          init: ->(s, _scope) { s&.project&.issues&.maximum(:iid) },
+          presence: false
       end
     end
 

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' do
+    it_behaves_like 'AtomicInternalId', validate_presence: false do
       let(:internal_id_attribute) { :iid }
       let(:instance) { build(:ci_pipeline) }
       let(:scope) { :project }

spec/models/concerns/atomic_internal_id_spec.rb

@@ -176,6 +176,69 @@ RSpec.describe AtomicInternalId do
     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

@@ -50,7 +50,12 @@ RSpec.describe Issuable do
       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/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' do
+  it_behaves_like 'AtomicInternalId', validate_presence: true 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' do
+    it_behaves_like 'AtomicInternalId', validate_presence: true 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
+RSpec.shared_examples 'AtomicInternalId' do |validate_presence: true|
   describe '.has_internal_id' do
     describe 'Module inclusion' do
       subject { described_class }
@@ -8,6 +8,64 @@ RSpec.shared_examples 'AtomicInternalId' do
       it { is_expected.to include_module(AtomicInternalId) }
     end
 
+    describe 'Validation' do
+      context 'when presence validation is required' do
+        before do
+          skip unless validate_presence
+        end
+
+        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
+
+      context 'when presence validation is not required' do
+        before do
+          skip if validate_presence
+        end
+
+        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
+
     describe 'Creating an instance' do
       subject { instance.save! }