1 files changed, 28 insertions, 13 deletions

diff --git a/app/models/internal_id.rb b/app/models/internal_id.rb
index 4c0469d849a..c735e593da7 100644
--- a/app/models/internal_id.rb
+++ b/app/models/internal_id.rb
@@ -61,13 +61,13 @@ class InternalId < ApplicationRecord
 
   class << self
     def track_greatest(subject, scope, usage, new_value, init)
-      InternalIdGenerator.new(subject, scope, usage)
-        .track_greatest(init, new_value)
+      InternalIdGenerator.new(subject, scope, usage, init)
+        .track_greatest(new_value)
     end
 
     def generate_next(subject, scope, usage, init)
-      InternalIdGenerator.new(subject, scope, usage)
-        .generate(init)
+      InternalIdGenerator.new(subject, scope, usage, init)
+        .generate
     end
 
     def reset(subject, scope, usage, value)
@@ -99,15 +99,18 @@ class InternalId < ApplicationRecord
     # 4) In the absence of a record in the internal_ids table, one will be created
     #    and last_value will be calculated on the fly.
     #
-    # subject: The instance we're generating an internal id for. Gets passed to init if called.
+    # subject: The instance or class we're generating an internal id for.
     # scope: Attributes that define the scope for id generation.
+    #        Valid keys are `project/project_id` and `namespace/namespace_id`.
     # usage: Symbol to define the usage of the internal id, see InternalId.usages
-    attr_reader :subject, :scope, :scope_attrs, :usage
+    # init: Proc that accepts the subject and the scope and returns Integer|NilClass
+    attr_reader :subject, :scope, :scope_attrs, :usage, :init
 
-    def initialize(subject, scope, usage)
+    def initialize(subject, scope, usage, init = nil)
       @subject = subject
       @scope = scope
       @usage = usage
+      @init = init
 
       raise ArgumentError, 'Scope is not well-defined, need at least one column for scope (given: 0)' if scope.empty?
 
@@ -119,13 +122,13 @@ class InternalId < ApplicationRecord
     # Generates next internal id and returns it
     # init: Block that gets called to initialize InternalId record if not present
     #       Make sure to not throw exceptions in the absence of records (if this is expected).
-    def generate(init)
+    def generate
       subject.transaction do
         # Create a record in internal_ids if one does not yet exist
         # and increment its last value
         #
         # Note this will acquire a ROW SHARE lock on the InternalId record
-        (lookup || create_record(init)).increment_and_save!
+        record.increment_and_save!
       end
     end
 
@@ -148,12 +151,20 @@ class InternalId < ApplicationRecord
     # and set its new_value if it is higher than the current last_value
     #
     # Note this will acquire a ROW SHARE lock on the InternalId record
-    def track_greatest(init, new_value)
+    def track_greatest(new_value)
       subject.transaction do
-        (lookup || create_record(init)).track_greatest_and_save!(new_value)
+        record.track_greatest_and_save!(new_value)
       end
     end
 
+    def record
+      @record ||= (lookup || create_record)
+    end
+
+    def with_lock(&block)
+      record.with_lock(&block)
+    end
+
     private
 
     # Retrieve InternalId record for (project, usage) combination, if it exists
@@ -171,12 +182,16 @@ class InternalId < ApplicationRecord
     # was faster in doing this, we'll realize once we hit the unique key constraint
     # violation. We can safely roll-back the nested transaction and perform
     # a lookup instead to retrieve the record.
-    def create_record(init)
+    def create_record
+      raise ArgumentError, 'Cannot initialize without init!' unless init
+
+      instance = subject.is_a?(::Class) ? nil : subject
+
       subject.transaction(requires_new: true) do
         InternalId.create!(
           **scope,
           usage: usage_value,
-          last_value: init.call(subject) || 0
+          last_value: init.call(instance, scope) || 0
         )
       end
     rescue ActiveRecord::RecordNotUnique