From a99281cbb04016652892b16433f296724c27bc32 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Fri, 8 Aug 2025 02:14:01 +0800
Subject: [PATCH 01/29] Capture changes

---
 app/models/cells/outstanding_lease.rb         |  3 +-
 app/models/concerns/cells/inspect_changes.rb  | 85 +++++++++++++++++++
 .../concerns/cells/transaction_callback.rb    | 58 +++++++++++++
 app/models/concerns/cells/unique.rb           | 45 ++++++++++
 app/models/email.rb                           |  6 ++
 app/models/gpg_key.rb                         |  5 ++
 app/models/gpg_key_subkey.rb                  |  5 ++
 app/models/key.rb                             |  4 +
 app/models/organizations/organization.rb      |  3 +
 app/models/route.rb                           |  3 +
 app/models/user.rb                            |  4 +
 ...ctive_record_cells_transaction_callback.rb | 10 +++
 config/initializers/doorkeeper.rb             |  9 ++
 db/docs/cells_outstanding_leases.yml          |  4 +-
 ...2145304_create_cells_outstanding_leases.rb | 30 +++++++
 db/schema_migrations/20250822145304           |  1 +
 db/structure.sql                              |  2 +-
 17 files changed, 273 insertions(+), 4 deletions(-)
 create mode 100644 app/models/concerns/cells/inspect_changes.rb
 create mode 100644 app/models/concerns/cells/transaction_callback.rb
 create mode 100644 app/models/concerns/cells/unique.rb
 create mode 100644 config/initializers/active_record_cells_transaction_callback.rb
 create mode 100644 db/migrate/20250822145304_create_cells_outstanding_leases.rb
 create mode 100644 db/schema_migrations/20250822145304

diff --git a/app/models/cells/outstanding_lease.rb b/app/models/cells/outstanding_lease.rb
index f74763e6b828ae..2577fc25c33983 100644
--- a/app/models/cells/outstanding_lease.rb
+++ b/app/models/cells/outstanding_lease.rb
@@ -2,6 +2,7 @@
 
 module Cells
   class OutstandingLease < ApplicationRecord
-    self.primary_key = :uuid
+    self.table_name = 'cells_outstanding_leases'
+    self.primary_key = 'uuid'
   end
 end
diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
new file mode 100644
index 00000000000000..876360d29e0933
--- /dev/null
+++ b/app/models/concerns/cells/inspect_changes.rb
@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module Cells
+  class InspectChanges
+    attr_reader :record
+
+    def initialize(new_record)
+      @record = new_record
+    end
+
+    def inspect_changes
+      return if !record.is_a?(Unique) || inspected
+
+      prepare_inspecting
+
+      inspect_record
+      inspect_associations
+
+      if TransactionCallback.visited
+        # There's an actual transaction, and TransactionCallback will
+        # handle this case right before committing
+      else
+        # Since the transaction is optimized away, TransactionCallback
+        # cannot handle this case, and we can just assume changes are
+        # going to be made, we treat this as if before_commit_all
+        Unique.claim_and_commit(visited)
+      end
+    end
+
+    private
+
+    def visited
+      @visited ||= TransactionCallback.visited || {}
+    end
+
+    def inspect_changes_key
+      @inspect_changes_key ||= [record.class.table_name, record.id]
+    end
+
+    def inspected
+      visited[inspect_changes_key]
+    end
+
+    def prepare_inspecting
+      visited[inspect_changes_key] = {}
+    end
+
+    def inspect_record
+      if record.destroyed?
+        inspect_deleted_attributes
+      else
+        inspect_saved_changes
+      end
+    end
+
+    def inspect_deleted_attributes
+      record.class.cells_unique_attributes.each_key do |column|
+        inspected[column] = [record.attributes[column.to_s]]
+      end
+    end
+
+    def inspect_saved_changes
+      record.saved_changes.each do |column, (old_value, new_value)|
+        if record.class.cells_unique_attributes[column.to_sym] # rubocop:disable Style/Next -- I think this is more clear
+          # XXX: new_value in GpgKey is transformed into upcase,
+          # which mismatches with what we're storing in the database
+          new_value = record.attributes[column]
+          inspected[column] = [old_value, new_value]
+        end
+      end
+    end
+
+    def inspect_associations
+      record.class.reflect_on_all_associations.each do |reflection|
+        association_proxy = record.association(reflection.name)
+
+        if association_proxy.loaded? # rubocop:disable Style/Next -- I think this is more clear
+          Array.wrap(association_proxy.target).each do |association_record|
+            InspectChanges.new(association_record).inspect_changes
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/app/models/concerns/cells/transaction_callback.rb b/app/models/concerns/cells/transaction_callback.rb
new file mode 100644
index 00000000000000..69047adbeac274
--- /dev/null
+++ b/app/models/concerns/cells/transaction_callback.rb
@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Cells
+  module TransactionCallback
+    THREAD_VARIABLE_NAME = :cells_unique_transaction_callback_visited_stack
+
+    def self.visited
+      visited_stack.first # Flatten to the first
+    end
+
+    def self.visited_stack
+      Thread.current.thread_variable_get(THREAD_VARIABLE_NAME) ||
+        Thread.current.thread_variable_set(THREAD_VARIABLE_NAME, [])
+    end
+
+    def begin_isolated_db_transaction(...)
+      TransactionCallback.visited_stack << {}
+      super
+    end
+
+    def begin_db_transaction
+      TransactionCallback.visited_stack << {}
+      super
+    end
+
+    def commit_db_transaction
+      # Don't pop the visited stack yet, and only do that after it's
+      # successfully committed, because errors can happen while claiming,
+      # and when that happens, we'll roll it back, and we need to look at
+      # the stack to know if it's a unbalanced ROLLBACK or not, therefore
+      # we only pop after successful COMMIT or ROLLBACk.
+      Unique.claim_and_commit(TransactionCallback.visited_stack.last) do
+        super
+        TransactionCallback.visited_stack.pop
+      end
+    end
+
+    def rollback_db_transaction
+      Unique.rollback do
+        super
+
+        # Rails can be issuing unbalanced ROLLBACK (no corresponding BEGIN),
+        # and in that case we ignore and do not need to rollback the last
+        # lease because it's already committed! This can happen if there are
+        # errors internally coming from Rails while committing to the local
+        # database. We determine this by checking the visited stack. If the
+        # stack is empty, it means it's not balanced and we do not want to
+        # rollback on Topology service in this case.
+        if TransactionCallback.visited_stack.empty?
+          false
+        else
+          TransactionCallback.visited_stack.pop
+          true
+        end
+      end
+    end
+  end
+end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
new file mode 100644
index 00000000000000..9ad269eb9bf528
--- /dev/null
+++ b/app/models/concerns/cells/unique.rb
@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Cells
+  module Unique
+    extend ActiveSupport::Concern
+
+    included do
+      next unless Gitlab.config.cell.enabled
+
+      before_commit do
+        InspectChanges.new(self).inspect_changes
+      end
+    end
+
+    class_methods do
+      def cells_unique_attributes
+        @cells_unique_attributes ||= {}
+      end
+
+      def cells_unique_attribute(name, args = {})
+        cells_unique_attributes[name] = args
+      end
+    end
+
+    def self.claim_and_commit(changes)
+      if changes.empty?
+        yield if block_given?
+      else
+        pp changes
+
+        claim_client = Gitlab::TopologyServiceClient::ClaimService.new
+        claim_client.claim_and_commit(changes) do
+          yield if block_given?
+        end
+      end
+    end
+
+    def self.rollback
+      claim_client = Gitlab::TopologyServiceClient::ClaimService.new
+      claim_client.rollback do
+        yield
+      end
+    end
+  end
+end
diff --git a/app/models/email.rb b/app/models/email.rb
index c2e6919322358d..120f7744d97fe3 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -3,6 +3,12 @@
 class Email < ApplicationRecord
   include Sortable
   include Gitlab::SQL::Pattern
+  include Cells::Unique
+
+  # TODO: Figure out why sometimes this wasn't released on Topology service
+  # when the user is deleted. Some of the operations maybe doesn't trigger
+  # before_commit hook. Maybe it's deleted by CASCADE from PostgreSQL?
+  cells_unique_attribute :email
 
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 43c2a70b8670a4..623ae909b3dc93 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -5,6 +5,11 @@ class GpgKey < ApplicationRecord
   KEY_SUFFIX = '-----END PGP PUBLIC KEY BLOCK-----'
 
   include ShaAttribute
+  include Cells::Unique
+
+  cells_unique_attribute :key
+  cells_unique_attribute :fingerprint
+  cells_unique_attribute :primary_keyid
 
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 110bf45113652e..f2ab9cc9180ca0 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -2,6 +2,11 @@
 
 class GpgKeySubkey < ApplicationRecord
   include ShaAttribute
+  include Cells::Unique
+
+  # TODO: This can be claimed fine, but for deleting, it's not deleted on
+  # Topology service for some reason. We need to figure out why
+  # cells_unique_attribute :fingerprint
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index 78c4518b46c331..fa21919865e295 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -8,6 +8,10 @@ class Key < ApplicationRecord
   include FromUnion
   include Todoable
   include CreatedAtFilterable
+  include Cells::Unique
+
+  cells_unique_attribute :key
+  cells_unique_attribute :fingerprint_sha256
 
   sha256_attribute :fingerprint_sha256
 
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 37253257299f51..529a040cb5a1a2 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -8,6 +8,9 @@ class Organization < ApplicationRecord
     include Gitlab::Routing.url_helpers
     include FeatureGate
     include Organizations::Isolatable
+    include Cells::Unique
+
+    cells_unique_attribute :path
 
     DEFAULT_ORGANIZATION_ID = 1
 
diff --git a/app/models/route.rb b/app/models/route.rb
index 7b12bdbadc1fc9..fec34861261d2b 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -4,6 +4,9 @@ class Route < ApplicationRecord
   include CaseSensitivity
   include Gitlab::SQL::Pattern
   include EachBatch
+  include Cells::Unique
+
+  cells_unique_attribute :path
 
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
diff --git a/app/models/user.rb b/app/models/user.rb
index b4403061ddafd6..2bbb5990a34769 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -36,6 +36,10 @@ class User < ApplicationRecord
   include Gitlab::InternalEventsTracking
   include Ci::PipelineScheduleOwnershipValidator
   include Users::DependentAssociations
+  include Cells::Unique
+
+  # cells_unique_attribute :email # This clashes with emails.email
+  cells_unique_attribute :username
 
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
diff --git a/config/initializers/active_record_cells_transaction_callback.rb b/config/initializers/active_record_cells_transaction_callback.rb
new file mode 100644
index 00000000000000..e213d795fb7d7d
--- /dev/null
+++ b/config/initializers/active_record_cells_transaction_callback.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+ActiveSupport.on_load(:active_record) do
+  # Prepend into AbstractAdapter doesn't fully work, so we prepend to
+  # actual adapter we use.
+  if Gitlab.config.cell.enabled
+    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+      .prepend(Cells::TransactionCallback)
+  end
+end
diff --git a/config/initializers/doorkeeper.rb b/config/initializers/doorkeeper.rb
index 215ed6f4ca9607..cdc3459d1ba0f1 100644
--- a/config/initializers/doorkeeper.rb
+++ b/config/initializers/doorkeeper.rb
@@ -151,6 +151,15 @@
   custom_access_token_attributes [:organization_id]
 
   enable_dynamic_scopes
+
+  # Following doorkeeper monkey patches are to claim required resources
+
+  Doorkeeper::Application.class_eval do
+    include Cells::Unique
+
+    cells_unique_attribute :uid
+  end
+
   # Following doorkeeper monkey patches are to identify the organization on best effort basis
 
   Doorkeeper::OAuth::PasswordAccessTokenRequest.class_eval do
diff --git a/db/docs/cells_outstanding_leases.yml b/db/docs/cells_outstanding_leases.yml
index 85e4ead9706cf2..2f17ceb1f446f7 100644
--- a/db/docs/cells_outstanding_leases.yml
+++ b/db/docs/cells_outstanding_leases.yml
@@ -5,7 +5,7 @@ classes:
 feature_categories:
 - cell
 description: Cells outstanding leases for claiming
-introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/203567
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/200701
 milestone: '18.4'
-gitlab_schema: gitlab_main_cell_local
+gitlab_schema: gitlab_main_cell_setting
 table_size: small
diff --git a/db/migrate/20250822145304_create_cells_outstanding_leases.rb b/db/migrate/20250822145304_create_cells_outstanding_leases.rb
new file mode 100644
index 00000000000000..8e38d1b246535c
--- /dev/null
+++ b/db/migrate/20250822145304_create_cells_outstanding_leases.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+# See https://docs.gitlab.com/ee/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class CreateCellsOutstandingLeases < Gitlab::Database::Migration[2.3]
+  # When using the methods "add_concurrent_index" or "remove_concurrent_index"
+  # you must disable the use of transactions
+  # as these methods can not run in an existing transaction.
+  # When using "add_concurrent_index" or "remove_concurrent_index" methods make sure
+  # that either of them is the _only_ method called in the migration,
+  # any other changes should go in a separate migration.
+  # This ensures that upon failure _only_ the index creation or removing fails
+  # and can be retried or reverted easily.
+  #
+  # To disable transactions uncomment the following line and remove these
+  # comments:
+  # disable_ddl_transaction!
+  milestone '18.4'
+
+  # Add dependent 'batched_background_migrations.queued_migration_version' values.
+  # DEPENDENT_BATCHED_BACKGROUND_MIGRATIONS = []
+
+  def change
+    create_table :cells_outstanding_leases, id: false do |t|
+      t.primary_key :uuid, :uuid
+      t.timestamps_with_timezone null: false
+    end
+  end
+end
diff --git a/db/schema_migrations/20250822145304 b/db/schema_migrations/20250822145304
new file mode 100644
index 00000000000000..4e64198a0f3573
--- /dev/null
+++ b/db/schema_migrations/20250822145304
@@ -0,0 +1 @@
+0acfd536a9af28741684976c1cf1a11362bc6beada54b11545a81dffd01701d0
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index 387820a3fa389b..d8b9e5e5c4855a 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -13299,7 +13299,7 @@ CREATE SEQUENCE catalog_verified_namespaces_id_seq
 ALTER SEQUENCE catalog_verified_namespaces_id_seq OWNED BY catalog_verified_namespaces.id;
 
 CREATE TABLE cells_outstanding_leases (
-    uuid uuid NOT NULL,
+    uuid uuid DEFAULT gen_random_uuid() NOT NULL,
     created_at timestamp with time zone NOT NULL,
     updated_at timestamp with time zone NOT NULL
 );
-- 
GitLab


From 6fbd98d4f61e22e2d1eeb77cd6e3e27bcc8c295b Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Fri, 29 Aug 2025 16:55:02 +0200
Subject: [PATCH 02/29] Fixes rubocop offences

---
 app/models/cells/outstanding_lease.rb         |  3 +--
 app/models/concerns/cells/unique.rb           |  2 --
 ...2145304_create_cells_outstanding_leases.rb | 20 +------------------
 3 files changed, 2 insertions(+), 23 deletions(-)

diff --git a/app/models/cells/outstanding_lease.rb b/app/models/cells/outstanding_lease.rb
index 2577fc25c33983..f74763e6b828ae 100644
--- a/app/models/cells/outstanding_lease.rb
+++ b/app/models/cells/outstanding_lease.rb
@@ -2,7 +2,6 @@
 
 module Cells
   class OutstandingLease < ApplicationRecord
-    self.table_name = 'cells_outstanding_leases'
-    self.primary_key = 'uuid'
+    self.primary_key = :uuid
   end
 end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 9ad269eb9bf528..3f0fe3b3c0149a 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -26,8 +26,6 @@ def self.claim_and_commit(changes)
       if changes.empty?
         yield if block_given?
       else
-        pp changes
-
         claim_client = Gitlab::TopologyServiceClient::ClaimService.new
         claim_client.claim_and_commit(changes) do
           yield if block_given?
diff --git a/db/migrate/20250822145304_create_cells_outstanding_leases.rb b/db/migrate/20250822145304_create_cells_outstanding_leases.rb
index 8e38d1b246535c..1110dee19c4ab2 100644
--- a/db/migrate/20250822145304_create_cells_outstanding_leases.rb
+++ b/db/migrate/20250822145304_create_cells_outstanding_leases.rb
@@ -1,28 +1,10 @@
 # frozen_string_literal: true
 
-# See https://docs.gitlab.com/ee/development/migration_style_guide.html
-# for more information on how to write migrations for GitLab.
-
 class CreateCellsOutstandingLeases < Gitlab::Database::Migration[2.3]
-  # When using the methods "add_concurrent_index" or "remove_concurrent_index"
-  # you must disable the use of transactions
-  # as these methods can not run in an existing transaction.
-  # When using "add_concurrent_index" or "remove_concurrent_index" methods make sure
-  # that either of them is the _only_ method called in the migration,
-  # any other changes should go in a separate migration.
-  # This ensures that upon failure _only_ the index creation or removing fails
-  # and can be retried or reverted easily.
-  #
-  # To disable transactions uncomment the following line and remove these
-  # comments:
-  # disable_ddl_transaction!
   milestone '18.4'
 
-  # Add dependent 'batched_background_migrations.queued_migration_version' values.
-  # DEPENDENT_BATCHED_BACKGROUND_MIGRATIONS = []
-
   def change
-    create_table :cells_outstanding_leases, id: false do |t|
+    create_table :cells_outstanding_leases, id: false do |t| # rubocop:disable Migration/EnsureFactoryForTable -- No factory needed
       t.primary_key :uuid, :uuid
       t.timestamps_with_timezone null: false
     end
-- 
GitLab


From bb83155aa3dbbacd266dc7e99f04317ce50ff5a3 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Thu, 4 Sep 2025 12:34:45 +0200
Subject: [PATCH 03/29] Add timeout for before_commit

---
 app/models/concerns/cells/unique.rb | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 3f0fe3b3c0149a..bee3820e0cc98f 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -4,11 +4,15 @@ module Cells
   module Unique
     extend ActiveSupport::Concern
 
+    CLAIM_TIMEOUT_MS = 10000
+
     included do
       next unless Gitlab.config.cell.enabled
 
       before_commit do
-        InspectChanges.new(self).inspect_changes
+        with_claim_timeout do
+          InspectChanges.new(self).inspect_changes
+        end
       end
     end
 
@@ -39,5 +43,12 @@ def self.rollback
         yield
       end
     end
+
+    private
+
+    def with_claim_timeout(timeout_ms = CLAIM_TIMEOUT_MS)
+      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
+      yield
+    end
   end
 end
-- 
GitLab


From 7034a5872da98145c50d467df981fa0f2056c43e Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Thu, 4 Sep 2025 15:56:27 +0200
Subject: [PATCH 04/29] Move attribute types to class

---
 app/models/email.rb                      | 3 ++-
 app/models/gpg_key.rb                    | 6 +++---
 app/models/gpg_key_subkey.rb             | 2 +-
 app/models/key.rb                        | 4 ++--
 app/models/organizations/organization.rb | 3 ++-
 app/models/route.rb                      | 2 +-
 app/models/user.rb                       | 2 +-
 7 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/app/models/email.rb b/app/models/email.rb
index 120f7744d97fe3..3bea4279c399bc 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -8,7 +8,8 @@ class Email < ApplicationRecord
   # TODO: Figure out why sometimes this wasn't released on Topology service
   # when the user is deleted. Some of the operations maybe doesn't trigger
   # before_commit hook. Maybe it's deleted by CASCADE from PostgreSQL?
-  cells_unique_attribute :email
+  cells_unique_attribute :email,
+    type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: EMAIL?
 
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 623ae909b3dc93..870ee7d8822c84 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -7,9 +7,9 @@ class GpgKey < ApplicationRecord
   include ShaAttribute
   include Cells::Unique
 
-  cells_unique_attribute :key
-  cells_unique_attribute :fingerprint
-  cells_unique_attribute :primary_keyid
+  cells_unique_attribute :key, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
+  cells_unique_attribute :fingerprint, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
+  cells_unique_attribute :primary_keyid, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
 
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index f2ab9cc9180ca0..a4c4d575ab08a7 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -6,7 +6,7 @@ class GpgKeySubkey < ApplicationRecord
 
   # TODO: This can be claimed fine, but for deleting, it's not deleted on
   # Topology service for some reason. We need to figure out why
-  # cells_unique_attribute :fingerprint
+  # cells_unique_attribute :fingerprint, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index fa21919865e295..4a7a662d6a6974 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -10,8 +10,8 @@ class Key < ApplicationRecord
   include CreatedAtFilterable
   include Cells::Unique
 
-  cells_unique_attribute :key
-  cells_unique_attribute :fingerprint_sha256
+  cells_unique_attribute :key, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: KEY?
+  cells_unique_attribute :fingerprint_sha256, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: KEY_FINGERPRINT?
 
   sha256_attribute :fingerprint_sha256
 
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 529a040cb5a1a2..25cc742ae429ac 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -10,7 +10,8 @@ class Organization < ApplicationRecord
     include Organizations::Isolatable
     include Cells::Unique
 
-    cells_unique_attribute :path
+    cells_unique_attribute :path,
+      type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: PATH?
 
     DEFAULT_ORGANIZATION_ID = 1
 
diff --git a/app/models/route.rb b/app/models/route.rb
index fec34861261d2b..73555588b7e089 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -6,7 +6,7 @@ class Route < ApplicationRecord
   include EachBatch
   include Cells::Unique
 
-  cells_unique_attribute :path
+  cells_unique_attribute :path, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::ROUTES
 
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
diff --git a/app/models/user.rb b/app/models/user.rb
index 2bbb5990a34769..4cc1dec8179918 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -39,7 +39,7 @@ class User < ApplicationRecord
   include Cells::Unique
 
   # cells_unique_attribute :email # This clashes with emails.email
-  cells_unique_attribute :username
+  cells_unique_attribute :username, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: USERNAME?
 
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
-- 
GitLab


From 8cbe833df352fe52e720f8a8d56f87b738bfa285 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Thu, 4 Sep 2025 16:36:08 +0200
Subject: [PATCH 05/29] Move owner_type and source_table to cells_config

---
 app/models/concerns/cells/unique.rb      | 20 ++++++++++++++++++--
 app/models/email.rb                      |  5 +++--
 app/models/gpg_key.rb                    |  8 +++++---
 app/models/gpg_key_subkey.rb             |  4 +++-
 app/models/key.rb                        |  6 ++++--
 app/models/organizations/organization.rb |  5 +++--
 app/models/route.rb                      |  4 +++-
 app/models/user.rb                       |  4 +++-
 config/initializers/doorkeeper.rb        |  2 +-
 9 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index bee3820e0cc98f..f3efa2a07357b9 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -5,6 +5,9 @@ module Unique
     extend ActiveSupport::Concern
 
     CLAIM_TIMEOUT_MS = 10000
+    BUCKET_TYPE = Gitlab::Cells::TopologyService::ClaimBucket::BucketType
+    OWNER_TYPE = Gitlab::Cells::TopologyService::ClaimOwner::OwnerType
+    SOURCE_TABLE = Gitlab::Cells::TopologyService::ClaimSource::SourceTable
 
     included do
       next unless Gitlab.config.cell.enabled
@@ -21,8 +24,21 @@ def cells_unique_attributes
         @cells_unique_attributes ||= {}
       end
 
-      def cells_unique_attribute(name, args = {})
-        cells_unique_attributes[name] = args
+      def cells_config(owner_type:, source_table:)
+        @cells_owner_type = owner_type
+        @cells_source_table = source_table
+      end
+
+      def cells_owner_type
+        @cells_owner_type
+      end
+
+      def cells_source_table
+        @cells_source_table
+      end
+
+      def cells_unique_attribute(name, type:)
+        cells_unique_attributes[name] = { type: type }
       end
     end
 
diff --git a/app/models/email.rb b/app/models/email.rb
index 3bea4279c399bc..81115d81b9b5ec 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -8,8 +8,9 @@ class Email < ApplicationRecord
   # TODO: Figure out why sometimes this wasn't released on Topology service
   # when the user is deleted. Some of the operations maybe doesn't trigger
   # before_commit hook. Maybe it's deleted by CASCADE from PostgreSQL?
-  cells_unique_attribute :email,
-    type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: EMAIL?
+  cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
+
+  cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
 
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 870ee7d8822c84..fd97c65425b58e 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -7,9 +7,11 @@ class GpgKey < ApplicationRecord
   include ShaAttribute
   include Cells::Unique
 
-  cells_unique_attribute :key, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
-  cells_unique_attribute :fingerprint, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
-  cells_unique_attribute :primary_keyid, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
+  cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED
+  cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
+  cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
+
+  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index a4c4d575ab08a7..a87f26735360c4 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -6,7 +6,9 @@ class GpgKeySubkey < ApplicationRecord
 
   # TODO: This can be claimed fine, but for deleting, it's not deleted on
   # Topology service for some reason. We need to figure out why
-  # cells_unique_attribute :fingerprint, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED
+  # cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
+
+  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index 4a7a662d6a6974..e6f5bb7a417d81 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -10,8 +10,10 @@ class Key < ApplicationRecord
   include CreatedAtFilterable
   include Cells::Unique
 
-  cells_unique_attribute :key, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: KEY?
-  cells_unique_attribute :fingerprint_sha256, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: KEY_FINGERPRINT?
+  cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
+  cells_unique_attribute :fingerprint_sha256, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
+
+  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::ORGANIZATIONS
 
   sha256_attribute :fingerprint_sha256
 
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 25cc742ae429ac..03aa4c97db06d3 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -10,8 +10,9 @@ class Organization < ApplicationRecord
     include Organizations::Isolatable
     include Cells::Unique
 
-    cells_unique_attribute :path,
-      type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: PATH?
+    cells_unique_attribute :path, type: BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
+
+    cells_config owner_type: OWNER_TYPE::ORGANIZATION, source_table: SOURCE_TABLE::ORGANIZATIONS
 
     DEFAULT_ORGANIZATION_ID = 1
 
diff --git a/app/models/route.rb b/app/models/route.rb
index 73555588b7e089..aab42fd78f8de4 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -6,7 +6,9 @@ class Route < ApplicationRecord
   include EachBatch
   include Cells::Unique
 
-  cells_unique_attribute :path, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::ROUTES
+  cells_unique_attribute :path, type: BUCKET_TYPE::ROUTES
+
+  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::ROUTES
 
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
diff --git a/app/models/user.rb b/app/models/user.rb
index 4cc1dec8179918..1afbd0128af9d2 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -39,7 +39,9 @@ class User < ApplicationRecord
   include Cells::Unique
 
   # cells_unique_attribute :email # This clashes with emails.email
-  cells_unique_attribute :username, type: Gitlab::Cells::TopologyService::ClaimBucket::BucketType::UNSPECIFIED # TODO: USERNAME?
+  cells_unique_attribute :username, type: BUCKET_TYPE::UNSPECIFIED # TODO: USERNAME?
+
+  cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
 
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
diff --git a/config/initializers/doorkeeper.rb b/config/initializers/doorkeeper.rb
index cdc3459d1ba0f1..625d56439145c7 100644
--- a/config/initializers/doorkeeper.rb
+++ b/config/initializers/doorkeeper.rb
@@ -157,7 +157,7 @@
   Doorkeeper::Application.class_eval do
     include Cells::Unique
 
-    cells_unique_attribute :uid
+    cells_unique_attribute :uid, type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
   end
 
   # Following doorkeeper monkey patches are to identify the organization on best effort basis
-- 
GitLab


From 6dc17f8648d8c67769e82f0327535477b1b6dc07 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Fri, 5 Sep 2025 14:38:36 +0200
Subject: [PATCH 06/29] Add unique specs

---
 db/docs/cells_outstanding_leases.yml          |   4 +-
 ...2145304_create_cells_outstanding_leases.rb |  12 -
 db/schema_migrations/20250822145304           |   1 -
 db/structure.sql                              |   2 +-
 .../concerns/cells/inspect_changes_spec.rb    | 188 +++++++++++++
 .../cells/transaction_callback_spec.rb        | 253 ++++++++++++++++++
 spec/models/concerns/cells/unique_spec.rb     | 141 ++++++++++
 7 files changed, 585 insertions(+), 16 deletions(-)
 delete mode 100644 db/migrate/20250822145304_create_cells_outstanding_leases.rb
 delete mode 100644 db/schema_migrations/20250822145304
 create mode 100644 spec/models/concerns/cells/inspect_changes_spec.rb
 create mode 100644 spec/models/concerns/cells/transaction_callback_spec.rb
 create mode 100644 spec/models/concerns/cells/unique_spec.rb

diff --git a/db/docs/cells_outstanding_leases.yml b/db/docs/cells_outstanding_leases.yml
index 2f17ceb1f446f7..85e4ead9706cf2 100644
--- a/db/docs/cells_outstanding_leases.yml
+++ b/db/docs/cells_outstanding_leases.yml
@@ -5,7 +5,7 @@ classes:
 feature_categories:
 - cell
 description: Cells outstanding leases for claiming
-introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/200701
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/203567
 milestone: '18.4'
-gitlab_schema: gitlab_main_cell_setting
+gitlab_schema: gitlab_main_cell_local
 table_size: small
diff --git a/db/migrate/20250822145304_create_cells_outstanding_leases.rb b/db/migrate/20250822145304_create_cells_outstanding_leases.rb
deleted file mode 100644
index 1110dee19c4ab2..00000000000000
--- a/db/migrate/20250822145304_create_cells_outstanding_leases.rb
+++ /dev/null
@@ -1,12 +0,0 @@
-# frozen_string_literal: true
-
-class CreateCellsOutstandingLeases < Gitlab::Database::Migration[2.3]
-  milestone '18.4'
-
-  def change
-    create_table :cells_outstanding_leases, id: false do |t| # rubocop:disable Migration/EnsureFactoryForTable -- No factory needed
-      t.primary_key :uuid, :uuid
-      t.timestamps_with_timezone null: false
-    end
-  end
-end
diff --git a/db/schema_migrations/20250822145304 b/db/schema_migrations/20250822145304
deleted file mode 100644
index 4e64198a0f3573..00000000000000
--- a/db/schema_migrations/20250822145304
+++ /dev/null
@@ -1 +0,0 @@
-0acfd536a9af28741684976c1cf1a11362bc6beada54b11545a81dffd01701d0
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index d8b9e5e5c4855a..387820a3fa389b 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -13299,7 +13299,7 @@ CREATE SEQUENCE catalog_verified_namespaces_id_seq
 ALTER SEQUENCE catalog_verified_namespaces_id_seq OWNED BY catalog_verified_namespaces.id;
 
 CREATE TABLE cells_outstanding_leases (
-    uuid uuid DEFAULT gen_random_uuid() NOT NULL,
+    uuid uuid NOT NULL,
     created_at timestamp with time zone NOT NULL,
     updated_at timestamp with time zone NOT NULL
 );
diff --git a/spec/models/concerns/cells/inspect_changes_spec.rb b/spec/models/concerns/cells/inspect_changes_spec.rb
new file mode 100644
index 00000000000000..53e06322e8e1aa
--- /dev/null
+++ b/spec/models/concerns/cells/inspect_changes_spec.rb
@@ -0,0 +1,188 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Cells::InspectChanges, feature_category: :cell do
+  let(:user) { create(:user) }
+  let(:inspector) { described_class.new(user) }
+
+  let(:dummy_class) do
+    Class.new(ActiveRecord::Base) do
+      self.table_name = 'users'
+      include Cells::Unique
+
+      cells_unique_attribute :name, type: :string
+      cells_unique_attribute :email, type: :string
+    end
+  end
+
+  describe '#inspect_changes' do
+    context 'when record is not Unique' do
+      let(:project) { create(:project) }
+      let(:inspector) { described_class.new(project) }
+
+      it 'returns without processing' do
+        expect(Cells::TransactionCallback).not_to receive(:visited)
+        inspector.inspect_changes
+      end
+    end
+
+    context 'when record is Unique but already inspected' do
+      before do
+        visited = { ['users', user.id] => { 'name' => %w[Old New] } }
+        allow(Cells::TransactionCallback).to receive(:visited).and_return(visited)
+      end
+
+      it 'returns without reprocessing' do
+        expect(inspector).not_to receive(:inspect_record)
+        inspector.inspect_changes
+      end
+    end
+
+    context 'when record needs inspection' do
+      before do
+        allow(Cells::TransactionCallback).to receive(:visited).and_return(visited)
+      end
+
+      context 'with transaction' do
+        let(:visited) { {} }
+
+        it 'adds record to visited without calling claim_and_commit' do
+          expect(Cells::Unique).not_to receive(:claim_and_commit)
+
+          inspector.inspect_changes
+
+          expect(visited.keys).to match_array([['users', user.id], ["organizations", user.organization_id]])
+        end
+      end
+
+      context 'without transaction' do
+        let(:visited) { nil }
+
+        it 'calls claim_and_commit' do
+          expect(Cells::Unique).to receive(:claim_and_commit).at_least(:once)
+          inspector.inspect_changes
+        end
+      end
+    end
+  end
+
+  describe 'change tracking' do
+    context 'with associations' do
+      before do
+        allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
+        allow(Cells::Unique).to receive(:claim_and_commit)
+      end
+
+      context 'when loaded associations' do
+        before do
+          create(:email, user: user)
+          create(:email, user: user)
+        end
+
+        it 'recursively inspects associations' do
+          inspector.inspect_changes
+
+          visited = inspector.send(:visited)
+          expect(visited.keys).to include(['users', user.id])
+          # expect(visited.keys).to include(['emails', user.email_ids.first]) TODO
+          # expect(visited.keys).to include(['emails', user.email_ids.second]) TODO
+        end
+      end
+
+      context 'with unloaded associations' do
+        let!(:email) { create(:email, user: user) }
+
+        before do
+          user.reload # Ensure association is not loaded
+        end
+
+        it 'skips unloaded associations' do
+          inspector.inspect_changes
+
+          visited = inspector.send(:visited)
+          expect(visited.keys).to eq([['users', user.id]])
+        end
+      end
+    end
+
+    context 'without associations' do
+      let(:user) { dummy_class.new(name: 'John', email: 'john@example.com') }
+
+      before do
+        allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
+        allow(Cells::Unique).to receive(:claim_and_commit)
+        allow(user).to receive(:destroyed?).and_return(false)
+      end
+
+      context 'with saved changes' do
+        before do
+          allow(user).to receive(:destroyed?).and_return(false)
+          allow(user).to receive_messages(destroyed?: false, saved_changes:
+           { 'name' => %w[Jane John],
+             'email' => ['jane@example.com',
+               'john@example.com'],
+             'age' => [29, 30] })
+        end
+
+        it 'tracks only unique attribute changes' do
+          inspector.inspect_changes
+
+          inspected = inspector.send(:inspected)
+          expect(inspected['name']).to eq(%w[Jane John])
+          expect(inspected['email']).to eq(['jane@example.com', 'john@example.com'])
+          expect(inspected['age']).to be_nil
+        end
+      end
+
+      context 'with destroyed record' do
+        before do
+          allow(user).to receive(:destroyed?).and_return(true)
+        end
+
+        it 'tracks deleted attribute values' do
+          inspector.inspect_changes
+
+          inspected = inspector.send(:inspected)
+          expect(inspected[:name]).to eq(['John'])
+          expect(inspected[:email]).to eq(['john@example.com'])
+        end
+      end
+    end
+  end
+
+  describe 'special cases' do
+    before do
+      allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
+      allow(Cells::Unique).to receive(:claim_and_commit)
+    end
+
+    context 'with GpgKey transformation' do
+      let(:gpg_key) { create(:gpg_key) }
+      let(:inspector) { described_class.new(gpg_key) }
+
+      it 'uses actual stored value after transformation' do
+        expect(gpg_key.fingerprint).to eq(GpgHelpers::User1.fingerprint)
+
+        gpg_key.update!(key: GpgHelpers::User2.public_key)
+
+        inspector.inspect_changes
+
+        inspected = inspector.send(:inspected)
+        expect(inspected['fingerprint']).to match_array(
+          [GpgHelpers::User1.fingerprint.downcase, GpgHelpers::User2.fingerprint.downcase]
+        )
+      end
+    end
+
+    context 'with nil association' do
+      before do
+        user.user_detail # Load nil association
+      end
+
+      it 'handles gracefully' do
+        expect { inspector.inspect_changes }.not_to raise_error
+      end
+    end
+  end
+end
diff --git a/spec/models/concerns/cells/transaction_callback_spec.rb b/spec/models/concerns/cells/transaction_callback_spec.rb
new file mode 100644
index 00000000000000..90d47427ea6fe1
--- /dev/null
+++ b/spec/models/concerns/cells/transaction_callback_spec.rb
@@ -0,0 +1,253 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Cells::TransactionCallback, feature_category: :cell do
+  let(:adapter_class) do
+    Class.new do
+      prepend Cells::TransactionCallback
+
+      attr_reader :transaction_state
+
+      def initialize
+        @transaction_state = []
+      end
+
+      def begin_isolated_db_transaction(...)
+        @transaction_state << :isolated_begin
+      end
+
+      def begin_db_transaction
+        @transaction_state << :begin
+      end
+
+      def commit_db_transaction
+        @transaction_state << :commit
+      end
+
+      def rollback_db_transaction
+        @transaction_state << :rollback
+      end
+    end
+  end
+
+  let(:connection) { adapter_class.new }
+
+  around do |example|
+    original_value = Thread.current.thread_variable_get(described_class::THREAD_VARIABLE_NAME)
+    Thread.current.thread_variable_set(described_class::THREAD_VARIABLE_NAME, nil)
+    example.run
+  ensure
+    Thread.current.thread_variable_set(described_class::THREAD_VARIABLE_NAME, original_value)
+  end
+
+  describe 'thread-local visited stack' do
+    it 'initializes empty stack per thread' do
+      expect(described_class.visited).to be_nil
+      expect(described_class.visited_stack).to eq([])
+    end
+
+    it 'returns first element as visited' do
+      described_class.visited_stack << { first: 'data' }
+      described_class.visited_stack << { second: 'data' }
+
+      expect(described_class.visited).to eq({ first: 'data' })
+    end
+
+    it 'maintains separate stacks per thread' do
+      described_class.visited_stack << { main: 'thread' }
+
+      other_thread_value = Thread.new do
+        described_class.visited_stack << { other: 'thread' }
+        described_class.visited
+      end.value
+
+      expect(described_class.visited).to eq({ main: 'thread' })
+      expect(other_thread_value).to eq({ other: 'thread' })
+    end
+  end
+
+  describe 'transaction lifecycle' do
+    describe '#begin_db_transaction' do
+      it 'pushes empty hash to stack' do
+        connection.begin_db_transaction
+
+        expect(described_class.visited_stack).to eq([{}])
+        expect(connection.transaction_state).to include(:begin)
+      end
+
+      it 'handles nested transactions' do
+        connection.begin_db_transaction
+        connection.begin_db_transaction
+
+        expect(described_class.visited_stack.size).to eq(2)
+        expect(described_class.visited_stack).to eq([{}, {}])
+      end
+    end
+
+    describe '#begin_isolated_db_transaction' do
+      it 'pushes empty hash to stack' do
+        connection.begin_isolated_db_transaction
+
+        expect(described_class.visited_stack).to eq([{}])
+        expect(connection.transaction_state).to include(:isolated_begin)
+      end
+    end
+
+    describe '#commit_db_transaction' do
+      before do
+        connection.begin_db_transaction
+        described_class.visited_stack.last[['users', 1]] = { 'email' => %w[old new] }
+      end
+
+      it 'claims and commits with visited data' do
+        expect(Cells::Unique).to receive(:claim_and_commit)
+          .with({ ['users', 1] => { 'email' => %w[old new] } })
+          .and_yield
+
+        connection.commit_db_transaction
+
+        expect(described_class.visited_stack).to be_empty
+        expect(connection.transaction_state).to include(:commit)
+      end
+
+      it 'maintains stack if commit fails' do
+        connection.begin_db_transaction
+        described_class.visited_stack.last[['users', 1]] = { 'email' => %w[old new] }
+
+        expect(Cells::Unique).to receive(:claim_and_commit).and_yield
+
+        allow(connection).to receive(:commit_db_transaction).and_wrap_original do |original_method, *args|
+          original_method.call(*args)
+          raise ActiveRecord::StatementInvalid, 'Connection lost'
+        end
+
+        expect do
+          connection.commit_db_transaction
+        end.to raise_error(ActiveRecord::StatementInvalid)
+
+        expect(described_class.visited_stack.size).to eq(1)
+      end
+
+      it 'handles empty visited data' do
+        described_class.visited_stack.last.clear
+
+        expect(Cells::Unique).to receive(:claim_and_commit).with({}).and_yield
+
+        connection.commit_db_transaction
+        expect(described_class.visited_stack).to be_empty
+      end
+    end
+
+    describe '#rollback_db_transaction' do
+      context 'with balanced transaction' do
+        before do
+          connection.begin_db_transaction
+        end
+
+        it 'rolls back and pops stack' do
+          expect(Cells::Unique).to receive(:rollback).and_yield
+
+          result = connection.rollback_db_transaction
+
+          expect(result).to be true
+          expect(described_class.visited_stack).to be_empty
+          expect(connection.transaction_state).to include(:rollback)
+        end
+      end
+
+      context 'with unbalanced rollback' do
+        it 'handles rollback without corresponding begin' do
+          expect(Cells::Unique).to receive(:rollback).and_yield
+
+          result = connection.rollback_db_transaction
+
+          expect(result).to be false
+          expect(described_class.visited_stack).to be_empty
+          expect(connection.transaction_state).to include(:rollback)
+        end
+      end
+
+      context 'with nested transactions' do
+        before do
+          connection.begin_db_transaction
+          connection.begin_db_transaction
+          described_class.visited_stack[0][['users', 1]] = 'outer'
+          described_class.visited_stack[1][['users', 2]] = 'inner'
+        end
+
+        it 'only rolls back innermost transaction' do
+          expect(described_class.visited_stack.size).to eq(2)
+          expect(Cells::Unique).to receive(:rollback).and_yield
+
+          connection.rollback_db_transaction
+
+          expect(described_class.visited_stack.size).to eq(1)
+          expect(described_class.visited_stack.first).to eq({ ['users', 1] => 'outer' })
+        end
+      end
+    end
+  end
+
+  describe 'integration scenarios' do
+    it 'handles successful transaction flow' do
+      connection.begin_db_transaction
+
+      described_class.visited[['users', 1]] = { 'name' => %w[old new] }
+
+      expect(Cells::Unique).to receive(:claim_and_commit)
+        .with({ ['users', 1] => { 'name' => %w[old new] } })
+        .and_yield
+
+      connection.commit_db_transaction
+
+      expect(described_class.visited_stack).to be_empty
+    end
+
+    it 'handles rollback after partial work' do
+      connection.begin_db_transaction
+
+      described_class.visited[['users', 1]] = { 'name' => %w[old new] }
+
+      expect(Cells::Unique).to receive(:rollback).and_yield
+
+      connection.rollback_db_transaction
+
+      expect(described_class.visited_stack).to be_empty
+    end
+
+    it 'handles nested transaction with inner rollback' do
+      connection.begin_db_transaction
+      described_class.visited[['outer', 1]] = 'data'
+
+      connection.begin_isolated_db_transaction
+      described_class.visited_stack.last[['inner', 1]] = 'data'
+
+      expect(Cells::Unique).to receive(:rollback).and_yield
+      connection.rollback_db_transaction
+
+      expect(Cells::Unique).to receive(:claim_and_commit)
+        .with({ ['outer', 1] => 'data' })
+        .and_yield
+
+      connection.commit_db_transaction
+
+      expect(described_class.visited_stack).to be_empty
+    end
+  end
+
+  describe 'error handling' do
+    it 'maintains consistency when claim service fails' do
+      connection.begin_db_transaction
+      described_class.visited[['users', 1]] = 'data'
+
+      expect(Cells::Unique).to receive(:claim_and_commit)
+        .and_raise(StandardError, 'Claim service unavailable')
+
+      expect { connection.commit_db_transaction }
+        .to raise_error(StandardError, 'Claim service unavailable')
+
+      expect(described_class.visited_stack.size).to eq(1)
+    end
+  end
+end
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
new file mode 100644
index 00000000000000..83e809381ea0c5
--- /dev/null
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -0,0 +1,141 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Cells::Unique, feature_category: :cell do
+  let(:dummy_class) do
+    Class.new(ActiveRecord::Base) do
+      self.table_name = 'users'
+      include Cells::Unique
+    end
+  end
+
+  describe 'configuration' do
+    it 'sets and retrieves cell configuration' do
+      dummy_class.cells_config(owner_type: :user, source_table: :users)
+      dummy_class.cells_unique_attribute(:email, type: :string)
+      dummy_class.cells_unique_attribute(:username, type: :string)
+
+      expect(dummy_class.cells_owner_type).to eq(:user)
+      expect(dummy_class.cells_source_table).to eq(:users)
+      expect(dummy_class.cells_unique_attributes).to eq(
+        email: { type: :string },
+        username: { type: :string }
+      )
+    end
+  end
+
+  describe '.claim_and_commit' do
+    let(:claim_service) { instance_double(Gitlab::TopologyServiceClient::ClaimService) }
+
+    before do
+      allow(Gitlab::TopologyServiceClient::ClaimService).to receive(:new).and_return(claim_service)
+    end
+
+    it 'yields directly when changes are empty' do
+      expect(claim_service).not_to receive(:claim_and_commit)
+
+      yielded = false
+      described_class.claim_and_commit({}) { yielded = true }
+
+      expect(yielded).to be true
+    end
+
+    it 'calls claim service when changes exist' do
+      changes = { ['users', 1] => { 'email' => %w[old new] } }
+      expect(claim_service).to receive(:claim_and_commit).with(changes)
+
+      described_class.claim_and_commit(changes)
+    end
+  end
+
+  describe '.rollback' do
+    let(:claim_service) { instance_double(Gitlab::TopologyServiceClient::ClaimService) }
+
+    it 'delegates to claim service' do
+      allow(Gitlab::TopologyServiceClient::ClaimService).to receive(:new).and_return(claim_service)
+      expect(claim_service).to receive(:rollback)
+
+      described_class.rollback
+    end
+  end
+
+  describe 'callbacks' do
+    let(:instance) { dummy_class.new }
+    let(:connection) { instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) }
+
+    context 'when cells are enabled' do
+      before do
+        allow(Gitlab.config.cell).to receive(:enabled).and_return(true)
+        allow(instance).to receive(:connection).and_return(connection)
+      end
+
+      describe 'before_commit callback' do
+        it 'inspects changes with timeout' do
+          inspector = instance_double(Cells::InspectChanges)
+
+          expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 10000")
+          expect(Cells::InspectChanges).to receive(:new).with(instance).and_return(inspector)
+          expect(inspector).to receive(:inspect_changes)
+
+          instance.run_callbacks(:before_commit)
+        end
+      end
+
+      describe 'instance methods' do
+        describe '#with_claim_timeout (private)' do
+          it 'sets the statement timeout to default value' do
+            expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 10000")
+
+            instance.send(:with_claim_timeout) { :result }
+          end
+
+          it 'sets the statement timeout to custom value' do
+            expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 5000")
+
+            instance.send(:with_claim_timeout, 5000) { :result }
+          end
+
+          it 'yields to the block' do
+            allow(connection).to receive(:exec_query)
+
+            block_called = false
+            result = instance.send(:with_claim_timeout) do
+              block_called = true
+              :block_result
+            end
+
+            expect(block_called).to be true
+            expect(result).to eq(:block_result)
+          end
+
+          it 'executes query before yielding' do
+            call_order = []
+
+            expect(connection).to receive(:exec_query) do |query|
+              call_order << :exec_query
+              expect(query).to eq("SET LOCAL statement_timeout = 10000")
+            end
+
+            instance.send(:with_claim_timeout) do
+              call_order << :block
+            end
+
+            expect(call_order).to eq([:exec_query, :block])
+          end
+        end
+      end
+    end
+
+    context 'when cells are disabled' do
+      before do
+        allow(Gitlab.config.cell).to receive(:enabled).and_return(false)
+      end
+
+      it 'does not inspect changes' do
+        expect(Cells::InspectChanges).not_to receive(:new)
+        instance.run_callbacks(:commit)
+      end
+    end
+  end
+end
-- 
GitLab


From f214ccb56922071a2595e63e47865e5ae598345b Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Mon, 15 Sep 2025 14:16:40 +0200
Subject: [PATCH 07/29] Adds specs for unique attributes

---
 app/models/email.rb                             |  3 ---
 app/models/gpg_key_subkey.rb                    |  5 +----
 spec/models/email_spec.rb                       |  5 +++++
 spec/models/gpg_key_spec.rb                     |  5 +++++
 spec/models/gpg_key_subkey_spec.rb              |  5 +++++
 spec/models/key_spec.rb                         |  5 +++++
 spec/models/organizations/organization_spec.rb  |  5 +++++
 spec/models/route_spec.rb                       |  5 +++++
 spec/models/user_spec.rb                        |  5 +++++
 .../cells/unique_shared_examples.rb             | 17 +++++++++++++++++
 10 files changed, 53 insertions(+), 7 deletions(-)
 create mode 100644 spec/support/shared_examples/cells/unique_shared_examples.rb

diff --git a/app/models/email.rb b/app/models/email.rb
index 81115d81b9b5ec..7edbd89dc21dc7 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -5,9 +5,6 @@ class Email < ApplicationRecord
   include Gitlab::SQL::Pattern
   include Cells::Unique
 
-  # TODO: Figure out why sometimes this wasn't released on Topology service
-  # when the user is deleted. Some of the operations maybe doesn't trigger
-  # before_commit hook. Maybe it's deleted by CASCADE from PostgreSQL?
   cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
   cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index a87f26735360c4..03cdff865730e3 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -4,10 +4,7 @@ class GpgKeySubkey < ApplicationRecord
   include ShaAttribute
   include Cells::Unique
 
-  # TODO: This can be claimed fine, but for deleting, it's not deleted on
-  # Topology service for some reason. We need to figure out why
-  # cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
-
+  cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
   cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha_attribute :keyid
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index bdc4945bc729e0..e4d820552a7df5 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -5,6 +5,11 @@
 RSpec.describe Email do
   let_it_be(:user) { create(:user) }
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::USER,
+    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    unique_attributes: [:email]
+
   describe 'modules' do
     subject { described_class }
 
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index fbd7db170ce127..b71b3bf4b7eb7f 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -5,6 +5,11 @@
 RSpec.describe GpgKey, feature_category: :source_code_management do
   subject { build(:gpg_key) }
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    unique_attributes: [:key, :fingerprint, :primary_keyid]
+
   describe "associations" do
     it { is_expected.to belong_to(:user) }
     it { is_expected.to have_many(:subkeys) }
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index c1d9e2bde43bd9..295f0712cec532 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -14,4 +14,9 @@
     it { is_expected.to validate_presence_of(:fingerprint) }
     it { is_expected.to validate_presence_of(:keyid) }
   end
+
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index 1701bbd17ad108..66ca553c308736 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -5,6 +5,11 @@
 RSpec.describe Key, :mailer, feature_category: :system_access do
   it_behaves_like 'having unique enum values'
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
+    unique_attributes: [:key, :fingerprint_sha256]
+
   describe "Associations" do
     it { is_expected.to belong_to(:user) }
     it { is_expected.to have_many(:todos).dependent(:destroy) }
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index d0c79229379737..afde6944e25b33 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -5,6 +5,11 @@
 RSpec.describe Organizations::Organization, type: :model, feature_category: :organization do
   let_it_be_with_refind(:organization) { create(:organization) }
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::ORGANIZATION,
+    source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
+    unique_attributes: [:path]
+
   describe 'associations' do
     it { is_expected.to have_one(:organization_detail).inverse_of(:organization).autosave(true) }
 
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index 16bb8fe9b3f340..2207e72b2d8745 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -6,6 +6,11 @@
   let(:group) { create(:group, path: 'git_lab', name: 'git_lab') }
   let(:route) { group.route }
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    source_table: Cells::Unique::SOURCE_TABLE::ROUTES,
+    unique_attributes: [:path]
+
   describe 'relationships' do
     it { is_expected.to belong_to(:source) }
     it { is_expected.to belong_to(:namespace) }
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index 07d0f2de9b8387..e7194c5361aff2 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -49,6 +49,11 @@
 
   it_behaves_like 'having unique enum values'
 
+  it_behaves_like 'cells unique model',
+    owner_type: Cells::Unique::OWNER_TYPE::USER,
+    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    unique_attributes: [:username]
+
   describe 'modules' do
     subject { described_class }
 
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
new file mode 100644
index 00000000000000..f5f152f52e705a
--- /dev/null
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+RSpec.shared_examples 'cells unique model' do |owner_type:, source_table:, unique_attributes:|
+  it 'has the expected owner_type' do
+    expect(described_class.cells_owner_type).to eq(owner_type)
+  end
+
+  it 'has the expected source_table' do
+    expect(described_class.cells_source_table).to eq(source_table)
+  end
+
+  it 'has the expected unique attributes' do
+    unique_attributes.each do |attr_name|
+      expect(described_class.cells_unique_attributes).to have_key(attr_name)
+    end
+  end
+end
-- 
GitLab


From 8ac282d960e97d1ca6375421d2839fcd56869144 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Thu, 25 Sep 2025 23:20:21 +0800
Subject: [PATCH 08/29] Use after_save and after_destroy so we avoid
 before_commit

Which isn't really a public API so we try to avoid relying on it.
---
 app/models/concerns/cells/unique.rb | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index f3efa2a07357b9..05b9370ce35ce2 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -12,7 +12,13 @@ module Unique
     included do
       next unless Gitlab.config.cell.enabled
 
-      before_commit do
+      after_save do
+        with_claim_timeout do
+          InspectChanges.new(self).inspect_changes
+        end
+      end
+
+      after_destroy do
         with_claim_timeout do
           InspectChanges.new(self).inspect_changes
         end
-- 
GitLab


From f1f620feac273f512b41cfb5057aca87c2e0f9ef Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Fri, 26 Sep 2025 01:44:08 +0800
Subject: [PATCH 09/29] Define cells_unique_has_many so we can unclaim for
 destroy

Even for records that would be deleted on cascade.
---
 app/models/concerns/cells/inspect_changes.rb | 41 +++++++++++++++++---
 app/models/concerns/cells/unique.rb          | 24 +++++++-----
 app/models/gpg_key.rb                        |  2 +
 app/models/user.rb                           |  2 +
 4 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
index 876360d29e0933..8a9cffe0119052 100644
--- a/app/models/concerns/cells/inspect_changes.rb
+++ b/app/models/concerns/cells/inspect_changes.rb
@@ -4,8 +4,9 @@ module Cells
   class InspectChanges
     attr_reader :record
 
-    def initialize(new_record)
+    def initialize(new_record, destroy: false)
       @record = new_record
+      @destroy = destroy
     end
 
     def inspect_changes
@@ -46,13 +47,17 @@ def prepare_inspecting
     end
 
     def inspect_record
-      if record.destroyed?
+      if destroying?
         inspect_deleted_attributes
       else
         inspect_saved_changes
       end
     end
 
+    def destroying?
+      @destroy
+    end
+
     def inspect_deleted_attributes
       record.class.cells_unique_attributes.each_key do |column|
         inspected[column] = [record.attributes[column.to_s]]
@@ -71,15 +76,41 @@ def inspect_saved_changes
     end
 
     def inspect_associations
+      # We know the following associations will be deleted yet the hooks
+      # won't be called, so we laod it at once and inspect the changes
+      # as if they would get changed.
+      preload_nested_associations(record) if destroying?
+
       record.class.reflect_on_all_associations.each do |reflection|
         association_proxy = record.association(reflection.name)
 
         if association_proxy.loaded? # rubocop:disable Style/Next -- I think this is more clear
-          Array.wrap(association_proxy.target).each do |association_record|
-            InspectChanges.new(association_record).inspect_changes
-          end
+          inspect_loaded_associations(association_proxy)
         end
       end
     end
+
+    def inspect_loaded_associations(association_proxy)
+      Array.wrap(association_proxy.target).each do |association_record|
+        InspectChanges.new(association_record, destroy: destroying?).inspect_changes
+      end
+    end
+
+    def preload_nested_associations(record)
+      ActiveRecord::Associations::Preloader.new(
+        records: [record],
+        associations: lookup_associations(record.class)
+      ).call
+    end
+
+    def lookup_associations(model)
+      model.cells_unique_has_many_associations.map do |name|
+        { name => lookup_associations(reflect_association_model(model, name)) }
+      end
+    end
+
+    def reflect_association_model(model, name)
+      model.reflect_on_association(name).class_name.constantize
+    end
   end
 end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 05b9370ce35ce2..d98cd04c24383d 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -18,40 +18,44 @@ module Unique
         end
       end
 
-      after_destroy do
+      before_destroy do
         with_claim_timeout do
-          InspectChanges.new(self).inspect_changes
+          InspectChanges.new(self, destroy: true).inspect_changes
         end
       end
     end
 
     class_methods do
+      attr_reader :cells_owner_type, :cells_source_table
+
       def cells_unique_attributes
         @cells_unique_attributes ||= {}
       end
 
+      def cells_unique_has_many_associations
+        @cells_unique_has_many_associations ||= []
+      end
+
       def cells_config(owner_type:, source_table:)
         @cells_owner_type = owner_type
         @cells_source_table = source_table
       end
 
-      def cells_owner_type
-        @cells_owner_type
-      end
-
-      def cells_source_table
-        @cells_source_table
-      end
-
       def cells_unique_attribute(name, type:)
         cells_unique_attributes[name] = { type: type }
       end
+
+      def cells_unique_has_many(name)
+        cells_unique_has_many_associations << name
+      end
     end
 
     def self.claim_and_commit(changes)
       if changes.empty?
         yield if block_given?
       else
+        pp changes
+
         claim_client = Gitlab::TopologyServiceClient::ClaimService.new
         claim_client.claim_and_commit(changes) do
           yield if block_given?
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index fd97c65425b58e..4fa06af9af70a0 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -13,6 +13,8 @@ class GpgKey < ApplicationRecord
 
   cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
 
+  cells_unique_has_many :subkeys
+
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
 
diff --git a/app/models/user.rb b/app/models/user.rb
index 1afbd0128af9d2..af602ecaebb0db 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -40,6 +40,8 @@ class User < ApplicationRecord
 
   # cells_unique_attribute :email # This clashes with emails.email
   cells_unique_attribute :username, type: BUCKET_TYPE::UNSPECIFIED # TODO: USERNAME?
+  cells_unique_has_many :gpg_keys
+  cells_unique_has_many :emails
 
   cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
 
-- 
GitLab


From 309d23a87f532387df72fd23b99f87cc2e8c8a21 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Mon, 29 Sep 2025 16:11:20 +0200
Subject: [PATCH 10/29] Add Unique claims behind FF

---
 .../gitlab/feature_flag_without_actor.yml     |   1 +
 app/models/concerns/cells/unique.rb           |  29 ++---
 .../beta/cells_unique_claims.yml              |  10 ++
 spec/models/concerns/cells/unique_spec.rb     | 103 +++++++-----------
 4 files changed, 63 insertions(+), 80 deletions(-)
 create mode 100644 config/feature_flags/beta/cells_unique_claims.yml

diff --git a/.rubocop_todo/gitlab/feature_flag_without_actor.yml b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
index 35ed8033b0fb09..53f6925c8b727d 100644
--- a/.rubocop_todo/gitlab/feature_flag_without_actor.yml
+++ b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
@@ -21,6 +21,7 @@ Gitlab/FeatureFlagWithoutActor:
     - 'app/models/clusters/instance.rb'
     - 'app/models/concerns/counter_attribute.rb'
     - 'app/models/concerns/reset_on_column_errors.rb'
+    - 'app/models/concerns/cells/unique.rb'
     - 'app/models/concerns/web_hooks/auto_disabling.rb'
     - 'app/models/merge_request.rb'
     - 'app/models/namespace.rb'
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index d98cd04c24383d..911649f2bfc518 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -4,7 +4,6 @@ module Cells
   module Unique
     extend ActiveSupport::Concern
 
-    CLAIM_TIMEOUT_MS = 10000
     BUCKET_TYPE = Gitlab::Cells::TopologyService::ClaimBucket::BucketType
     OWNER_TYPE = Gitlab::Cells::TopologyService::ClaimOwner::OwnerType
     SOURCE_TABLE = Gitlab::Cells::TopologyService::ClaimSource::SourceTable
@@ -12,17 +11,8 @@ module Unique
     included do
       next unless Gitlab.config.cell.enabled
 
-      after_save do
-        with_claim_timeout do
-          InspectChanges.new(self).inspect_changes
-        end
-      end
-
-      before_destroy do
-        with_claim_timeout do
-          InspectChanges.new(self, destroy: true).inspect_changes
-        end
-      end
+      after_save inspect_save_changes, if: -> { Feature.enabled?(:cells_unique_claims) }
+      before_destroy inspect_destroy_changes, if: -> { Feature.enabled?(:cells_unique_claims) }
     end
 
     class_methods do
@@ -48,6 +38,14 @@ def cells_unique_attribute(name, type:)
       def cells_unique_has_many(name)
         cells_unique_has_many_associations << name
       end
+
+      def inspect_save_changes
+        InspectChanges.new(self).inspect_changes
+      end
+
+      def inspect_destroy_changes
+        InspectChanges.new(self, destroy: true).inspect_changes
+      end
     end
 
     def self.claim_and_commit(changes)
@@ -69,12 +67,5 @@ def self.rollback
         yield
       end
     end
-
-    private
-
-    def with_claim_timeout(timeout_ms = CLAIM_TIMEOUT_MS)
-      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
-      yield
-    end
   end
 end
diff --git a/config/feature_flags/beta/cells_unique_claims.yml b/config/feature_flags/beta/cells_unique_claims.yml
new file mode 100644
index 00000000000000..a32bd5090bf770
--- /dev/null
+++ b/config/feature_flags/beta/cells_unique_claims.yml
@@ -0,0 +1,10 @@
+---
+name: cells_unique_claims
+description:
+feature_issue_url: https://gitlab.com/gitlab-com/gl-infra/tenant-scale/cells-infrastructure/team/-/issues/365
+introduced_by_url:
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/572819
+milestone: '18.5'
+group: group::cells infrastructure
+type: beta
+default_enabled: false
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index 83e809381ea0c5..60cebbbf117bf9 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -3,25 +3,17 @@
 require 'spec_helper'
 
 RSpec.describe Cells::Unique, feature_category: :cell do
-  let(:dummy_class) do
-    Class.new(ActiveRecord::Base) do
-      self.table_name = 'users'
-      include Cells::Unique
-    end
-  end
+  let(:test_klass) { ::User }
 
   describe 'configuration' do
     it 'sets and retrieves cell configuration' do
-      dummy_class.cells_config(owner_type: :user, source_table: :users)
-      dummy_class.cells_unique_attribute(:email, type: :string)
-      dummy_class.cells_unique_attribute(:username, type: :string)
-
-      expect(dummy_class.cells_owner_type).to eq(:user)
-      expect(dummy_class.cells_source_table).to eq(:users)
-      expect(dummy_class.cells_unique_attributes).to eq(
-        email: { type: :string },
-        username: { type: :string }
-      )
+      test_klass.cells_config(owner_type: :user, source_table: :users)
+      test_klass.cells_unique_attribute(:email, type: :string)
+      test_klass.cells_unique_attribute(:username, type: :string)
+
+      expect(test_klass.cells_owner_type).to eq(:user)
+      expect(test_klass.cells_source_table).to eq(:users)
+      expect(test_klass.cells_unique_attributes).to eq(email: { type: :string }, username: { type: :string })
     end
   end
 
@@ -61,68 +53,52 @@
   end
 
   describe 'callbacks' do
-    let(:instance) { dummy_class.new }
-    let(:connection) { instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) }
+    let(:instance) { build(:user) }
 
     context 'when cells are enabled' do
       before do
         allow(Gitlab.config.cell).to receive(:enabled).and_return(true)
-        allow(instance).to receive(:connection).and_return(connection)
       end
 
-      describe 'before_commit callback' do
-        it 'inspects changes with timeout' do
+      describe 'after_save callback' do
+        it 'calls inspect_save_changes after saving' do
           inspector = instance_double(Cells::InspectChanges)
 
-          expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 10000")
-          expect(Cells::InspectChanges).to receive(:new).with(instance).and_return(inspector)
+          expect(Cells::InspectChanges).to receive(:new).with(instance).at_least(:once).and_return(inspector)
           expect(inspector).to receive(:inspect_changes)
 
-          instance.run_callbacks(:before_commit)
+          instance.save!
         end
       end
 
-      describe 'instance methods' do
-        describe '#with_claim_timeout (private)' do
-          it 'sets the statement timeout to default value' do
-            expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 10000")
-
-            instance.send(:with_claim_timeout) { :result }
-          end
-
-          it 'sets the statement timeout to custom value' do
-            expect(connection).to receive(:exec_query).with("SET LOCAL statement_timeout = 5000")
-
-            instance.send(:with_claim_timeout, 5000) { :result }
-          end
-
-          it 'yields to the block' do
-            allow(connection).to receive(:exec_query)
+      describe 'before_destroy callback' do
+        before do
+          instance.save!
+        end
 
-            block_called = false
-            result = instance.send(:with_claim_timeout) do
-              block_called = true
-              :block_result
-            end
+        it 'calls inspect_destroy_changes before destroying' do
+          inspector = instance_double(Cells::InspectChanges)
 
-            expect(block_called).to be true
-            expect(result).to eq(:block_result)
-          end
+          expect(Cells::InspectChanges).to receive(:new).with(instance, destroy: true).and_return(inspector)
+          expect(inspector).to receive(:inspect_changes)
 
-          it 'executes query before yielding' do
-            call_order = []
+          instance.destroy!
+        end
+      end
 
-            expect(connection).to receive(:exec_query) do |query|
-              call_order << :exec_query
-              expect(query).to eq("SET LOCAL statement_timeout = 10000")
-            end
+      context 'when feature flag is disabled' do
+        before do
+          stub_feature_flags(cells_unique_claims: false)
+        end
 
-            instance.send(:with_claim_timeout) do
-              call_order << :block
-            end
+        it 'does not call inspect_save_changes on save' do
+          expect(Cells::InspectChanges).not_to receive(:new)
+          instance.save!
+        end
 
-            expect(call_order).to eq([:exec_query, :block])
-          end
+        it 'does not call inspect_destroy_changes on destroy' do
+          expect(Cells::InspectChanges).not_to receive(:new)
+          instance.destroy!
         end
       end
     end
@@ -132,9 +108,14 @@
         allow(Gitlab.config.cell).to receive(:enabled).and_return(false)
       end
 
-      it 'does not inspect changes' do
+      it 'does not call inspect_save_changes on save' do
+        expect(Cells::InspectChanges).not_to receive(:new)
+        instance.save!
+      end
+
+      it 'does not call inspect_destroy_changes on destroy' do
         expect(Cells::InspectChanges).not_to receive(:new)
-        instance.run_callbacks(:commit)
+        instance.destroy!
       end
     end
   end
-- 
GitLab


From 17a5142b9f3358d858eda2927cec15528bde261d Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Tue, 30 Sep 2025 13:40:45 +0200
Subject: [PATCH 11/29] Fixes failing pipeline

---
 app/models/concerns/cells/unique.rb       | 30 ++++++++++++++---------
 spec/models/concerns/cells/unique_spec.rb | 23 +++++++++--------
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 911649f2bfc518..db8c3c84f39b7d 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -9,10 +9,8 @@ module Unique
     SOURCE_TABLE = Gitlab::Cells::TopologyService::ClaimSource::SourceTable
 
     included do
-      next unless Gitlab.config.cell.enabled
-
-      after_save inspect_save_changes, if: -> { Feature.enabled?(:cells_unique_claims) }
-      before_destroy inspect_destroy_changes, if: -> { Feature.enabled?(:cells_unique_claims) }
+      after_save :inspect_save_changes
+      before_destroy :inspect_destroy_changes
     end
 
     class_methods do
@@ -38,14 +36,6 @@ def cells_unique_attribute(name, type:)
       def cells_unique_has_many(name)
         cells_unique_has_many_associations << name
       end
-
-      def inspect_save_changes
-        InspectChanges.new(self).inspect_changes
-      end
-
-      def inspect_destroy_changes
-        InspectChanges.new(self, destroy: true).inspect_changes
-      end
     end
 
     def self.claim_and_commit(changes)
@@ -67,5 +57,21 @@ def self.rollback
         yield
       end
     end
+
+    private
+
+    def inspect_save_changes
+      return unless Gitlab.config.cell.enabled
+      return unless Feature.enabled?(:cells_unique_claims)
+
+      InspectChanges.new(self).inspect_changes
+    end
+
+    def inspect_destroy_changes
+      return unless Gitlab.config.cell.enabled
+      return unless Feature.enabled?(:cells_unique_claims)
+
+      InspectChanges.new(self, destroy: true).inspect_changes
+    end
   end
 end
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index 60cebbbf117bf9..d4121c29301125 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -3,6 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Cells::Unique, feature_category: :cell do
+  let_it_be(:user) { create(:user) }
   let(:test_klass) { ::User }
 
   describe 'configuration' do
@@ -62,27 +63,25 @@
 
       describe 'after_save callback' do
         it 'calls inspect_save_changes after saving' do
-          inspector = instance_double(Cells::InspectChanges)
+          email_inspector = instance_double(Cells::InspectChanges)
+          user_inspector = instance_double(Cells::InspectChanges)
 
-          expect(Cells::InspectChanges).to receive(:new).with(instance).at_least(:once).and_return(inspector)
-          expect(inspector).to receive(:inspect_changes)
+          expect(Cells::InspectChanges).to receive(:new).with(an_instance_of(Email)).and_return(email_inspector)
+          expect(Cells::InspectChanges).to receive(:new).with(instance).and_return(user_inspector)
+          expect(email_inspector).to receive(:inspect_changes)
+          expect(user_inspector).to receive(:inspect_changes)
 
           instance.save!
         end
       end
 
       describe 'before_destroy callback' do
-        before do
-          instance.save!
-        end
-
         it 'calls inspect_destroy_changes before destroying' do
           inspector = instance_double(Cells::InspectChanges)
 
-          expect(Cells::InspectChanges).to receive(:new).with(instance, destroy: true).and_return(inspector)
+          expect(Cells::InspectChanges).to receive(:new).with(user, destroy: true).and_return(inspector)
           expect(inspector).to receive(:inspect_changes)
-
-          instance.destroy!
+          user.destroy!
         end
       end
 
@@ -98,7 +97,7 @@
 
         it 'does not call inspect_destroy_changes on destroy' do
           expect(Cells::InspectChanges).not_to receive(:new)
-          instance.destroy!
+          user.destroy!
         end
       end
     end
@@ -115,7 +114,7 @@
 
       it 'does not call inspect_destroy_changes on destroy' do
         expect(Cells::InspectChanges).not_to receive(:new)
-        instance.destroy!
+        user.destroy!
       end
     end
   end
-- 
GitLab


From 8e7b4d85e619e207bc05b08ac0102e29303f100a Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Tue, 30 Sep 2025 17:00:52 +0200
Subject: [PATCH 12/29] Update FF url

---
 config/feature_flags/beta/cells_unique_claims.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/config/feature_flags/beta/cells_unique_claims.yml b/config/feature_flags/beta/cells_unique_claims.yml
index a32bd5090bf770..da3d94089711b8 100644
--- a/config/feature_flags/beta/cells_unique_claims.yml
+++ b/config/feature_flags/beta/cells_unique_claims.yml
@@ -2,7 +2,7 @@
 name: cells_unique_claims
 description:
 feature_issue_url: https://gitlab.com/gitlab-com/gl-infra/tenant-scale/cells-infrastructure/team/-/issues/365
-introduced_by_url:
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/200701
 rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/572819
 milestone: '18.5'
 group: group::cells infrastructure
-- 
GitLab


From d5e8c02ff5eaa1a697d5738f9dd7861d334c3169 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Thu, 2 Oct 2025 23:25:58 +0800
Subject: [PATCH 13/29] Use `let_it_be`

Co-authored-by: DANGER_GITLAB_API_TOKEN <project_278964_bot_d10f7ffdb05c889168416678645d29bb@noreply.gitlab.com>
---
 spec/models/concerns/cells/inspect_changes_spec.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spec/models/concerns/cells/inspect_changes_spec.rb b/spec/models/concerns/cells/inspect_changes_spec.rb
index 53e06322e8e1aa..ca9a0e7bea6cc3 100644
--- a/spec/models/concerns/cells/inspect_changes_spec.rb
+++ b/spec/models/concerns/cells/inspect_changes_spec.rb
@@ -18,7 +18,7 @@
 
   describe '#inspect_changes' do
     context 'when record is not Unique' do
-      let(:project) { create(:project) }
+      let_it_be(:project) { create(:project) }
       let(:inspector) { described_class.new(project) }
 
       it 'returns without processing' do
-- 
GitLab


From 4337b19c67b5c780228ffefcf9d8a714ba14afb5 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Sat, 4 Oct 2025 00:22:09 +0800
Subject: [PATCH 14/29] Ensure feature flag cells_unique_claims is consistent
 across Current

---
 .rubocop_todo/gitlab/feature_flag_without_actor.yml | 1 -
 app/models/concerns/cells/transaction_callback.rb   | 8 ++++++++
 app/models/concerns/cells/unique.rb                 | 4 ++--
 app/models/current.rb                               | 7 +++++++
 spec/models/concerns/cells/unique_spec.rb           | 1 +
 5 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/.rubocop_todo/gitlab/feature_flag_without_actor.yml b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
index 53f6925c8b727d..35ed8033b0fb09 100644
--- a/.rubocop_todo/gitlab/feature_flag_without_actor.yml
+++ b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
@@ -21,7 +21,6 @@ Gitlab/FeatureFlagWithoutActor:
     - 'app/models/clusters/instance.rb'
     - 'app/models/concerns/counter_attribute.rb'
     - 'app/models/concerns/reset_on_column_errors.rb'
-    - 'app/models/concerns/cells/unique.rb'
     - 'app/models/concerns/web_hooks/auto_disabling.rb'
     - 'app/models/merge_request.rb'
     - 'app/models/namespace.rb'
diff --git a/app/models/concerns/cells/transaction_callback.rb b/app/models/concerns/cells/transaction_callback.rb
index 69047adbeac274..736f01cdf2326b 100644
--- a/app/models/concerns/cells/transaction_callback.rb
+++ b/app/models/concerns/cells/transaction_callback.rb
@@ -14,16 +14,22 @@ def self.visited_stack
     end
 
     def begin_isolated_db_transaction(...)
+      return super unless Current.feature_cells_unique_claims
+
       TransactionCallback.visited_stack << {}
       super
     end
 
     def begin_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       TransactionCallback.visited_stack << {}
       super
     end
 
     def commit_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       # Don't pop the visited stack yet, and only do that after it's
       # successfully committed, because errors can happen while claiming,
       # and when that happens, we'll roll it back, and we need to look at
@@ -36,6 +42,8 @@ def commit_db_transaction
     end
 
     def rollback_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       Unique.rollback do
         super
 
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index db8c3c84f39b7d..ceb44e268fea7d 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -62,14 +62,14 @@ def self.rollback
 
     def inspect_save_changes
       return unless Gitlab.config.cell.enabled
-      return unless Feature.enabled?(:cells_unique_claims)
+      return unless Current.feature_cells_unique_claims
 
       InspectChanges.new(self).inspect_changes
     end
 
     def inspect_destroy_changes
       return unless Gitlab.config.cell.enabled
-      return unless Feature.enabled?(:cells_unique_claims)
+      return unless Current.feature_cells_unique_claims
 
       InspectChanges.new(self, destroy: true).inspect_changes
     end
diff --git a/app/models/current.rb b/app/models/current.rb
index 6747b08526f7a6..a4f78841b0c8ff 100644
--- a/app/models/current.rb
+++ b/app/models/current.rb
@@ -18,6 +18,9 @@ def message
   attribute :organization, :organization_assigned
   attribute :token_info
 
+  attribute :feature_cells_unique_claims
+  resets :reset_feature_cells_unique_claims
+
   def organization=(value)
     # We want to explicitly allow only one organization assignment per thread
     # This fits the request/response cycle, but of course for rake tasks/background jobs that use the same thread,
@@ -45,6 +48,10 @@ def organization
     super
   end
 
+  def reset_feature_cells_unique_claims
+    self.feature_cells_unique_claims = Feature.enabled?(:cells_unique_claims) # rubocop:disable Gitlab/FeatureFlagWithoutActor -- We can't easily tie this to an actor
+  end
+
   private
 
   # Do not allow to reset this
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index d4121c29301125..392bcd492a0d0c 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -88,6 +88,7 @@
       context 'when feature flag is disabled' do
         before do
           stub_feature_flags(cells_unique_claims: false)
+          Current.reset
         end
 
         it 'does not call inspect_save_changes on save' do
-- 
GitLab


From 142fc4aaeaba89cdb829e7bbfc521f9f3c9a621d Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Thu, 2 Oct 2025 00:35:55 +0800
Subject: [PATCH 15/29] Use the sharding keys to determine the owner

We also update that instead of storing only the changes, we store the
record as well, so that we can use the record to load the sharding keys.
This also simplifies some codes because we can get data from the record
instead of passing a few values around.
---
 app/models/application_record.rb                   |  4 ++++
 app/models/concerns/cells/inspect_changes.rb       | 14 +++++++++++++-
 app/models/concerns/cells/transaction_callback.rb  |  6 +++++-
 app/models/concerns/cells/unique.rb                |  9 ++++++---
 app/models/email.rb                                |  2 +-
 app/models/gpg_key.rb                              |  2 +-
 app/models/gpg_key_subkey.rb                       |  2 +-
 app/models/key.rb                                  |  2 +-
 app/models/namespace.rb                            |  1 +
 app/models/organizations/organization.rb           |  2 +-
 app/models/route.rb                                |  2 +-
 app/models/user.rb                                 |  2 +-
 spec/models/concerns/cells/unique_spec.rb          |  4 ++--
 spec/models/email_spec.rb                          |  2 +-
 spec/models/gpg_key_spec.rb                        |  2 +-
 spec/models/gpg_key_subkey_spec.rb                 |  2 +-
 spec/models/key_spec.rb                            |  2 +-
 spec/models/organizations/organization_spec.rb     |  2 +-
 spec/models/route_spec.rb                          |  2 +-
 spec/models/user_spec.rb                           |  2 +-
 .../cells/unique_shared_examples.rb                |  4 ++--
 21 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 7cec15efa8d468..8195ee8d5f08f1 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -165,6 +165,10 @@ def self.sharding_keys
     @sharding_keys ||= Gitlab::Database::Dictionary.entry(table_name).sharding_key || {}
   end
 
+  def sharding_keys
+    @sharding_keys ||= self.class.sharding_keys.transform_keys(&method(:public_send)).invert
+  end
+
   def readable_by?(user)
     Ability.allowed?(user, :"read_#{to_ability_name}", self)
   end
diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
index 8a9cffe0119052..0d4c96d8d34f93 100644
--- a/app/models/concerns/cells/inspect_changes.rb
+++ b/app/models/concerns/cells/inspect_changes.rb
@@ -2,6 +2,18 @@
 
 module Cells
   class InspectChanges
+    class Visit
+      extend Forwardable
+
+      attr_reader :record
+      def_delegators :@map, :[], :[]=, :filter_map
+
+      def initialize(new_record)
+        @record = new_record
+        @map = {}
+      end
+    end
+
     attr_reader :record
 
     def initialize(new_record, destroy: false)
@@ -43,7 +55,7 @@ def inspected
     end
 
     def prepare_inspecting
-      visited[inspect_changes_key] = {}
+      visited[inspect_changes_key] = Visit.new(record)
     end
 
     def inspect_record
diff --git a/app/models/concerns/cells/transaction_callback.rb b/app/models/concerns/cells/transaction_callback.rb
index 736f01cdf2326b..84264115eeebe4 100644
--- a/app/models/concerns/cells/transaction_callback.rb
+++ b/app/models/concerns/cells/transaction_callback.rb
@@ -8,6 +8,10 @@ def self.visited
       visited_stack.first # Flatten to the first
     end
 
+    def self.last_visited
+      visited_stack.last
+    end
+
     def self.visited_stack
       Thread.current.thread_variable_get(THREAD_VARIABLE_NAME) ||
         Thread.current.thread_variable_set(THREAD_VARIABLE_NAME, [])
@@ -35,7 +39,7 @@ def commit_db_transaction
       # and when that happens, we'll roll it back, and we need to look at
       # the stack to know if it's a unbalanced ROLLBACK or not, therefore
       # we only pop after successful COMMIT or ROLLBACk.
-      Unique.claim_and_commit(TransactionCallback.visited_stack.last) do
+      Unique.claim_and_commit(TransactionCallback.last_visited.values) do
         super
         TransactionCallback.visited_stack.pop
       end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index ceb44e268fea7d..c330659a8c097f 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -11,10 +11,11 @@ module Unique
     included do
       after_save :inspect_save_changes
       before_destroy :inspect_destroy_changes
+      cells_config
     end
 
     class_methods do
-      attr_reader :cells_owner_type, :cells_source_table
+      attr_reader :cells_type, :cells_source_table
 
       def cells_unique_attributes
         @cells_unique_attributes ||= {}
@@ -24,8 +25,10 @@ def cells_unique_has_many_associations
         @cells_unique_has_many_associations ||= []
       end
 
-      def cells_config(owner_type:, source_table:)
-        @cells_owner_type = owner_type
+      def cells_config(
+        type: OWNER_TYPE::UNSPECIFIED,
+        source_table: SOURCE_TABLE::UNSPECIFIED)
+        @cells_type = type
         @cells_source_table = source_table
       end
 
diff --git a/app/models/email.rb b/app/models/email.rb
index 7edbd89dc21dc7..677b273c3ef879 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -7,7 +7,7 @@ class Email < ApplicationRecord
 
   cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
-  cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
+  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
 
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 4fa06af9af70a0..a01cd55c69df01 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -11,7 +11,7 @@ class GpgKey < ApplicationRecord
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
   cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
 
-  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
+  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
 
   cells_unique_has_many :subkeys
 
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 03cdff865730e3..3c041a6e6b6f79 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -5,7 +5,7 @@ class GpgKeySubkey < ApplicationRecord
   include Cells::Unique
 
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
-  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::UNSPECIFIED
+  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index e6f5bb7a417d81..828ed84946dba7 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -13,7 +13,7 @@ class Key < ApplicationRecord
   cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
   cells_unique_attribute :fingerprint_sha256, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
 
-  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::ORGANIZATIONS
+  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha256_attribute :fingerprint_sha256
 
diff --git a/app/models/namespace.rb b/app/models/namespace.rb
index 2740c045c76fb1..34433c0b81345e 100644
--- a/app/models/namespace.rb
+++ b/app/models/namespace.rb
@@ -22,6 +22,7 @@ class Namespace < ApplicationRecord
   include UseSqlFunctionForPrimaryKeyLookups
   include SafelyChangeColumnDefault
   include Todoable
+  include Cells::Unique
 
   extend Gitlab::Utils::Override
 
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 03aa4c97db06d3..70a718363fc9e0 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -12,7 +12,7 @@ class Organization < ApplicationRecord
 
     cells_unique_attribute :path, type: BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
 
-    cells_config owner_type: OWNER_TYPE::ORGANIZATION, source_table: SOURCE_TABLE::ORGANIZATIONS
+    cells_config type: OWNER_TYPE::ORGANIZATION, source_table: SOURCE_TABLE::ORGANIZATIONS
 
     DEFAULT_ORGANIZATION_ID = 1
 
diff --git a/app/models/route.rb b/app/models/route.rb
index aab42fd78f8de4..700c8b5e212c5d 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -8,7 +8,7 @@ class Route < ApplicationRecord
 
   cells_unique_attribute :path, type: BUCKET_TYPE::ROUTES
 
-  cells_config owner_type: OWNER_TYPE::UNSPECIFIED, source_table: SOURCE_TABLE::ROUTES
+  cells_config source_table: SOURCE_TABLE::ROUTES
 
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
diff --git a/app/models/user.rb b/app/models/user.rb
index af602ecaebb0db..85ac44ea536274 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -43,7 +43,7 @@ class User < ApplicationRecord
   cells_unique_has_many :gpg_keys
   cells_unique_has_many :emails
 
-  cells_config owner_type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
+  cells_config type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
 
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index 392bcd492a0d0c..b2cc920966a4be 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -8,11 +8,11 @@
 
   describe 'configuration' do
     it 'sets and retrieves cell configuration' do
-      test_klass.cells_config(owner_type: :user, source_table: :users)
+      test_klass.cells_config(type: :user, source_table: :users)
       test_klass.cells_unique_attribute(:email, type: :string)
       test_klass.cells_unique_attribute(:username, type: :string)
 
-      expect(test_klass.cells_owner_type).to eq(:user)
+      expect(test_klass.cells_type).to eq(:user)
       expect(test_klass.cells_source_table).to eq(:users)
       expect(test_klass.cells_unique_attributes).to eq(email: { type: :string }, username: { type: :string })
     end
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index e4d820552a7df5..e0fb9df7e0f8ad 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -6,7 +6,7 @@
   let_it_be(:user) { create(:user) }
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::USER,
+    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
     source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
     unique_attributes: [:email]
 
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index b71b3bf4b7eb7f..acad9633d120c7 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -6,7 +6,7 @@
   subject { build(:gpg_key) }
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
     source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint, :primary_keyid]
 
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index 295f0712cec532..e9d8a17cce7b6e 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -16,7 +16,7 @@
   end
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
     source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
     unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index 66ca553c308736..3b6a7b142f2414 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -6,7 +6,7 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
     source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
     unique_attributes: [:key, :fingerprint_sha256]
 
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index afde6944e25b33..3297bfa30aee1e 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -6,7 +6,7 @@
   let_it_be_with_refind(:organization) { create(:organization) }
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::ORGANIZATION,
+    type: Cells::Unique::OWNER_TYPE::ORGANIZATION,
     source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
     unique_attributes: [:path]
 
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index 2207e72b2d8745..4db4bb72cba463 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -7,7 +7,7 @@
   let(:route) { group.route }
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
+    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
     source_table: Cells::Unique::SOURCE_TABLE::ROUTES,
     unique_attributes: [:path]
 
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index e7194c5361aff2..d4b23b23bd1d6c 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -50,7 +50,7 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    owner_type: Cells::Unique::OWNER_TYPE::USER,
+    type: Cells::Unique::OWNER_TYPE::USER,
     source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
     unique_attributes: [:username]
 
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
index f5f152f52e705a..17a13d413936fe 100644
--- a/spec/support/shared_examples/cells/unique_shared_examples.rb
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
-RSpec.shared_examples 'cells unique model' do |owner_type:, source_table:, unique_attributes:|
+RSpec.shared_examples 'cells unique model' do |type:, source_table:, unique_attributes:|
   it 'has the expected owner_type' do
-    expect(described_class.cells_owner_type).to eq(owner_type)
+    expect(described_class.cells_type).to eq(type)
   end
 
   it 'has the expected source_table' do
-- 
GitLab


From 05f39528cf8e2aeb3db176fcd381228439ede66b Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Tue, 7 Oct 2025 21:30:00 +0800
Subject: [PATCH 16/29] Use attributes_in_database to read the actual data

---
 app/models/concerns/cells/inspect_changes.rb | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
index 0d4c96d8d34f93..e5c2c523951c0f 100644
--- a/app/models/concerns/cells/inspect_changes.rb
+++ b/app/models/concerns/cells/inspect_changes.rb
@@ -72,16 +72,16 @@ def destroying?
 
     def inspect_deleted_attributes
       record.class.cells_unique_attributes.each_key do |column|
-        inspected[column] = [record.attributes[column.to_s]]
+        inspected[column] = [record.attribute_in_database(column)]
       end
     end
 
     def inspect_saved_changes
       record.saved_changes.each do |column, (old_value, new_value)|
         if record.class.cells_unique_attributes[column.to_sym] # rubocop:disable Style/Next -- I think this is more clear
-          # XXX: new_value in GpgKey is transformed into upcase,
-          # which mismatches with what we're storing in the database
-          new_value = record.attributes[column]
+          # new_value in GpgKey is transformed into upcase, which
+          # mismatches with what we're storing in the database
+          new_value = record.attribute_in_database(column)
           inspected[column] = [old_value, new_value]
         end
       end
-- 
GitLab


From 195d9beb1f9ce6bb8286f476721980c43ae90672 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Wed, 8 Oct 2025 00:15:32 +0800
Subject: [PATCH 17/29] Adapt to latest topology service client

---
 app/models/application_record.rb              |  4 +-
 app/models/concerns/cells/inspect_changes.rb  |  5 ++-
 app/models/concerns/cells/unique.rb           | 34 +++++++++++---
 app/models/current.rb                         |  4 ++
 app/models/email.rb                           |  2 -
 app/models/gpg_key.rb                         |  2 -
 app/models/gpg_key_subkey.rb                  |  1 -
 app/models/key.rb                             |  2 -
 app/models/namespace.rb                       |  2 +
 app/models/organizations/organization.rb      |  2 -
 app/models/route.rb                           |  2 -
 app/models/user.rb                            |  2 -
 .../concerns/cells/inspect_changes_spec.rb    | 45 ++++++++++++-------
 .../cells/transaction_callback_spec.rb        |  8 ++--
 spec/models/concerns/cells/unique_spec.rb     |  4 +-
 spec/models/email_spec.rb                     |  4 +-
 spec/models/gpg_key_spec.rb                   |  4 +-
 spec/models/gpg_key_subkey_spec.rb            |  4 +-
 spec/models/key_spec.rb                       |  4 +-
 .../models/organizations/organization_spec.rb |  4 +-
 spec/models/route_spec.rb                     |  4 +-
 spec/models/user_spec.rb                      |  4 +-
 .../cells/unique_shared_examples.rb           |  6 +--
 23 files changed, 91 insertions(+), 62 deletions(-)

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 8195ee8d5f08f1..c3c5379150d366 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -166,7 +166,9 @@ def self.sharding_keys
   end
 
   def sharding_keys
-    @sharding_keys ||= self.class.sharding_keys.transform_keys(&method(:public_send)).invert
+    @sharding_keys ||= self.class.sharding_keys.transform_keys do |key|
+      public_send(key) # rubocop:disable GitlabSecurity/PublicSend -- this works by definition
+    end.invert
   end
 
   def readable_by?(user)
diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
index e5c2c523951c0f..e152f7c29080b0 100644
--- a/app/models/concerns/cells/inspect_changes.rb
+++ b/app/models/concerns/cells/inspect_changes.rb
@@ -6,6 +6,7 @@ class Visit
       extend Forwardable
 
       attr_reader :record
+
       def_delegators :@map, :[], :[]=, :filter_map
 
       def initialize(new_record)
@@ -72,7 +73,7 @@ def destroying?
 
     def inspect_deleted_attributes
       record.class.cells_unique_attributes.each_key do |column|
-        inspected[column] = [record.attribute_in_database(column)]
+        inspected[column.to_s] = [record.attribute_in_database(column)]
       end
     end
 
@@ -96,7 +97,7 @@ def inspect_associations
       record.class.reflect_on_all_associations.each do |reflection|
         association_proxy = record.association(reflection.name)
 
-        if association_proxy.loaded? # rubocop:disable Style/Next -- I think this is more clear
+        if association_proxy.loaded? # rubocop:disable Style/IfUnlessModifier -- I think this is more clear
           inspect_loaded_associations(association_proxy)
         end
       end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index c330659a8c097f..1c6a5c86b2bb28 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -4,9 +4,9 @@ module Cells
   module Unique
     extend ActiveSupport::Concern
 
-    BUCKET_TYPE = Gitlab::Cells::TopologyService::ClaimBucket::BucketType
-    OWNER_TYPE = Gitlab::Cells::TopologyService::ClaimOwner::OwnerType
-    SOURCE_TABLE = Gitlab::Cells::TopologyService::ClaimSource::SourceTable
+    BUCKET_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Bucket::Type
+    SOURCE_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Source::Type
+    SUBJECT_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Subject::Type
 
     included do
       after_save :inspect_save_changes
@@ -15,7 +15,7 @@ module Unique
     end
 
     class_methods do
-      attr_reader :cells_type, :cells_source_table
+      attr_reader :cells_type, :cells_source_type
 
       def cells_unique_attributes
         @cells_unique_attributes ||= {}
@@ -26,10 +26,30 @@ def cells_unique_has_many_associations
       end
 
       def cells_config(
-        type: OWNER_TYPE::UNSPECIFIED,
-        source_table: SOURCE_TABLE::UNSPECIFIED)
+        type: default_type,
+        source_type: default_source_type)
         @cells_type = type
-        @cells_source_table = source_table
+        @cells_source_type = source_type
+      end
+
+      def default_source_type
+        source_type_name = "RAILS_TABLE_#{table_name.upcase}"
+
+        if SOURCE_TYPE.const_defined?(source_type_name)
+          SOURCE_TYPE.const_get(source_type_name, false)
+        else
+          SOURCE_TYPE::UNSPECIFIED
+        end
+      end
+
+      def default_type
+        subject_type_name = name.demodulize.upcase
+
+        if SUBJECT_TYPE.const_defined?(subject_type_name)
+          SUBJECT_TYPE.const_get(subject_type_name, false)
+        else
+          SUBJECT_TYPE::UNSPECIFIED
+        end
       end
 
       def cells_unique_attribute(name, type:)
diff --git a/app/models/current.rb b/app/models/current.rb
index a4f78841b0c8ff..e261872e61790b 100644
--- a/app/models/current.rb
+++ b/app/models/current.rb
@@ -21,6 +21,10 @@ def message
   attribute :feature_cells_unique_claims
   resets :reset_feature_cells_unique_claims
 
+  def initialize
+    reset
+  end
+
   def organization=(value)
     # We want to explicitly allow only one organization assignment per thread
     # This fits the request/response cycle, but of course for rake tasks/background jobs that use the same thread,
diff --git a/app/models/email.rb b/app/models/email.rb
index 677b273c3ef879..471124ffa15bb3 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -7,8 +7,6 @@ class Email < ApplicationRecord
 
   cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
-  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
-
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
 
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index a01cd55c69df01..30b0096cfe7a2e 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -11,8 +11,6 @@ class GpgKey < ApplicationRecord
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
   cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
 
-  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
-
   cells_unique_has_many :subkeys
 
   sha_attribute :primary_keyid
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 3c041a6e6b6f79..5749818cb7ccc3 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -5,7 +5,6 @@ class GpgKeySubkey < ApplicationRecord
   include Cells::Unique
 
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
-  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index 828ed84946dba7..67fd9a465df979 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -13,8 +13,6 @@ class Key < ApplicationRecord
   cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
   cells_unique_attribute :fingerprint_sha256, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
 
-  cells_config source_table: SOURCE_TABLE::UNSPECIFIED
-
   sha256_attribute :fingerprint_sha256
 
   belongs_to :user
diff --git a/app/models/namespace.rb b/app/models/namespace.rb
index 34433c0b81345e..5969ce15b6de4b 100644
--- a/app/models/namespace.rb
+++ b/app/models/namespace.rb
@@ -24,6 +24,8 @@ class Namespace < ApplicationRecord
   include Todoable
   include Cells::Unique
 
+  cells_config type: SUBJECT_TYPE::GROUP
+
   extend Gitlab::Utils::Override
 
   ignore_columns :description, :description_html, :cached_markdown_version, remove_with: '18.3', remove_after: '2025-07-17'
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 70a718363fc9e0..fdf7808b796e0a 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -12,8 +12,6 @@ class Organization < ApplicationRecord
 
     cells_unique_attribute :path, type: BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
 
-    cells_config type: OWNER_TYPE::ORGANIZATION, source_table: SOURCE_TABLE::ORGANIZATIONS
-
     DEFAULT_ORGANIZATION_ID = 1
 
     scope :without_default, -> { id_not_in(DEFAULT_ORGANIZATION_ID) }
diff --git a/app/models/route.rb b/app/models/route.rb
index 700c8b5e212c5d..ec429506b31c02 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -8,8 +8,6 @@ class Route < ApplicationRecord
 
   cells_unique_attribute :path, type: BUCKET_TYPE::ROUTES
 
-  cells_config source_table: SOURCE_TABLE::ROUTES
-
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
   validates :source, presence: true
diff --git a/app/models/user.rb b/app/models/user.rb
index 85ac44ea536274..b8c2f364e7d8f0 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -43,8 +43,6 @@ class User < ApplicationRecord
   cells_unique_has_many :gpg_keys
   cells_unique_has_many :emails
 
-  cells_config type: OWNER_TYPE::USER, source_table: SOURCE_TABLE::UNSPECIFIED
-
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
   DEFAULT_NOTIFICATION_LEVEL = :participating
diff --git a/spec/models/concerns/cells/inspect_changes_spec.rb b/spec/models/concerns/cells/inspect_changes_spec.rb
index ca9a0e7bea6cc3..0b0a5264e9f94b 100644
--- a/spec/models/concerns/cells/inspect_changes_spec.rb
+++ b/spec/models/concerns/cells/inspect_changes_spec.rb
@@ -4,15 +4,26 @@
 
 RSpec.describe Cells::InspectChanges, feature_category: :cell do
   let(:user) { create(:user) }
-  let(:inspector) { described_class.new(user) }
+  let(:inspector) { described_class.new(user, destroy: destroy) }
+  let(:destroy) { false }
 
   let(:dummy_class) do
     Class.new(ActiveRecord::Base) do
+      def self.name
+        'Namespace::ClassName'
+      end
+
       self.table_name = 'users'
       include Cells::Unique
 
-      cells_unique_attribute :name, type: :string
-      cells_unique_attribute :email, type: :string
+      cells_unique_attribute :name,
+        type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
+      cells_unique_attribute :email,
+        type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
+
+      def attribute_in_database(column)
+        public_send(column)
+      end
     end
   end
 
@@ -52,7 +63,11 @@
 
           inspector.inspect_changes
 
-          expect(visited.keys).to match_array([['users', user.id], ["organizations", user.organization_id]])
+          expect(visited.keys).to match_array([
+            ['users', user.id],
+            ['emails', user.email_ids.first],
+            ['organizations', user.organization_id]
+          ])
         end
       end
 
@@ -112,17 +127,17 @@
       before do
         allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
         allow(Cells::Unique).to receive(:claim_and_commit)
-        allow(user).to receive(:destroyed?).and_return(false)
       end
 
       context 'with saved changes' do
         before do
-          allow(user).to receive(:destroyed?).and_return(false)
-          allow(user).to receive_messages(destroyed?: false, saved_changes:
-           { 'name' => %w[Jane John],
-             'email' => ['jane@example.com',
-               'john@example.com'],
-             'age' => [29, 30] })
+          allow(user).to receive(:saved_changes).and_return(
+            {
+              'name' => %w[Jane John],
+              'email' => ['jane@example.com',
+                'john@example.com'],
+              'age' => [29, 30]
+            })
         end
 
         it 'tracks only unique attribute changes' do
@@ -136,16 +151,14 @@
       end
 
       context 'with destroyed record' do
-        before do
-          allow(user).to receive(:destroyed?).and_return(true)
-        end
+        let(:destroy) { true }
 
         it 'tracks deleted attribute values' do
           inspector.inspect_changes
 
           inspected = inspector.send(:inspected)
-          expect(inspected[:name]).to eq(['John'])
-          expect(inspected[:email]).to eq(['john@example.com'])
+          expect(inspected['name']).to eq(['John'])
+          expect(inspected['email']).to eq(['john@example.com'])
         end
       end
     end
diff --git a/spec/models/concerns/cells/transaction_callback_spec.rb b/spec/models/concerns/cells/transaction_callback_spec.rb
index 90d47427ea6fe1..bf1fa134551f82 100644
--- a/spec/models/concerns/cells/transaction_callback_spec.rb
+++ b/spec/models/concerns/cells/transaction_callback_spec.rb
@@ -102,7 +102,7 @@ def rollback_db_transaction
 
       it 'claims and commits with visited data' do
         expect(Cells::Unique).to receive(:claim_and_commit)
-          .with({ ['users', 1] => { 'email' => %w[old new] } })
+          .with([{ 'email' => %w[old new] }])
           .and_yield
 
         connection.commit_db_transaction
@@ -132,7 +132,7 @@ def rollback_db_transaction
       it 'handles empty visited data' do
         described_class.visited_stack.last.clear
 
-        expect(Cells::Unique).to receive(:claim_and_commit).with({}).and_yield
+        expect(Cells::Unique).to receive(:claim_and_commit).with([]).and_yield
 
         connection.commit_db_transaction
         expect(described_class.visited_stack).to be_empty
@@ -196,7 +196,7 @@ def rollback_db_transaction
       described_class.visited[['users', 1]] = { 'name' => %w[old new] }
 
       expect(Cells::Unique).to receive(:claim_and_commit)
-        .with({ ['users', 1] => { 'name' => %w[old new] } })
+        .with([{ 'name' => %w[old new] }])
         .and_yield
 
       connection.commit_db_transaction
@@ -227,7 +227,7 @@ def rollback_db_transaction
       connection.rollback_db_transaction
 
       expect(Cells::Unique).to receive(:claim_and_commit)
-        .with({ ['outer', 1] => 'data' })
+        .with(['data'])
         .and_yield
 
       connection.commit_db_transaction
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index b2cc920966a4be..a3a95fa7d0a0bf 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -8,12 +8,12 @@
 
   describe 'configuration' do
     it 'sets and retrieves cell configuration' do
-      test_klass.cells_config(type: :user, source_table: :users)
+      test_klass.cells_config(type: :user, source_type: :rails_table_users)
       test_klass.cells_unique_attribute(:email, type: :string)
       test_klass.cells_unique_attribute(:username, type: :string)
 
       expect(test_klass.cells_type).to eq(:user)
-      expect(test_klass.cells_source_table).to eq(:users)
+      expect(test_klass.cells_source_type).to eq(:rails_table_users)
       expect(test_klass.cells_unique_attributes).to eq(email: { type: :string }, username: { type: :string })
     end
   end
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index e0fb9df7e0f8ad..d8d8375fb47a46 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -6,8 +6,8 @@
   let_it_be(:user) { create(:user) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
-    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_EMAILS,
     unique_attributes: [:email]
 
   describe 'modules' do
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index acad9633d120c7..0af57bfa542e7e 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -6,8 +6,8 @@
   subject { build(:gpg_key) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
-    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint, :primary_keyid]
 
   describe "associations" do
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index e9d8a17cce7b6e..0eda583f420ec6 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -16,7 +16,7 @@
   end
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
-    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index 3b6a7b142f2414..ea2aa7c444dcc2 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -6,8 +6,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
-    source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
+    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint_sha256]
 
   describe "Associations" do
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index 3297bfa30aee1e..634487c6c42412 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -6,8 +6,8 @@
   let_it_be_with_refind(:organization) { create(:organization) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::ORGANIZATION,
-    source_table: Cells::Unique::SOURCE_TABLE::ORGANIZATIONS,
+    type: Cells::Unique::SUBJECT_TYPE::ORGANIZATION,
+    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
     unique_attributes: [:path]
 
   describe 'associations' do
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index 4db4bb72cba463..a15e7e93d5ff63 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -7,8 +7,8 @@
   let(:route) { group.route }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::UNSPECIFIED,
-    source_table: Cells::Unique::SOURCE_TABLE::ROUTES,
+    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_ROUTES,
     unique_attributes: [:path]
 
   describe 'relationships' do
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index d4b23b23bd1d6c..81442a6aeeec4b 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -50,8 +50,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::OWNER_TYPE::USER,
-    source_table: Cells::Unique::SOURCE_TABLE::UNSPECIFIED,
+    type: Cells::Unique::SUBJECT_TYPE::USER,
+    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_USERS,
     unique_attributes: [:username]
 
   describe 'modules' do
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
index 17a13d413936fe..7c16a72f8a1f17 100644
--- a/spec/support/shared_examples/cells/unique_shared_examples.rb
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -1,12 +1,12 @@
 # frozen_string_literal: true
 
-RSpec.shared_examples 'cells unique model' do |type:, source_table:, unique_attributes:|
+RSpec.shared_examples 'cells unique model' do |type:, source_type:, unique_attributes:|
   it 'has the expected owner_type' do
     expect(described_class.cells_type).to eq(type)
   end
 
-  it 'has the expected source_table' do
-    expect(described_class.cells_source_table).to eq(source_table)
+  it 'has the expected source_type' do
+    expect(described_class.cells_source_type).to eq(source_type)
   end
 
   it 'has the expected unique attributes' do
-- 
GitLab


From cc113dea76687917a0988fff22b9c1a9b7b0553e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Mon, 6 Oct 2025 13:03:06 +0200
Subject: [PATCH 18/29] Introduce `Cells::TransactionRecord` as a way to hook
 into transaction

---
 app/models/cells/outstanding_lease.rb         |  51 +++++++
 app/models/cells/transaction_record.rb        |  89 ++++++++++++
 app/models/concerns/cells/inspect_changes.rb  | 129 ------------------
 .../concerns/cells/transaction_callback.rb    |  70 ----------
 app/models/concerns/cells/unique.rb           | 117 ++++++++--------
 app/models/current.rb                         |  11 +-
 app/models/email.rb                           |   2 +
 app/models/gpg_key.rb                         |   2 +
 app/models/gpg_key_subkey.rb                  |   1 +
 app/models/key.rb                             |   2 +
 app/models/organizations/organization.rb      |   2 +
 app/models/route.rb                           |   2 +
 app/models/user.rb                            |   2 +
 ...ctive_record_cells_transaction_callback.rb |  10 --
 14 files changed, 223 insertions(+), 267 deletions(-)
 create mode 100644 app/models/cells/transaction_record.rb
 delete mode 100644 app/models/concerns/cells/inspect_changes.rb
 delete mode 100644 app/models/concerns/cells/transaction_callback.rb
 delete mode 100644 config/initializers/active_record_cells_transaction_callback.rb

diff --git a/app/models/cells/outstanding_lease.rb b/app/models/cells/outstanding_lease.rb
index f74763e6b828ae..caf509046270a0 100644
--- a/app/models/cells/outstanding_lease.rb
+++ b/app/models/cells/outstanding_lease.rb
@@ -3,5 +3,56 @@
 module Cells
   class OutstandingLease < ApplicationRecord
     self.primary_key = :uuid
+
+    def self.create_from_request!(create_records:, destroy_records:)
+      req = Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateRequest.new(
+        create_records: create_records,
+        destroy_records: destroy_records,
+        cell_id: claim_service.cell_id
+      )
+
+      pp("send_begin_update req #{req}")
+      res = self.claim_service.begin_update(req)
+
+      pp("send_begin_update res #{res.lease_uuid}")
+      self.create!(uuid: res.lease_uuid.value)
+    end
+
+    def send_commit_update!
+      req = Gitlab::Cells::TopologyService::Claims::V1::CommitUpdateRequest.new(
+        lease_uuid: Gitlab::Cells::TopologyService::Types::V1::UUID.new(value: uuid),
+        cell_id: self.class.claim_service.cell_id
+      )
+
+      pp("send_commit_update #{self.object_id} #{req}")
+      self.class.claim_service.commit_update(
+        req
+      )
+
+      pp("send_commit_update done #{self.object_id} #{req}")
+    end
+
+    def send_rollback_update!
+      return unless uuid
+
+      req = Gitlab::Cells::TopologyService::Claims::V1::RollbackUpdateRequest.new(
+        lease_uuid: Gitlab::Cells::TopologyService::Types::V1::UUID.new(value: uuid),
+        cell_id: self.class.claim_service.cell_id
+      )
+
+      pp("send_rollback_update #{self.object_id} #{req}")
+
+      self.class.claim_service.rollback_update(
+        req
+      )
+
+      pp("send_rollback_update done #{self.object_id} #{req}")
+    end
+
+    private
+
+    def self.claim_service
+      @claim_service ||= ::Gitlab::TopologyServiceClient::ClaimService.instance
+    end
   end
 end
diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
new file mode 100644
index 00000000000000..d33516886a9097
--- /dev/null
+++ b/app/models/cells/transaction_record.rb
@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module Cells
+  class TransactionRecord
+    TRANSACTION_RECORD_KEY = :@cells_current_transaction
+
+    # TODO: hack to assign to current transaction
+    # Figure out a better way to do this
+    # Current? but what if nested transactions?
+    # what is the transactions gets rolled back?
+    def self.current_transaction(connection)
+      return unless Current.cells_claim_leases?
+
+      # The Cells::OutstandingLease requires a transaction to be open
+      # to ensure that the lease is only created if the transaction
+      # within a transaction and not outside of one
+      unless connection.current_transaction
+        raise 'The Cells::TransactionRecord requires transaction to be open'
+      end
+
+      current_transaction = connection.current_transaction
+
+      instance = current_transaction.instance_variable_get(TRANSACTION_RECORD_KEY)
+      return instance if instance
+
+      TransactionRecord.new(connection, current_transaction).tap do |instance|
+        current_transaction.instance_variable_set(TRANSACTION_RECORD_KEY, instance)
+        current_transaction.add_record(instance)
+      end
+    end
+
+    def initialize(connection, transaction)
+      @connection = connection
+      @transaction = transaction
+      @create_records = []
+      @destroy_records = []
+      @outstanding_lease = nil
+      @done = false
+    end
+
+    def create_record(metadata)
+      @create_records << metadata
+      pp("create_record #{self.object_id} #{metadata}")
+    end
+
+    def destroy_record(metadata)
+      @destroy_records << metadata
+      pp("destroy_record #{self.object_id} #{metadata}")
+    end
+
+    # Callbacks required by the ActiveRecord::ConnectionAdapters::TransactionState
+
+    def trigger_transactional_callbacks?
+      true
+    end
+  
+    def before_committed!
+      raise 'Already done' if @done
+      raise 'Already created lease' if @outstanding_lease
+      raise 'Cross-database lease' if Cells::OutstandingLease.connection != @connection
+
+      @outstanding_lease = Cells::OutstandingLease.create_from_request!(
+        create_records: @create_records,
+        destroy_records: @destroy_records
+      )
+    end
+
+    def rolledback!(force_restore_state: false, should_run_callbacks: true)
+      raise 'Already done' if @done
+
+      # It is possible that lease might be not created yet,
+      # since the transaction might be rolledback prematurely
+      return unless @outstanding_lease
+
+      @outstanding_lease.send_rollback_update!
+      @outstanding_lease.destroy! # the lease is no longer needed
+      @done = true
+    end
+
+    def committed!(should_run_callbacks: true)
+      raise 'Already done' if @done
+      raise 'No lease created' unless @outstanding_lease
+
+      @outstanding_lease.send_commit_update!
+      @outstanding_lease.destroy! # the lease is no longer needed
+      @done = true
+    end
+  end
+end
diff --git a/app/models/concerns/cells/inspect_changes.rb b/app/models/concerns/cells/inspect_changes.rb
deleted file mode 100644
index e152f7c29080b0..00000000000000
--- a/app/models/concerns/cells/inspect_changes.rb
+++ /dev/null
@@ -1,129 +0,0 @@
-# frozen_string_literal: true
-
-module Cells
-  class InspectChanges
-    class Visit
-      extend Forwardable
-
-      attr_reader :record
-
-      def_delegators :@map, :[], :[]=, :filter_map
-
-      def initialize(new_record)
-        @record = new_record
-        @map = {}
-      end
-    end
-
-    attr_reader :record
-
-    def initialize(new_record, destroy: false)
-      @record = new_record
-      @destroy = destroy
-    end
-
-    def inspect_changes
-      return if !record.is_a?(Unique) || inspected
-
-      prepare_inspecting
-
-      inspect_record
-      inspect_associations
-
-      if TransactionCallback.visited
-        # There's an actual transaction, and TransactionCallback will
-        # handle this case right before committing
-      else
-        # Since the transaction is optimized away, TransactionCallback
-        # cannot handle this case, and we can just assume changes are
-        # going to be made, we treat this as if before_commit_all
-        Unique.claim_and_commit(visited)
-      end
-    end
-
-    private
-
-    def visited
-      @visited ||= TransactionCallback.visited || {}
-    end
-
-    def inspect_changes_key
-      @inspect_changes_key ||= [record.class.table_name, record.id]
-    end
-
-    def inspected
-      visited[inspect_changes_key]
-    end
-
-    def prepare_inspecting
-      visited[inspect_changes_key] = Visit.new(record)
-    end
-
-    def inspect_record
-      if destroying?
-        inspect_deleted_attributes
-      else
-        inspect_saved_changes
-      end
-    end
-
-    def destroying?
-      @destroy
-    end
-
-    def inspect_deleted_attributes
-      record.class.cells_unique_attributes.each_key do |column|
-        inspected[column.to_s] = [record.attribute_in_database(column)]
-      end
-    end
-
-    def inspect_saved_changes
-      record.saved_changes.each do |column, (old_value, new_value)|
-        if record.class.cells_unique_attributes[column.to_sym] # rubocop:disable Style/Next -- I think this is more clear
-          # new_value in GpgKey is transformed into upcase, which
-          # mismatches with what we're storing in the database
-          new_value = record.attribute_in_database(column)
-          inspected[column] = [old_value, new_value]
-        end
-      end
-    end
-
-    def inspect_associations
-      # We know the following associations will be deleted yet the hooks
-      # won't be called, so we laod it at once and inspect the changes
-      # as if they would get changed.
-      preload_nested_associations(record) if destroying?
-
-      record.class.reflect_on_all_associations.each do |reflection|
-        association_proxy = record.association(reflection.name)
-
-        if association_proxy.loaded? # rubocop:disable Style/IfUnlessModifier -- I think this is more clear
-          inspect_loaded_associations(association_proxy)
-        end
-      end
-    end
-
-    def inspect_loaded_associations(association_proxy)
-      Array.wrap(association_proxy.target).each do |association_record|
-        InspectChanges.new(association_record, destroy: destroying?).inspect_changes
-      end
-    end
-
-    def preload_nested_associations(record)
-      ActiveRecord::Associations::Preloader.new(
-        records: [record],
-        associations: lookup_associations(record.class)
-      ).call
-    end
-
-    def lookup_associations(model)
-      model.cells_unique_has_many_associations.map do |name|
-        { name => lookup_associations(reflect_association_model(model, name)) }
-      end
-    end
-
-    def reflect_association_model(model, name)
-      model.reflect_on_association(name).class_name.constantize
-    end
-  end
-end
diff --git a/app/models/concerns/cells/transaction_callback.rb b/app/models/concerns/cells/transaction_callback.rb
deleted file mode 100644
index 84264115eeebe4..00000000000000
--- a/app/models/concerns/cells/transaction_callback.rb
+++ /dev/null
@@ -1,70 +0,0 @@
-# frozen_string_literal: true
-
-module Cells
-  module TransactionCallback
-    THREAD_VARIABLE_NAME = :cells_unique_transaction_callback_visited_stack
-
-    def self.visited
-      visited_stack.first # Flatten to the first
-    end
-
-    def self.last_visited
-      visited_stack.last
-    end
-
-    def self.visited_stack
-      Thread.current.thread_variable_get(THREAD_VARIABLE_NAME) ||
-        Thread.current.thread_variable_set(THREAD_VARIABLE_NAME, [])
-    end
-
-    def begin_isolated_db_transaction(...)
-      return super unless Current.feature_cells_unique_claims
-
-      TransactionCallback.visited_stack << {}
-      super
-    end
-
-    def begin_db_transaction
-      return super unless Current.feature_cells_unique_claims
-
-      TransactionCallback.visited_stack << {}
-      super
-    end
-
-    def commit_db_transaction
-      return super unless Current.feature_cells_unique_claims
-
-      # Don't pop the visited stack yet, and only do that after it's
-      # successfully committed, because errors can happen while claiming,
-      # and when that happens, we'll roll it back, and we need to look at
-      # the stack to know if it's a unbalanced ROLLBACK or not, therefore
-      # we only pop after successful COMMIT or ROLLBACk.
-      Unique.claim_and_commit(TransactionCallback.last_visited.values) do
-        super
-        TransactionCallback.visited_stack.pop
-      end
-    end
-
-    def rollback_db_transaction
-      return super unless Current.feature_cells_unique_claims
-
-      Unique.rollback do
-        super
-
-        # Rails can be issuing unbalanced ROLLBACK (no corresponding BEGIN),
-        # and in that case we ignore and do not need to rollback the last
-        # lease because it's already committed! This can happen if there are
-        # errors internally coming from Rails while committing to the local
-        # database. We determine this by checking the visited stack. If the
-        # stack is empty, it means it's not balanced and we do not want to
-        # rollback on Topology service in this case.
-        if TransactionCallback.visited_stack.empty?
-          false
-        else
-          TransactionCallback.visited_stack.pop
-          true
-        end
-      end
-    end
-  end
-end
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 1c6a5c86b2bb28..8ef94a87a8d25b 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -5,17 +5,16 @@ module Unique
     extend ActiveSupport::Concern
 
     BUCKET_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Bucket::Type
-    SOURCE_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Source::Type
     SUBJECT_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Subject::Type
+    SOURCE_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Source::Type
 
     included do
-      after_save :inspect_save_changes
-      before_destroy :inspect_destroy_changes
-      cells_config
+      after_save :claim_save_changes
+      before_destroy :claim_destroy_changes
     end
 
     class_methods do
-      attr_reader :cells_type, :cells_source_type
+      attr_reader :cells_subject_type, :cells_subject_key, :cells_source_type
 
       def cells_unique_attributes
         @cells_unique_attributes ||= {}
@@ -25,31 +24,10 @@ def cells_unique_has_many_associations
         @cells_unique_has_many_associations ||= []
       end
 
-      def cells_config(
-        type: default_type,
-        source_type: default_source_type)
-        @cells_type = type
-        @cells_source_type = source_type
-      end
-
-      def default_source_type
-        source_type_name = "RAILS_TABLE_#{table_name.upcase}"
-
-        if SOURCE_TYPE.const_defined?(source_type_name)
-          SOURCE_TYPE.const_get(source_type_name, false)
-        else
-          SOURCE_TYPE::UNSPECIFIED
-        end
-      end
-
-      def default_type
-        subject_type_name = name.demodulize.upcase
-
-        if SUBJECT_TYPE.const_defined?(subject_type_name)
-          SUBJECT_TYPE.const_get(subject_type_name, false)
-        else
-          SUBJECT_TYPE::UNSPECIFIED
-        end
+      def cells_config(subject_type:, subject_key: nil, source_type: nil)
+        @cells_subject_type = subject_type
+        @cells_subject_key = subject_key || :organization_id
+        @cells_source_type = source_type || const_get("Gitlab::Cells::TopologyService::Claims::V1::Source::Type::RAILS_TABLE_#{table_name.upcase}")
       end
 
       def cells_unique_attribute(name, type:)
@@ -61,40 +39,71 @@ def cells_unique_has_many(name)
       end
     end
 
-    def self.claim_and_commit(changes)
-      if changes.empty?
-        yield if block_given?
-      else
-        pp changes
+    private
+
+    def claim_save_changes
+      transaction_record = ::Cells::TransactionRecord.current_transaction(connection)
+      return unless transaction_record
 
-        claim_client = Gitlab::TopologyServiceClient::ClaimService.new
-        claim_client.claim_and_commit(changes) do
-          yield if block_given?
+      pp("claim_save_changes #{self.object_id}: #{self.class.name} #{self.id}")
+
+      default_metadata = claim_metadata
+
+      self.class.cells_unique_attributes.each do |attribute, config|
+        next unless saved_change_to_attribute?(attribute)
+
+        was, is = saved_change_to_attribute(attribute)
+        pp("  attribute '#{attribute}' changed from '#{was}' to '#{is}'")
+
+        if !was.nil? && was != is
+          transaction_record.destroy_record(default_metadata.merge({
+            bucket: {
+              type: config[:type],
+              value: attribute_in_database(attribute)
+            }
+          }))
         end
-      end
-    end
 
-    def self.rollback
-      claim_client = Gitlab::TopologyServiceClient::ClaimService.new
-      claim_client.rollback do
-        yield
+        if !is.nil?
+          transaction_record.create_record(default_metadata.merge(
+            bucket: {
+              type: config[:type],
+              value: self.public_send(attribute)
+            }
+          ))
+        end
       end
     end
 
-    private
+    def claim_destroy_changes
+      transaction_record = ::Cells::TransactionRecord.current_transaction(connection)
+      return unless transaction_record
 
-    def inspect_save_changes
-      return unless Gitlab.config.cell.enabled
-      return unless Current.feature_cells_unique_claims
+      pp("claim_destroy_changes #{self.object_id}: #{self.class.name} #{self.id}")
 
-      InspectChanges.new(self).inspect_changes
-    end
+      default_metadata = claim_metadata
 
-    def inspect_destroy_changes
-      return unless Gitlab.config.cell.enabled
-      return unless Current.feature_cells_unique_claims
+      self.class.cells_unique_attributes.each do |attribute, config|
+        transaction_record.destroy_record(default_metadata.merge(
+          bucket: {
+            type: config[:type],
+            value: self.public_send(attribute)
+          }
+        ))
+      end
+    end
 
-      InspectChanges.new(self, destroy: true).inspect_changes
+    def claim_metadata
+      {
+        subject: {
+          type: self.class.cells_subject_type,
+          id: self.public_send(self.class.cells_subject_key)
+        },
+        source: {
+          type: self.class.cells_source_type,
+          rails_primary_key_id: self.public_send(self.class.primary_key)
+        },
+      }
     end
   end
 end
diff --git a/app/models/current.rb b/app/models/current.rb
index e261872e61790b..c8c3ff32280ea1 100644
--- a/app/models/current.rb
+++ b/app/models/current.rb
@@ -18,11 +18,14 @@ def message
   attribute :organization, :organization_assigned
   attribute :token_info
 
-  attribute :feature_cells_unique_claims
-  resets :reset_feature_cells_unique_claims
+  attribute :cells_claim_leases
 
-  def initialize
-    reset
+  def cells_claim_leases?
+    if cells_claim_leases.nil?
+      self.cells_claim_leases = Gitlab.config.cell.enabled && Feature.enabled?(:cells_unique_claims)
+    end
+
+    cells_claim_leases
   end
 
   def organization=(value)
diff --git a/app/models/email.rb b/app/models/email.rb
index 471124ffa15bb3..ba00f710da0a85 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -7,6 +7,8 @@ class Email < ApplicationRecord
 
   cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
+  cells_config subject_type: SUBJECT_TYPE::USER, subject_key: :user_id
+
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
 
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 30b0096cfe7a2e..7e09ddc12cc858 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -11,6 +11,8 @@ class GpgKey < ApplicationRecord
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
   cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
 
+  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED
+
   cells_unique_has_many :subkeys
 
   sha_attribute :primary_keyid
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 5749818cb7ccc3..161c9db5c15555 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -5,6 +5,7 @@ class GpgKeySubkey < ApplicationRecord
   include Cells::Unique
 
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
+  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index 67fd9a465df979..6fcf3ace91045f 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -13,6 +13,8 @@ class Key < ApplicationRecord
   cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
   cells_unique_attribute :fingerprint_sha256, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
 
+  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED
+
   sha256_attribute :fingerprint_sha256
 
   belongs_to :user
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index fdf7808b796e0a..fedfd5931eff69 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -12,6 +12,8 @@ class Organization < ApplicationRecord
 
     cells_unique_attribute :path, type: BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
 
+    cells_config subject_type: SUBJECT_TYPE::ORGANIZATION
+
     DEFAULT_ORGANIZATION_ID = 1
 
     scope :without_default, -> { id_not_in(DEFAULT_ORGANIZATION_ID) }
diff --git a/app/models/route.rb b/app/models/route.rb
index ec429506b31c02..38839ee1df89a5 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -8,6 +8,8 @@ class Route < ApplicationRecord
 
   cells_unique_attribute :path, type: BUCKET_TYPE::ROUTES
 
+  cells_config subject_type: SUBJECT_TYPE::GROUP, subject_key: :namespace_id
+
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
   validates :source, presence: true
diff --git a/app/models/user.rb b/app/models/user.rb
index b8c2f364e7d8f0..b43e079a20c3a1 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -43,6 +43,8 @@ class User < ApplicationRecord
   cells_unique_has_many :gpg_keys
   cells_unique_has_many :emails
 
+  cells_config subject_type: SUBJECT_TYPE::USER
+
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
   DEFAULT_NOTIFICATION_LEVEL = :participating
diff --git a/config/initializers/active_record_cells_transaction_callback.rb b/config/initializers/active_record_cells_transaction_callback.rb
deleted file mode 100644
index e213d795fb7d7d..00000000000000
--- a/config/initializers/active_record_cells_transaction_callback.rb
+++ /dev/null
@@ -1,10 +0,0 @@
-# frozen_string_literal: true
-
-ActiveSupport.on_load(:active_record) do
-  # Prepend into AbstractAdapter doesn't fully work, so we prepend to
-  # actual adapter we use.
-  if Gitlab.config.cell.enabled
-    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-      .prepend(Cells::TransactionCallback)
-  end
-end
-- 
GitLab


From 66069bf35114b2c140febe10f1ac6dc57e867b41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Thu, 9 Oct 2025 14:03:02 +0200
Subject: [PATCH 19/29] Do not allow to create records once lease was created

---
 app/models/cells/transaction_record.rb | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
index d33516886a9097..9e99ee6e917579 100644
--- a/app/models/cells/transaction_record.rb
+++ b/app/models/cells/transaction_record.rb
@@ -39,11 +39,15 @@ def initialize(connection, transaction)
     end
 
     def create_record(metadata)
+      raise 'Lease already created' if @outstanding_lease
+
       @create_records << metadata
       pp("create_record #{self.object_id} #{metadata}")
     end
 
     def destroy_record(metadata)
+      raise 'Lease already created' if @outstanding_lease
+
       @destroy_records << metadata
       pp("destroy_record #{self.object_id} #{metadata}")
     end
-- 
GitLab


From 17bf308ecee6c4fa5cb8cc6a5ed9cd891ad99051 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Fri, 10 Oct 2025 00:39:27 +0800
Subject: [PATCH 20/29] Minimal changes to make this work

---
 app/models/concerns/cells/unique.rb | 2 +-
 app/models/gpg_key.rb               | 4 ++--
 app/models/gpg_key_subkey.rb        | 3 ++-
 app/models/namespace.rb             | 3 ---
 app/models/user.rb                  | 4 ++--
 5 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 8ef94a87a8d25b..19cf8311e1cae7 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -59,7 +59,7 @@ def claim_save_changes
           transaction_record.destroy_record(default_metadata.merge({
             bucket: {
               type: config[:type],
-              value: attribute_in_database(attribute)
+              value: was
             }
           }))
         end
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index 7e09ddc12cc858..e3e76132c07042 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -11,7 +11,7 @@ class GpgKey < ApplicationRecord
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
   cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
 
-  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED
+  cells_config subject_type: SUBJECT_TYPE::USER, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :user_id
 
   cells_unique_has_many :subkeys
 
@@ -20,7 +20,7 @@ class GpgKey < ApplicationRecord
 
   belongs_to :user
   has_many :gpg_signatures, class_name: 'CommitSignatures::GpgSignature'
-  has_many :subkeys, class_name: 'GpgKeySubkey'
+  has_many :subkeys, class_name: 'GpgKeySubkey', dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent -- needed to unclaim
 
   scope :with_subkeys, -> { includes(:subkeys) }
   scope :externally_invalid, -> { where(externally_verified: false) }
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 161c9db5c15555..6716ab33398ca3 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -5,7 +5,8 @@ class GpgKeySubkey < ApplicationRecord
   include Cells::Unique
 
   cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
-  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED
+
+  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :gpg_key_id
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/namespace.rb b/app/models/namespace.rb
index 5969ce15b6de4b..2740c045c76fb1 100644
--- a/app/models/namespace.rb
+++ b/app/models/namespace.rb
@@ -22,9 +22,6 @@ class Namespace < ApplicationRecord
   include UseSqlFunctionForPrimaryKeyLookups
   include SafelyChangeColumnDefault
   include Todoable
-  include Cells::Unique
-
-  cells_config type: SUBJECT_TYPE::GROUP
 
   extend Gitlab::Utils::Override
 
diff --git a/app/models/user.rb b/app/models/user.rb
index b43e079a20c3a1..6c1ec69d9ff567 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -173,9 +173,9 @@ def update_tracked_fields!(request)
   has_many :expired_today_and_unnotified_keys, -> { expired_today_and_not_notified }, class_name: 'Key'
   has_many :expiring_soon_and_unnotified_keys, -> { expiring_soon_and_not_notified }, class_name: 'Key'
   has_many :deploy_keys, -> { where(type: 'DeployKey') }, dependent: :nullify # rubocop:disable Cop/ActiveRecordDependent
-  has_many :gpg_keys
+  has_many :gpg_keys, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent -- needed to unclaim
 
-  has_many :emails
+  has_many :emails, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent -- needed to unclaim
   has_many :personal_access_tokens, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
   has_many :expiring_soon_and_unnotified_personal_access_tokens, -> { expiring_and_not_notified_without_impersonation }, class_name: 'PersonalAccessToken'
 
-- 
GitLab


From e86e153ca31a2fb03c77964457266fb5d95f8dd4 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Tue, 14 Oct 2025 01:06:59 +0800
Subject: [PATCH 21/29] Cleanup and apply a few straightforward suggestions

---
 app/models/application_record.rb              |   6 -
 app/models/cells/outstanding_lease.rb         |  22 +-
 app/models/cells/transaction_record.rb        |  14 +-
 app/models/concerns/cells/unique.rb           |  79 +++---
 app/models/current.rb                         |  16 +-
 app/models/email.rb                           |   4 +-
 app/models/gpg_key.rb                         |  10 +-
 app/models/gpg_key_subkey.rb                  |   4 +-
 app/models/key.rb                             |   6 +-
 app/models/organizations/organization.rb      |   4 +-
 app/models/route.rb                           |   4 +-
 app/models/user.rb                            |   9 +-
 config/initializers/doorkeeper.rb             |   2 +-
 .../concerns/cells/inspect_changes_spec.rb    | 201 --------------
 .../cells/transaction_callback_spec.rb        | 253 ------------------
 spec/models/concerns/cells/unique_spec.rb     |   8 +-
 spec/models/email_spec.rb                     |   4 +-
 spec/models/gpg_key_spec.rb                   |   4 +-
 spec/models/gpg_key_subkey_spec.rb            |   4 +-
 spec/models/key_spec.rb                       |   4 +-
 .../models/organizations/organization_spec.rb |   4 +-
 spec/models/route_spec.rb                     |   4 +-
 spec/models/user_spec.rb                      |   4 +-
 .../cells/unique_shared_examples.rb           |   2 +-
 24 files changed, 99 insertions(+), 573 deletions(-)
 delete mode 100644 spec/models/concerns/cells/inspect_changes_spec.rb
 delete mode 100644 spec/models/concerns/cells/transaction_callback_spec.rb

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index c3c5379150d366..7cec15efa8d468 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -165,12 +165,6 @@ def self.sharding_keys
     @sharding_keys ||= Gitlab::Database::Dictionary.entry(table_name).sharding_key || {}
   end
 
-  def sharding_keys
-    @sharding_keys ||= self.class.sharding_keys.transform_keys do |key|
-      public_send(key) # rubocop:disable GitlabSecurity/PublicSend -- this works by definition
-    end.invert
-  end
-
   def readable_by?(user)
     Ability.allowed?(user, :"read_#{to_ability_name}", self)
   end
diff --git a/app/models/cells/outstanding_lease.rb b/app/models/cells/outstanding_lease.rb
index caf509046270a0..2bf3764b2d4f65 100644
--- a/app/models/cells/outstanding_lease.rb
+++ b/app/models/cells/outstanding_lease.rb
@@ -4,6 +4,10 @@ module Cells
   class OutstandingLease < ApplicationRecord
     self.primary_key = :uuid
 
+    def self.claim_service
+      ::Gitlab::TopologyServiceClient::ClaimService.instance # rubocop:disable CodeReuse/ServiceClass -- this is a gRPC client
+    end
+
     def self.create_from_request!(create_records:, destroy_records:)
       req = Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateRequest.new(
         create_records: create_records,
@@ -12,10 +16,10 @@ def self.create_from_request!(create_records:, destroy_records:)
       )
 
       pp("send_begin_update req #{req}")
-      res = self.claim_service.begin_update(req)
+      res = claim_service.begin_update(req)
 
       pp("send_begin_update res #{res.lease_uuid}")
-      self.create!(uuid: res.lease_uuid.value)
+      create!(uuid: res.lease_uuid.value)
     end
 
     def send_commit_update!
@@ -24,12 +28,12 @@ def send_commit_update!
         cell_id: self.class.claim_service.cell_id
       )
 
-      pp("send_commit_update #{self.object_id} #{req}")
+      pp("send_commit_update #{object_id} #{req}")
       self.class.claim_service.commit_update(
         req
       )
 
-      pp("send_commit_update done #{self.object_id} #{req}")
+      pp("send_commit_update done #{object_id} #{req}")
     end
 
     def send_rollback_update!
@@ -40,19 +44,13 @@ def send_rollback_update!
         cell_id: self.class.claim_service.cell_id
       )
 
-      pp("send_rollback_update #{self.object_id} #{req}")
+      pp("send_rollback_update #{object_id} #{req}")
 
       self.class.claim_service.rollback_update(
         req
       )
 
-      pp("send_rollback_update done #{self.object_id} #{req}")
-    end
-
-    private
-
-    def self.claim_service
-      @claim_service ||= ::Gitlab::TopologyServiceClient::ClaimService.instance
+      pp("send_rollback_update done #{object_id} #{req}")
     end
   end
 end
diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
index 9e99ee6e917579..bade09dcffe2de 100644
--- a/app/models/cells/transaction_record.rb
+++ b/app/models/cells/transaction_record.rb
@@ -9,12 +9,12 @@ class TransactionRecord
     # Current? but what if nested transactions?
     # what is the transactions gets rolled back?
     def self.current_transaction(connection)
-      return unless Current.cells_claim_leases?
+      return unless Current.cells_claims_leases?
 
       # The Cells::OutstandingLease requires a transaction to be open
       # to ensure that the lease is only created if the transaction
       # within a transaction and not outside of one
-      unless connection.current_transaction
+      unless connection.current_transaction # rubocop:disable Style/IfUnlessModifier -- this is more clear
         raise 'The Cells::TransactionRecord requires transaction to be open'
       end
 
@@ -42,14 +42,14 @@ def create_record(metadata)
       raise 'Lease already created' if @outstanding_lease
 
       @create_records << metadata
-      pp("create_record #{self.object_id} #{metadata}")
+      pp("create_record #{object_id} #{metadata}")
     end
 
     def destroy_record(metadata)
       raise 'Lease already created' if @outstanding_lease
 
       @destroy_records << metadata
-      pp("destroy_record #{self.object_id} #{metadata}")
+      pp("destroy_record #{object_id} #{metadata}")
     end
 
     # Callbacks required by the ActiveRecord::ConnectionAdapters::TransactionState
@@ -57,7 +57,7 @@ def destroy_record(metadata)
     def trigger_transactional_callbacks?
       true
     end
-  
+
     def before_committed!
       raise 'Already done' if @done
       raise 'Already created lease' if @outstanding_lease
@@ -69,7 +69,7 @@ def before_committed!
       )
     end
 
-    def rolledback!(force_restore_state: false, should_run_callbacks: true)
+    def rolledback!(force_restore_state: false, should_run_callbacks: true) # rubocop:disable Lint/UnusedMethodArgument -- this needs to follow the interface
       raise 'Already done' if @done
 
       # It is possible that lease might be not created yet,
@@ -81,7 +81,7 @@ def rolledback!(force_restore_state: false, should_run_callbacks: true)
       @done = true
     end
 
-    def committed!(should_run_callbacks: true)
+    def committed!(should_run_callbacks: true) # rubocop:disable Lint/UnusedMethodArgument -- this needs to follow the interface
       raise 'Already done' if @done
       raise 'No lease created' unless @outstanding_lease
 
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index 19cf8311e1cae7..9ef1bcef908e68 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -4,52 +4,47 @@ module Cells
   module Unique
     extend ActiveSupport::Concern
 
-    BUCKET_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Bucket::Type
-    SUBJECT_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Subject::Type
-    SOURCE_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Source::Type
+    CLAIMS_BUCKET_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Bucket::Type
+    CLAIMS_SUBJECT_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Subject::Type
+    CLAIMS_SOURCE_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Source::Type
 
     included do
-      after_save :claim_save_changes
-      before_destroy :claim_destroy_changes
+      after_save :cells_claims_save_changes
+      before_destroy :cells_claims_destroy_changes
+
+      class_attribute :cells_claims_subject_type, instance_accessor: false
+      class_attribute :cells_claims_subject_key, instance_accessor: false
+      class_attribute :cells_claims_source_type, instance_accessor: false
+      class_attribute :cells_claims_attributes, instance_accessor: false, default: {}.freeze
     end
 
     class_methods do
-      attr_reader :cells_subject_type, :cells_subject_key, :cells_source_type
-
-      def cells_unique_attributes
-        @cells_unique_attributes ||= {}
+      def cells_claims_metadata(subject_type:, subject_key: nil, source_type: nil)
+        self.cells_claims_subject_type = subject_type
+        self.cells_claims_subject_key = subject_key || :organization_id
+        self.cells_claims_source_type = source_type ||
+          Gitlab::Cells::TopologyService::Claims::V1::Source::Type
+            .const_get("RAILS_TABLE_#{table_name.upcase}", false)
       end
 
-      def cells_unique_has_many_associations
-        @cells_unique_has_many_associations ||= []
-      end
-
-      def cells_config(subject_type:, subject_key: nil, source_type: nil)
-        @cells_subject_type = subject_type
-        @cells_subject_key = subject_key || :organization_id
-        @cells_source_type = source_type || const_get("Gitlab::Cells::TopologyService::Claims::V1::Source::Type::RAILS_TABLE_#{table_name.upcase}")
-      end
-
-      def cells_unique_attribute(name, type:)
-        cells_unique_attributes[name] = { type: type }
-      end
-
-      def cells_unique_has_many(name)
-        cells_unique_has_many_associations << name
+      def cells_claims_attribute(name, type:)
+        self.cells_claims_attributes = cells_claims_attributes
+          .merge(name => { type: type })
+          .freeze
       end
     end
 
     private
 
-    def claim_save_changes
+    def cells_claims_save_changes
       transaction_record = ::Cells::TransactionRecord.current_transaction(connection)
       return unless transaction_record
 
-      pp("claim_save_changes #{self.object_id}: #{self.class.name} #{self.id}")
+      pp("claim_save_changes #{object_id}: #{self.class.name} #{id}")
 
-      default_metadata = claim_metadata
+      default_metadata = cells_claims_default_metadata
 
-      self.class.cells_unique_attributes.each do |attribute, config|
+      self.class.cells_claims_attributes.each do |attribute, config|
         next unless saved_change_to_attribute?(attribute)
 
         was, is = saved_change_to_attribute(attribute)
@@ -64,45 +59,45 @@ def claim_save_changes
           }))
         end
 
-        if !is.nil?
+        if !is.nil? # rubocop:disable Style/NegatedIf, Style/Next -- this is more clear
           transaction_record.create_record(default_metadata.merge(
             bucket: {
               type: config[:type],
-              value: self.public_send(attribute)
+              value: public_send(attribute) # rubocop:disable GitlabSecurity/PublicSend -- developer hard coded
             }
           ))
         end
       end
     end
 
-    def claim_destroy_changes
+    def cells_claims_destroy_changes
       transaction_record = ::Cells::TransactionRecord.current_transaction(connection)
       return unless transaction_record
 
-      pp("claim_destroy_changes #{self.object_id}: #{self.class.name} #{self.id}")
+      pp("claim_destroy_changes #{object_id}: #{self.class.name} #{id}")
 
-      default_metadata = claim_metadata
+      default_metadata = cells_claims_default_metadata
 
-      self.class.cells_unique_attributes.each do |attribute, config|
+      self.class.cells_claims_attributes.each do |attribute, config|
         transaction_record.destroy_record(default_metadata.merge(
           bucket: {
             type: config[:type],
-            value: self.public_send(attribute)
+            value: public_send(attribute) # rubocop:disable GitlabSecurity/PublicSend -- developer hard coded
           }
         ))
       end
     end
 
-    def claim_metadata
+    def cells_claims_default_metadata
       {
         subject: {
-          type: self.class.cells_subject_type,
-          id: self.public_send(self.class.cells_subject_key)
+          type: self.class.cells_claims_subject_type,
+          id: public_send(self.class.cells_claims_subject_key) # rubocop:disable GitlabSecurity/PublicSend -- developer hard coded
         },
         source: {
-          type: self.class.cells_source_type,
-          rails_primary_key_id: self.public_send(self.class.primary_key)
-        },
+          type: self.class.cells_claims_source_type,
+          rails_primary_key_id: public_send(self.class.primary_key) # rubocop:disable GitlabSecurity/PublicSend -- this should work by definition
+        }
       }
     end
   end
diff --git a/app/models/current.rb b/app/models/current.rb
index c8c3ff32280ea1..e9ee743554ba04 100644
--- a/app/models/current.rb
+++ b/app/models/current.rb
@@ -18,14 +18,16 @@ def message
   attribute :organization, :organization_assigned
   attribute :token_info
 
-  attribute :cells_claim_leases
+  attribute :cells_claims_leases
 
-  def cells_claim_leases?
-    if cells_claim_leases.nil?
-      self.cells_claim_leases = Gitlab.config.cell.enabled && Feature.enabled?(:cells_unique_claims)
+  def cells_claims_leases?
+    if cells_claims_leases.nil?
+      self.cells_claims_leases =
+        Gitlab.config.cell.enabled &&
+        Feature.enabled?(:cells_unique_claims) # rubocop:disable Gitlab/FeatureFlagWithoutActor -- We can't easily tie this to an actor
     end
 
-    cells_claim_leases
+    cells_claims_leases
   end
 
   def organization=(value)
@@ -55,10 +57,6 @@ def organization
     super
   end
 
-  def reset_feature_cells_unique_claims
-    self.feature_cells_unique_claims = Feature.enabled?(:cells_unique_claims) # rubocop:disable Gitlab/FeatureFlagWithoutActor -- We can't easily tie this to an actor
-  end
-
   private
 
   # Do not allow to reset this
diff --git a/app/models/email.rb b/app/models/email.rb
index ba00f710da0a85..ce1a71c77a93c9 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -5,9 +5,9 @@ class Email < ApplicationRecord
   include Gitlab::SQL::Pattern
   include Cells::Unique
 
-  cells_unique_attribute :email, type: BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
+  cells_claims_attribute :email, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
-  cells_config subject_type: SUBJECT_TYPE::USER, subject_key: :user_id
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::USER, subject_key: :user_id
 
   belongs_to :user, optional: false
   belongs_to :banned_user, class_name: '::Users::BannedUser', foreign_key: 'user_id', inverse_of: 'emails'
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index e3e76132c07042..f0f118503017fe 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -7,13 +7,11 @@ class GpgKey < ApplicationRecord
   include ShaAttribute
   include Cells::Unique
 
-  cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED
-  cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
-  cells_unique_attribute :primary_keyid, type: BUCKET_TYPE::UNSPECIFIED
+  cells_claims_attribute :key, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
+  cells_claims_attribute :fingerprint, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
+  cells_claims_attribute :primary_keyid, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
 
-  cells_config subject_type: SUBJECT_TYPE::USER, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :user_id
-
-  cells_unique_has_many :subkeys
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::USER, source_type: CLAIMS_SOURCE_TYPE::UNSPECIFIED, subject_key: :user_id
 
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 6716ab33398ca3..12a2c931f25868 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -4,9 +4,9 @@ class GpgKeySubkey < ApplicationRecord
   include ShaAttribute
   include Cells::Unique
 
-  cells_unique_attribute :fingerprint, type: BUCKET_TYPE::UNSPECIFIED
+  cells_claims_attribute :fingerprint, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
 
-  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :gpg_key_id
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :gpg_key_id
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index 6fcf3ace91045f..aabae4f00fcb59 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -10,10 +10,10 @@ class Key < ApplicationRecord
   include CreatedAtFilterable
   include Cells::Unique
 
-  cells_unique_attribute :key, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
-  cells_unique_attribute :fingerprint_sha256, type: BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
+  cells_claims_attribute :key, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
+  cells_claims_attribute :fingerprint_sha256, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
 
-  cells_config subject_type: SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: CLAIMS_SOURCE_TYPE::UNSPECIFIED
 
   sha256_attribute :fingerprint_sha256
 
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index fedfd5931eff69..8ba83c7376c97d 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -10,9 +10,9 @@ class Organization < ApplicationRecord
     include Organizations::Isolatable
     include Cells::Unique
 
-    cells_unique_attribute :path, type: BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
+    cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
 
-    cells_config subject_type: SUBJECT_TYPE::ORGANIZATION
+    cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::ORGANIZATION
 
     DEFAULT_ORGANIZATION_ID = 1
 
diff --git a/app/models/route.rb b/app/models/route.rb
index 38839ee1df89a5..185127a2e0c34a 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -6,9 +6,9 @@ class Route < ApplicationRecord
   include EachBatch
   include Cells::Unique
 
-  cells_unique_attribute :path, type: BUCKET_TYPE::ROUTES
+  cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::ROUTES
 
-  cells_config subject_type: SUBJECT_TYPE::GROUP, subject_key: :namespace_id
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::GROUP, subject_key: :namespace_id
 
   belongs_to :source, polymorphic: true, inverse_of: :route # rubocop:disable Cop/PolymorphicAssociations
   belongs_to :namespace, inverse_of: :namespace_route
diff --git a/app/models/user.rb b/app/models/user.rb
index 6c1ec69d9ff567..4b345a71ffc323 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -38,12 +38,9 @@ class User < ApplicationRecord
   include Users::DependentAssociations
   include Cells::Unique
 
-  # cells_unique_attribute :email # This clashes with emails.email
-  cells_unique_attribute :username, type: BUCKET_TYPE::UNSPECIFIED # TODO: USERNAME?
-  cells_unique_has_many :gpg_keys
-  cells_unique_has_many :emails
-
-  cells_config subject_type: SUBJECT_TYPE::USER
+  # cells_claims_attribute :email # This clashes with emails.email
+  cells_claims_attribute :username, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: USERNAME?
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::USER
 
   ignore_column %i[role skype], remove_after: '2025-09-18', remove_with: '18.4'
 
diff --git a/config/initializers/doorkeeper.rb b/config/initializers/doorkeeper.rb
index 625d56439145c7..809541d888cc84 100644
--- a/config/initializers/doorkeeper.rb
+++ b/config/initializers/doorkeeper.rb
@@ -157,7 +157,7 @@
   Doorkeeper::Application.class_eval do
     include Cells::Unique
 
-    cells_unique_attribute :uid, type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
+    cells_claims_attribute :uid, type: Cells::Unique::CLAIMS_BUCKET_TYPE::UNSPECIFIED
   end
 
   # Following doorkeeper monkey patches are to identify the organization on best effort basis
diff --git a/spec/models/concerns/cells/inspect_changes_spec.rb b/spec/models/concerns/cells/inspect_changes_spec.rb
deleted file mode 100644
index 0b0a5264e9f94b..00000000000000
--- a/spec/models/concerns/cells/inspect_changes_spec.rb
+++ /dev/null
@@ -1,201 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Cells::InspectChanges, feature_category: :cell do
-  let(:user) { create(:user) }
-  let(:inspector) { described_class.new(user, destroy: destroy) }
-  let(:destroy) { false }
-
-  let(:dummy_class) do
-    Class.new(ActiveRecord::Base) do
-      def self.name
-        'Namespace::ClassName'
-      end
-
-      self.table_name = 'users'
-      include Cells::Unique
-
-      cells_unique_attribute :name,
-        type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
-      cells_unique_attribute :email,
-        type: Cells::Unique::BUCKET_TYPE::UNSPECIFIED
-
-      def attribute_in_database(column)
-        public_send(column)
-      end
-    end
-  end
-
-  describe '#inspect_changes' do
-    context 'when record is not Unique' do
-      let_it_be(:project) { create(:project) }
-      let(:inspector) { described_class.new(project) }
-
-      it 'returns without processing' do
-        expect(Cells::TransactionCallback).not_to receive(:visited)
-        inspector.inspect_changes
-      end
-    end
-
-    context 'when record is Unique but already inspected' do
-      before do
-        visited = { ['users', user.id] => { 'name' => %w[Old New] } }
-        allow(Cells::TransactionCallback).to receive(:visited).and_return(visited)
-      end
-
-      it 'returns without reprocessing' do
-        expect(inspector).not_to receive(:inspect_record)
-        inspector.inspect_changes
-      end
-    end
-
-    context 'when record needs inspection' do
-      before do
-        allow(Cells::TransactionCallback).to receive(:visited).and_return(visited)
-      end
-
-      context 'with transaction' do
-        let(:visited) { {} }
-
-        it 'adds record to visited without calling claim_and_commit' do
-          expect(Cells::Unique).not_to receive(:claim_and_commit)
-
-          inspector.inspect_changes
-
-          expect(visited.keys).to match_array([
-            ['users', user.id],
-            ['emails', user.email_ids.first],
-            ['organizations', user.organization_id]
-          ])
-        end
-      end
-
-      context 'without transaction' do
-        let(:visited) { nil }
-
-        it 'calls claim_and_commit' do
-          expect(Cells::Unique).to receive(:claim_and_commit).at_least(:once)
-          inspector.inspect_changes
-        end
-      end
-    end
-  end
-
-  describe 'change tracking' do
-    context 'with associations' do
-      before do
-        allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
-        allow(Cells::Unique).to receive(:claim_and_commit)
-      end
-
-      context 'when loaded associations' do
-        before do
-          create(:email, user: user)
-          create(:email, user: user)
-        end
-
-        it 'recursively inspects associations' do
-          inspector.inspect_changes
-
-          visited = inspector.send(:visited)
-          expect(visited.keys).to include(['users', user.id])
-          # expect(visited.keys).to include(['emails', user.email_ids.first]) TODO
-          # expect(visited.keys).to include(['emails', user.email_ids.second]) TODO
-        end
-      end
-
-      context 'with unloaded associations' do
-        let!(:email) { create(:email, user: user) }
-
-        before do
-          user.reload # Ensure association is not loaded
-        end
-
-        it 'skips unloaded associations' do
-          inspector.inspect_changes
-
-          visited = inspector.send(:visited)
-          expect(visited.keys).to eq([['users', user.id]])
-        end
-      end
-    end
-
-    context 'without associations' do
-      let(:user) { dummy_class.new(name: 'John', email: 'john@example.com') }
-
-      before do
-        allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
-        allow(Cells::Unique).to receive(:claim_and_commit)
-      end
-
-      context 'with saved changes' do
-        before do
-          allow(user).to receive(:saved_changes).and_return(
-            {
-              'name' => %w[Jane John],
-              'email' => ['jane@example.com',
-                'john@example.com'],
-              'age' => [29, 30]
-            })
-        end
-
-        it 'tracks only unique attribute changes' do
-          inspector.inspect_changes
-
-          inspected = inspector.send(:inspected)
-          expect(inspected['name']).to eq(%w[Jane John])
-          expect(inspected['email']).to eq(['jane@example.com', 'john@example.com'])
-          expect(inspected['age']).to be_nil
-        end
-      end
-
-      context 'with destroyed record' do
-        let(:destroy) { true }
-
-        it 'tracks deleted attribute values' do
-          inspector.inspect_changes
-
-          inspected = inspector.send(:inspected)
-          expect(inspected['name']).to eq(['John'])
-          expect(inspected['email']).to eq(['john@example.com'])
-        end
-      end
-    end
-  end
-
-  describe 'special cases' do
-    before do
-      allow(Cells::TransactionCallback).to receive(:visited).and_return(nil)
-      allow(Cells::Unique).to receive(:claim_and_commit)
-    end
-
-    context 'with GpgKey transformation' do
-      let(:gpg_key) { create(:gpg_key) }
-      let(:inspector) { described_class.new(gpg_key) }
-
-      it 'uses actual stored value after transformation' do
-        expect(gpg_key.fingerprint).to eq(GpgHelpers::User1.fingerprint)
-
-        gpg_key.update!(key: GpgHelpers::User2.public_key)
-
-        inspector.inspect_changes
-
-        inspected = inspector.send(:inspected)
-        expect(inspected['fingerprint']).to match_array(
-          [GpgHelpers::User1.fingerprint.downcase, GpgHelpers::User2.fingerprint.downcase]
-        )
-      end
-    end
-
-    context 'with nil association' do
-      before do
-        user.user_detail # Load nil association
-      end
-
-      it 'handles gracefully' do
-        expect { inspector.inspect_changes }.not_to raise_error
-      end
-    end
-  end
-end
diff --git a/spec/models/concerns/cells/transaction_callback_spec.rb b/spec/models/concerns/cells/transaction_callback_spec.rb
deleted file mode 100644
index bf1fa134551f82..00000000000000
--- a/spec/models/concerns/cells/transaction_callback_spec.rb
+++ /dev/null
@@ -1,253 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Cells::TransactionCallback, feature_category: :cell do
-  let(:adapter_class) do
-    Class.new do
-      prepend Cells::TransactionCallback
-
-      attr_reader :transaction_state
-
-      def initialize
-        @transaction_state = []
-      end
-
-      def begin_isolated_db_transaction(...)
-        @transaction_state << :isolated_begin
-      end
-
-      def begin_db_transaction
-        @transaction_state << :begin
-      end
-
-      def commit_db_transaction
-        @transaction_state << :commit
-      end
-
-      def rollback_db_transaction
-        @transaction_state << :rollback
-      end
-    end
-  end
-
-  let(:connection) { adapter_class.new }
-
-  around do |example|
-    original_value = Thread.current.thread_variable_get(described_class::THREAD_VARIABLE_NAME)
-    Thread.current.thread_variable_set(described_class::THREAD_VARIABLE_NAME, nil)
-    example.run
-  ensure
-    Thread.current.thread_variable_set(described_class::THREAD_VARIABLE_NAME, original_value)
-  end
-
-  describe 'thread-local visited stack' do
-    it 'initializes empty stack per thread' do
-      expect(described_class.visited).to be_nil
-      expect(described_class.visited_stack).to eq([])
-    end
-
-    it 'returns first element as visited' do
-      described_class.visited_stack << { first: 'data' }
-      described_class.visited_stack << { second: 'data' }
-
-      expect(described_class.visited).to eq({ first: 'data' })
-    end
-
-    it 'maintains separate stacks per thread' do
-      described_class.visited_stack << { main: 'thread' }
-
-      other_thread_value = Thread.new do
-        described_class.visited_stack << { other: 'thread' }
-        described_class.visited
-      end.value
-
-      expect(described_class.visited).to eq({ main: 'thread' })
-      expect(other_thread_value).to eq({ other: 'thread' })
-    end
-  end
-
-  describe 'transaction lifecycle' do
-    describe '#begin_db_transaction' do
-      it 'pushes empty hash to stack' do
-        connection.begin_db_transaction
-
-        expect(described_class.visited_stack).to eq([{}])
-        expect(connection.transaction_state).to include(:begin)
-      end
-
-      it 'handles nested transactions' do
-        connection.begin_db_transaction
-        connection.begin_db_transaction
-
-        expect(described_class.visited_stack.size).to eq(2)
-        expect(described_class.visited_stack).to eq([{}, {}])
-      end
-    end
-
-    describe '#begin_isolated_db_transaction' do
-      it 'pushes empty hash to stack' do
-        connection.begin_isolated_db_transaction
-
-        expect(described_class.visited_stack).to eq([{}])
-        expect(connection.transaction_state).to include(:isolated_begin)
-      end
-    end
-
-    describe '#commit_db_transaction' do
-      before do
-        connection.begin_db_transaction
-        described_class.visited_stack.last[['users', 1]] = { 'email' => %w[old new] }
-      end
-
-      it 'claims and commits with visited data' do
-        expect(Cells::Unique).to receive(:claim_and_commit)
-          .with([{ 'email' => %w[old new] }])
-          .and_yield
-
-        connection.commit_db_transaction
-
-        expect(described_class.visited_stack).to be_empty
-        expect(connection.transaction_state).to include(:commit)
-      end
-
-      it 'maintains stack if commit fails' do
-        connection.begin_db_transaction
-        described_class.visited_stack.last[['users', 1]] = { 'email' => %w[old new] }
-
-        expect(Cells::Unique).to receive(:claim_and_commit).and_yield
-
-        allow(connection).to receive(:commit_db_transaction).and_wrap_original do |original_method, *args|
-          original_method.call(*args)
-          raise ActiveRecord::StatementInvalid, 'Connection lost'
-        end
-
-        expect do
-          connection.commit_db_transaction
-        end.to raise_error(ActiveRecord::StatementInvalid)
-
-        expect(described_class.visited_stack.size).to eq(1)
-      end
-
-      it 'handles empty visited data' do
-        described_class.visited_stack.last.clear
-
-        expect(Cells::Unique).to receive(:claim_and_commit).with([]).and_yield
-
-        connection.commit_db_transaction
-        expect(described_class.visited_stack).to be_empty
-      end
-    end
-
-    describe '#rollback_db_transaction' do
-      context 'with balanced transaction' do
-        before do
-          connection.begin_db_transaction
-        end
-
-        it 'rolls back and pops stack' do
-          expect(Cells::Unique).to receive(:rollback).and_yield
-
-          result = connection.rollback_db_transaction
-
-          expect(result).to be true
-          expect(described_class.visited_stack).to be_empty
-          expect(connection.transaction_state).to include(:rollback)
-        end
-      end
-
-      context 'with unbalanced rollback' do
-        it 'handles rollback without corresponding begin' do
-          expect(Cells::Unique).to receive(:rollback).and_yield
-
-          result = connection.rollback_db_transaction
-
-          expect(result).to be false
-          expect(described_class.visited_stack).to be_empty
-          expect(connection.transaction_state).to include(:rollback)
-        end
-      end
-
-      context 'with nested transactions' do
-        before do
-          connection.begin_db_transaction
-          connection.begin_db_transaction
-          described_class.visited_stack[0][['users', 1]] = 'outer'
-          described_class.visited_stack[1][['users', 2]] = 'inner'
-        end
-
-        it 'only rolls back innermost transaction' do
-          expect(described_class.visited_stack.size).to eq(2)
-          expect(Cells::Unique).to receive(:rollback).and_yield
-
-          connection.rollback_db_transaction
-
-          expect(described_class.visited_stack.size).to eq(1)
-          expect(described_class.visited_stack.first).to eq({ ['users', 1] => 'outer' })
-        end
-      end
-    end
-  end
-
-  describe 'integration scenarios' do
-    it 'handles successful transaction flow' do
-      connection.begin_db_transaction
-
-      described_class.visited[['users', 1]] = { 'name' => %w[old new] }
-
-      expect(Cells::Unique).to receive(:claim_and_commit)
-        .with([{ 'name' => %w[old new] }])
-        .and_yield
-
-      connection.commit_db_transaction
-
-      expect(described_class.visited_stack).to be_empty
-    end
-
-    it 'handles rollback after partial work' do
-      connection.begin_db_transaction
-
-      described_class.visited[['users', 1]] = { 'name' => %w[old new] }
-
-      expect(Cells::Unique).to receive(:rollback).and_yield
-
-      connection.rollback_db_transaction
-
-      expect(described_class.visited_stack).to be_empty
-    end
-
-    it 'handles nested transaction with inner rollback' do
-      connection.begin_db_transaction
-      described_class.visited[['outer', 1]] = 'data'
-
-      connection.begin_isolated_db_transaction
-      described_class.visited_stack.last[['inner', 1]] = 'data'
-
-      expect(Cells::Unique).to receive(:rollback).and_yield
-      connection.rollback_db_transaction
-
-      expect(Cells::Unique).to receive(:claim_and_commit)
-        .with(['data'])
-        .and_yield
-
-      connection.commit_db_transaction
-
-      expect(described_class.visited_stack).to be_empty
-    end
-  end
-
-  describe 'error handling' do
-    it 'maintains consistency when claim service fails' do
-      connection.begin_db_transaction
-      described_class.visited[['users', 1]] = 'data'
-
-      expect(Cells::Unique).to receive(:claim_and_commit)
-        .and_raise(StandardError, 'Claim service unavailable')
-
-      expect { connection.commit_db_transaction }
-        .to raise_error(StandardError, 'Claim service unavailable')
-
-      expect(described_class.visited_stack.size).to eq(1)
-    end
-  end
-end
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index a3a95fa7d0a0bf..e8c798c9c53934 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -8,13 +8,13 @@
 
   describe 'configuration' do
     it 'sets and retrieves cell configuration' do
-      test_klass.cells_config(type: :user, source_type: :rails_table_users)
-      test_klass.cells_unique_attribute(:email, type: :string)
-      test_klass.cells_unique_attribute(:username, type: :string)
+      test_klass.cells_claims_metadata(type: :user, source_type: :rails_table_users)
+      test_klass.cells_claims_attribute(:email, type: :string)
+      test_klass.cells_claims_attribute(:username, type: :string)
 
       expect(test_klass.cells_type).to eq(:user)
       expect(test_klass.cells_source_type).to eq(:rails_table_users)
-      expect(test_klass.cells_unique_attributes).to eq(email: { type: :string }, username: { type: :string })
+      expect(test_klass.cells_claims_attributes).to eq(email: { type: :string }, username: { type: :string })
     end
   end
 
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index d8d8375fb47a46..494d0918ae7cd0 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -6,8 +6,8 @@
   let_it_be(:user) { create(:user) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_EMAILS,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_EMAILS,
     unique_attributes: [:email]
 
   describe 'modules' do
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index 0af57bfa542e7e..641afd571bd0cf 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -6,8 +6,8 @@
   subject { build(:gpg_key) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint, :primary_keyid]
 
   describe "associations" do
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index 0eda583f420ec6..fb8ca50d487846 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -16,7 +16,7 @@
   end
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index ea2aa7c444dcc2..edd53cc27ec124 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -6,8 +6,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::SOURCE_TYPE::UNSPECIFIED,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint_sha256]
 
   describe "Associations" do
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index 634487c6c42412..4b5b2d47c02ea6 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -6,8 +6,8 @@
   let_it_be_with_refind(:organization) { create(:organization) }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::ORGANIZATION,
-    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::ORGANIZATION,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
     unique_attributes: [:path]
 
   describe 'associations' do
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index a15e7e93d5ff63..5dbf8dc8266d4e 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -7,8 +7,8 @@
   let(:route) { group.route }
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_ROUTES,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ROUTES,
     unique_attributes: [:path]
 
   describe 'relationships' do
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index 81442a6aeeec4b..a20ff3119be565 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -50,8 +50,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells unique model',
-    type: Cells::Unique::SUBJECT_TYPE::USER,
-    source_type: Cells::Unique::SOURCE_TYPE::RAILS_TABLE_USERS,
+    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::USER,
+    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_USERS,
     unique_attributes: [:username]
 
   describe 'modules' do
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
index 7c16a72f8a1f17..f3bc7c35c24589 100644
--- a/spec/support/shared_examples/cells/unique_shared_examples.rb
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -11,7 +11,7 @@
 
   it 'has the expected unique attributes' do
     unique_attributes.each do |attr_name|
-      expect(described_class.cells_unique_attributes).to have_key(attr_name)
+      expect(described_class.cells_claims_attributes).to have_key(attr_name)
     end
   end
 end
-- 
GitLab


From 82bb04a3b394a4d3065dd2bf6119ad978bacf06f Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Tue, 14 Oct 2025 15:26:46 +0200
Subject: [PATCH 22/29] Rename Unique concern to Claimable

---
 app/models/concerns/cells/{unique.rb => claimable.rb}       | 2 +-
 app/models/email.rb                                         | 2 +-
 app/models/gpg_key.rb                                       | 5 +++--
 app/models/gpg_key_subkey.rb                                | 5 +++--
 app/models/key.rb                                           | 2 +-
 app/models/organizations/organization.rb                    | 2 +-
 app/models/route.rb                                         | 2 +-
 app/models/user.rb                                          | 2 +-
 config/initializers/doorkeeper.rb                           | 4 ++--
 .../concerns/cells/{unique_spec.rb => claimable_spec.rb}    | 2 +-
 spec/models/email_spec.rb                                   | 6 +++---
 spec/models/gpg_key_spec.rb                                 | 6 +++---
 spec/models/gpg_key_subkey_spec.rb                          | 6 +++---
 spec/models/key_spec.rb                                     | 6 +++---
 spec/models/organizations/organization_spec.rb              | 6 +++---
 spec/models/route_spec.rb                                   | 6 +++---
 spec/models/user_spec.rb                                    | 6 +++---
 .../support/shared_examples/cells/unique_shared_examples.rb | 2 +-
 18 files changed, 37 insertions(+), 35 deletions(-)
 rename app/models/concerns/cells/{unique.rb => claimable.rb} (99%)
 rename spec/models/concerns/cells/{unique_spec.rb => claimable_spec.rb} (98%)

diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/claimable.rb
similarity index 99%
rename from app/models/concerns/cells/unique.rb
rename to app/models/concerns/cells/claimable.rb
index 9ef1bcef908e68..932b5a7471d4c5 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/claimable.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Cells
-  module Unique
+  module Claimable
     extend ActiveSupport::Concern
 
     CLAIMS_BUCKET_TYPE = Gitlab::Cells::TopologyService::Claims::V1::Bucket::Type
diff --git a/app/models/email.rb b/app/models/email.rb
index ce1a71c77a93c9..7ac7a51b525b26 100644
--- a/app/models/email.rb
+++ b/app/models/email.rb
@@ -3,7 +3,7 @@
 class Email < ApplicationRecord
   include Sortable
   include Gitlab::SQL::Pattern
-  include Cells::Unique
+  include Cells::Claimable
 
   cells_claims_attribute :email, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: EMAIL?
 
diff --git a/app/models/gpg_key.rb b/app/models/gpg_key.rb
index f0f118503017fe..e55642d055446b 100644
--- a/app/models/gpg_key.rb
+++ b/app/models/gpg_key.rb
@@ -5,13 +5,14 @@ class GpgKey < ApplicationRecord
   KEY_SUFFIX = '-----END PGP PUBLIC KEY BLOCK-----'
 
   include ShaAttribute
-  include Cells::Unique
+  include Cells::Claimable
 
   cells_claims_attribute :key, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
   cells_claims_attribute :fingerprint, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
   cells_claims_attribute :primary_keyid, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
 
-  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::USER, source_type: CLAIMS_SOURCE_TYPE::UNSPECIFIED, subject_key: :user_id
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::USER, source_type: CLAIMS_SOURCE_TYPE::UNSPECIFIED,
+    subject_key: :user_id
 
   sha_attribute :primary_keyid
   sha_attribute :fingerprint
diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 12a2c931f25868..5643da92483261 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -2,11 +2,12 @@
 
 class GpgKeySubkey < ApplicationRecord
   include ShaAttribute
-  include Cells::Unique
+  include Cells::Claimable
 
   cells_claims_attribute :fingerprint, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
 
-  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED, subject_key: :gpg_key_id
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED,
+    subject_key: :gpg_key_id
 
   sha_attribute :keyid
   sha_attribute :fingerprint
diff --git a/app/models/key.rb b/app/models/key.rb
index aabae4f00fcb59..64152acea978d3 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -8,7 +8,7 @@ class Key < ApplicationRecord
   include FromUnion
   include Todoable
   include CreatedAtFilterable
-  include Cells::Unique
+  include Cells::Claimable
 
   cells_claims_attribute :key, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: KEY?
   cells_claims_attribute :fingerprint_sha256, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: KEY_FINGERPRINT?
diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index 8ba83c7376c97d..c43020209a99c4 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -8,7 +8,7 @@ class Organization < ApplicationRecord
     include Gitlab::Routing.url_helpers
     include FeatureGate
     include Organizations::Isolatable
-    include Cells::Unique
+    include Cells::Claimable
 
     cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
 
diff --git a/app/models/route.rb b/app/models/route.rb
index 185127a2e0c34a..e85d6ae9cfed6a 100644
--- a/app/models/route.rb
+++ b/app/models/route.rb
@@ -4,7 +4,7 @@ class Route < ApplicationRecord
   include CaseSensitivity
   include Gitlab::SQL::Pattern
   include EachBatch
-  include Cells::Unique
+  include Cells::Claimable
 
   cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::ROUTES
 
diff --git a/app/models/user.rb b/app/models/user.rb
index 4b345a71ffc323..f8d5ddcfeac0d8 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -36,7 +36,7 @@ class User < ApplicationRecord
   include Gitlab::InternalEventsTracking
   include Ci::PipelineScheduleOwnershipValidator
   include Users::DependentAssociations
-  include Cells::Unique
+  include Cells::Claimable
 
   # cells_claims_attribute :email # This clashes with emails.email
   cells_claims_attribute :username, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: USERNAME?
diff --git a/config/initializers/doorkeeper.rb b/config/initializers/doorkeeper.rb
index 809541d888cc84..9252a8cdf7f58a 100644
--- a/config/initializers/doorkeeper.rb
+++ b/config/initializers/doorkeeper.rb
@@ -155,9 +155,9 @@
   # Following doorkeeper monkey patches are to claim required resources
 
   Doorkeeper::Application.class_eval do
-    include Cells::Unique
+    include Cells::Claimable
 
-    cells_claims_attribute :uid, type: Cells::Unique::CLAIMS_BUCKET_TYPE::UNSPECIFIED
+    cells_claims_attribute :uid, type: Cells::Claimable::CLAIMS_BUCKET_TYPE::UNSPECIFIED
   end
 
   # Following doorkeeper monkey patches are to identify the organization on best effort basis
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/claimable_spec.rb
similarity index 98%
rename from spec/models/concerns/cells/unique_spec.rb
rename to spec/models/concerns/cells/claimable_spec.rb
index e8c798c9c53934..171aad87795c5c 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/claimable_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Cells::Unique, feature_category: :cell do
+RSpec.describe Cells::Claimable, feature_category: :cell do
   let_it_be(:user) { create(:user) }
   let(:test_klass) { ::User }
 
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index 494d0918ae7cd0..8d58611df77750 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -5,9 +5,9 @@
 RSpec.describe Email do
   let_it_be(:user) { create(:user) }
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_EMAILS,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_EMAILS,
     unique_attributes: [:email]
 
   describe 'modules' do
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index 641afd571bd0cf..f3994ccc84cdaf 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -5,9 +5,9 @@
 RSpec.describe GpgKey, feature_category: :source_code_management do
   subject { build(:gpg_key) }
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint, :primary_keyid]
 
   describe "associations" do
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index fb8ca50d487846..99e7f81063eb79 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -15,8 +15,8 @@
     it { is_expected.to validate_presence_of(:keyid) }
   end
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index edd53cc27ec124..2e6b2bc0e95750 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -5,9 +5,9 @@
 RSpec.describe Key, :mailer, feature_category: :system_access do
   it_behaves_like 'having unique enum values'
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint_sha256]
 
   describe "Associations" do
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index 4b5b2d47c02ea6..4a3111f4438eac 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -5,9 +5,9 @@
 RSpec.describe Organizations::Organization, type: :model, feature_category: :organization do
   let_it_be_with_refind(:organization) { create(:organization) }
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::ORGANIZATION,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::ORGANIZATION,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
     unique_attributes: [:path]
 
   describe 'associations' do
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index 5dbf8dc8266d4e..2bff53c8bd91a7 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -6,9 +6,9 @@
   let(:group) { create(:group, path: 'git_lab', name: 'git_lab') }
   let(:route) { group.route }
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ROUTES,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ROUTES,
     unique_attributes: [:path]
 
   describe 'relationships' do
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index a20ff3119be565..582edabb365271 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -49,9 +49,9 @@
 
   it_behaves_like 'having unique enum values'
 
-  it_behaves_like 'cells unique model',
-    type: Cells::Unique::CLAIMS_SUBJECT_TYPE::USER,
-    source_type: Cells::Unique::CLAIMS_SOURCE_TYPE::RAILS_TABLE_USERS,
+  it_behaves_like 'cells claimable model',
+    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::USER,
+    source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_USERS,
     unique_attributes: [:username]
 
   describe 'modules' do
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
index f3bc7c35c24589..883d184158ea23 100644
--- a/spec/support/shared_examples/cells/unique_shared_examples.rb
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-RSpec.shared_examples 'cells unique model' do |type:, source_type:, unique_attributes:|
+RSpec.shared_examples 'cells claimable model' do |type:, source_type:, unique_attributes:|
   it 'has the expected owner_type' do
     expect(described_class.cells_type).to eq(type)
   end
-- 
GitLab


From 1067532f80e45993fb7a6350a38f09927ca58420 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Tue, 14 Oct 2025 16:44:33 +0200
Subject: [PATCH 23/29] Fixes failing specs for claimable

---
 spec/models/concerns/cells/claimable_spec.rb | 194 ++++++++++++-------
 1 file changed, 127 insertions(+), 67 deletions(-)

diff --git a/spec/models/concerns/cells/claimable_spec.rb b/spec/models/concerns/cells/claimable_spec.rb
index 171aad87795c5c..22f23ab5c052dd 100644
--- a/spec/models/concerns/cells/claimable_spec.rb
+++ b/spec/models/concerns/cells/claimable_spec.rb
@@ -8,115 +8,175 @@
 
   describe 'configuration' do
     it 'sets and retrieves cell configuration' do
-      test_klass.cells_claims_metadata(type: :user, source_type: :rails_table_users)
+      test_klass.cells_claims_metadata(subject_type: :user, source_type: :rails_table_users)
       test_klass.cells_claims_attribute(:email, type: :string)
       test_klass.cells_claims_attribute(:username, type: :string)
 
-      expect(test_klass.cells_type).to eq(:user)
-      expect(test_klass.cells_source_type).to eq(:rails_table_users)
+      expect(test_klass.cells_claims_subject_type).to eq(:user)
+      expect(test_klass.cells_claims_source_type).to eq(:rails_table_users)
       expect(test_klass.cells_claims_attributes).to eq(email: { type: :string }, username: { type: :string })
     end
-  end
 
-  describe '.claim_and_commit' do
-    let(:claim_service) { instance_double(Gitlab::TopologyServiceClient::ClaimService) }
+    it 'uses default subject_key when not provided' do
+      test_klass.cells_claims_metadata(subject_type: :user, source_type: :rails_table_users)
 
-    before do
-      allow(Gitlab::TopologyServiceClient::ClaimService).to receive(:new).and_return(claim_service)
+      expect(test_klass.cells_claims_subject_key).to eq(:organization_id)
     end
 
-    it 'yields directly when changes are empty' do
-      expect(claim_service).not_to receive(:claim_and_commit)
-
-      yielded = false
-      described_class.claim_and_commit({}) { yielded = true }
+    it 'allows custom subject_key' do
+      test_klass.cells_claims_metadata(
+        subject_type: :user,
+        subject_key: :custom_id,
+        source_type: :rails_table_users
+      )
 
-      expect(yielded).to be true
+      expect(test_klass.cells_claims_subject_key).to eq(:custom_id)
     end
 
-    it 'calls claim service when changes exist' do
-      changes = { ['users', 1] => { 'email' => %w[old new] } }
-      expect(claim_service).to receive(:claim_and_commit).with(changes)
+    it 'derives source_type from table_name when not provided' do
+      allow(test_klass).to receive(:table_name).and_return('users')
+      test_klass.cells_claims_metadata(subject_type: :user)
 
-      described_class.claim_and_commit(changes)
+      expect(test_klass.cells_claims_source_type).to eq(
+        Gitlab::Cells::TopologyService::Claims::V1::Source::Type::RAILS_TABLE_USERS
+      )
     end
   end
 
-  describe '.rollback' do
-    let(:claim_service) { instance_double(Gitlab::TopologyServiceClient::ClaimService) }
-
-    it 'delegates to claim service' do
-      allow(Gitlab::TopologyServiceClient::ClaimService).to receive(:new).and_return(claim_service)
-      expect(claim_service).to receive(:rollback)
+  describe 'callbacks' do
+    let(:instance) { build(:user) }
+    let(:transaction_record) { instance_double(Cells::TransactionRecord) }
 
-      described_class.rollback
+    before do
+      allow(Cells::TransactionRecord).to receive(:current_transaction)
+        .with(instance.connection)
+        .and_return(transaction_record)
     end
-  end
 
-  describe 'callbacks' do
-    let(:instance) { build(:user) }
+    describe '#cells_claims_save_changes' do
+      context 'when transaction record exists' do
+        context 'when creating a new record' do
+          it 'creates claims for all configured attributes' do
+            instance.email = 'new@example.com'
+            instance.username = 'newuser'
+
+            expect(transaction_record).to receive(:create_record).once.with(
+              {
+                bucket: { type: 0, value: 'newuser' },
+                source: { type: 1, rails_primary_key_id: be_a(Integer) },
+                subject: { type: 4, id: instance.organization_id }
+              }
+            )
+            expect(transaction_record).to receive(:create_record).once.with(
+              {
+                bucket: { type: 0, value: 'new@example.com' },
+                source: { type: 2, rails_primary_key_id: be_a(Integer) },
+                subject: { type: 4, id: be_a(Integer) }
+              }
+            )
+
+            instance.save!
+          end
+        end
+
+        context 'when updating an existing record' do
+          let(:instance) { user }
+
+          it 'destroys old claim and creates new claim when attribute changes' do
+            old_username = instance.username
+            new_username = generate(:username)
+
+            expect(transaction_record)
+            .to receive(:destroy_record).with(a_hash_including(bucket: { type: 0, value: old_username }))
+            expect(transaction_record)
+              .to receive(:create_record).with(a_hash_including(bucket: { type: 0, value: new_username }))
+
+            instance.update!(username: new_username)
+          end
+
+          it 'does not process unchanged attributes' do
+            expect(transaction_record).not_to receive(:destroy_record)
+            expect(transaction_record).not_to receive(:create_record)
+
+            instance.save!
+          end
 
-    context 'when cells are enabled' do
-      before do
-        allow(Gitlab.config.cell).to receive(:enabled).and_return(true)
+          it 'only destroys claim when attribute changes to nil' do
+            old_username = instance.username
+
+            expect(transaction_record)
+              .to receive(:destroy_record).with(hash_including(bucket: { type: 0, value: old_username }))
+            expect(transaction_record).not_to receive(:create_record)
+
+            instance.update_attribute(:username, nil)
+          end
+        end
       end
 
-      describe 'after_save callback' do
-        it 'calls inspect_save_changes after saving' do
-          email_inspector = instance_double(Cells::InspectChanges)
-          user_inspector = instance_double(Cells::InspectChanges)
+      context 'when transaction record does not exist' do
+        before do
+          allow(Cells::TransactionRecord).to receive(:current_transaction).and_return(nil)
+        end
 
-          expect(Cells::InspectChanges).to receive(:new).with(an_instance_of(Email)).and_return(email_inspector)
-          expect(Cells::InspectChanges).to receive(:new).with(instance).and_return(user_inspector)
-          expect(email_inspector).to receive(:inspect_changes)
-          expect(user_inspector).to receive(:inspect_changes)
+        it 'does not process claims' do
+          expect(transaction_record).not_to receive(:create_record)
+          expect(transaction_record).not_to receive(:destroy_record)
 
           instance.save!
         end
       end
+    end
+
+    describe '#cells_claims_destroy_changes' do
+      let(:instance) { user }
 
-      describe 'before_destroy callback' do
-        it 'calls inspect_destroy_changes before destroying' do
-          inspector = instance_double(Cells::InspectChanges)
+      context 'when transaction record exists' do
+        it 'destroys claims for all configured attributes' do
+          old_username = instance.username
+          old_email = instance.email
 
-          expect(Cells::InspectChanges).to receive(:new).with(user, destroy: true).and_return(inspector)
-          expect(inspector).to receive(:inspect_changes)
-          user.destroy!
+          expect(transaction_record)
+            .to receive(:destroy_record).with(a_hash_including(bucket: { type: 0, value: old_username }))
+          expect(transaction_record)
+            .to receive(:destroy_record).with(a_hash_including(bucket: { type: 0, value: old_email }))
+
+          instance.destroy!
         end
       end
 
-      context 'when feature flag is disabled' do
+      context 'when transaction record does not exist' do
         before do
-          stub_feature_flags(cells_unique_claims: false)
-          Current.reset
+          allow(Cells::TransactionRecord).to receive(:current_transaction).and_return(nil)
         end
 
-        it 'does not call inspect_save_changes on save' do
-          expect(Cells::InspectChanges).not_to receive(:new)
-          instance.save!
-        end
+        it 'does not process claims' do
+          expect(transaction_record).not_to receive(:destroy_record)
 
-        it 'does not call inspect_destroy_changes on destroy' do
-          expect(Cells::InspectChanges).not_to receive(:new)
-          user.destroy!
+          instance.destroy!
         end
       end
     end
+  end
 
-    context 'when cells are disabled' do
-      before do
-        allow(Gitlab.config.cell).to receive(:enabled).and_return(false)
-      end
+  describe '#cells_claims_default_metadata' do
+    let(:instance) { user }
 
-      it 'does not call inspect_save_changes on save' do
-        expect(Cells::InspectChanges).not_to receive(:new)
-        instance.save!
-      end
+    before do
+      allow(instance.class).to receive_messages(
+        cells_claims_subject_type: :user,
+        cells_claims_subject_key: :organization_id,
+        cells_claims_source_type: :users
+      )
+      allow(instance).to receive(:organization_id).and_return(123)
+    end
 
-      it 'does not call inspect_destroy_changes on destroy' do
-        expect(Cells::InspectChanges).not_to receive(:new)
-        user.destroy!
-      end
+    it 'returns metadata with subject and source information' do
+      metadata = instance.send(:cells_claims_default_metadata)
+
+      expect(metadata).to eq({
+        subject: { type: :user, id: 123 },
+        source: { type: :users, rails_primary_key_id: instance.id }
+      })
     end
   end
 end
-- 
GitLab


From 6ab45c5fce4fd93e826c6041839b2851942867cf Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Wed, 15 Oct 2025 22:50:23 +0800
Subject: [PATCH 24/29] Use a prepended module to store the transaction record

This also addresses a problem that instead of being nil,
when there's no transaction, it would be a NullTransaction.
---
 app/models/cells/transaction_record.rb                 | 10 ++++++----
 .../active_record_connection_adapters_transaction.rb   |  4 ++++
 2 files changed, 10 insertions(+), 4 deletions(-)
 create mode 100644 config/initializers/active_record_connection_adapters_transaction.rb

diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
index bade09dcffe2de..29e40710214f94 100644
--- a/app/models/cells/transaction_record.rb
+++ b/app/models/cells/transaction_record.rb
@@ -2,7 +2,9 @@
 
 module Cells
   class TransactionRecord
-    TRANSACTION_RECORD_KEY = :@cells_current_transaction
+    module TransactionExtension
+      attr_accessor :cells_current_transaction_record
+    end
 
     # TODO: hack to assign to current transaction
     # Figure out a better way to do this
@@ -14,17 +16,17 @@ def self.current_transaction(connection)
       # The Cells::OutstandingLease requires a transaction to be open
       # to ensure that the lease is only created if the transaction
       # within a transaction and not outside of one
-      unless connection.current_transaction # rubocop:disable Style/IfUnlessModifier -- this is more clear
+      if connection.current_transaction.closed? # rubocop:disable Style/IfUnlessModifier -- this is more clear
         raise 'The Cells::TransactionRecord requires transaction to be open'
       end
 
       current_transaction = connection.current_transaction
 
-      instance = current_transaction.instance_variable_get(TRANSACTION_RECORD_KEY)
+      instance = current_transaction.cells_current_transaction_record
       return instance if instance
 
       TransactionRecord.new(connection, current_transaction).tap do |instance|
-        current_transaction.instance_variable_set(TRANSACTION_RECORD_KEY, instance)
+        current_transaction.cells_current_transaction_record = instance
         current_transaction.add_record(instance)
       end
     end
diff --git a/config/initializers/active_record_connection_adapters_transaction.rb b/config/initializers/active_record_connection_adapters_transaction.rb
new file mode 100644
index 00000000000000..eb4f509fa4de96
--- /dev/null
+++ b/config/initializers/active_record_connection_adapters_transaction.rb
@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+ActiveRecord::ConnectionAdapters::Transaction
+  .prepend(Cells::TransactionRecord::TransactionExtension)
-- 
GitLab


From 46d79fbd4e47f8d47e5d00ed70289485c4f526af Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Thu, 16 Oct 2025 00:23:33 +0800
Subject: [PATCH 25/29] Fix left out renaming

---
 app/models/gpg_key_subkey.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/app/models/gpg_key_subkey.rb b/app/models/gpg_key_subkey.rb
index 5643da92483261..ff944898199cdd 100644
--- a/app/models/gpg_key_subkey.rb
+++ b/app/models/gpg_key_subkey.rb
@@ -6,7 +6,7 @@ class GpgKeySubkey < ApplicationRecord
 
   cells_claims_attribute :fingerprint, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED
 
-  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: SOURCE_TYPE::UNSPECIFIED,
+  cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::UNSPECIFIED, source_type: CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     subject_key: :gpg_key_id
 
   sha_attribute :keyid
-- 
GitLab


From 7ac5cb235fdf9cfd417bbda5a76ae94b7041208f Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Thu, 16 Oct 2025 12:14:01 +0200
Subject: [PATCH 26/29] Fixes specs for claimable models

---
 app/models/organizations/organization.rb             |  2 +-
 spec/models/email_spec.rb                            |  3 ++-
 spec/models/gpg_key_spec.rb                          |  3 ++-
 spec/models/gpg_key_subkey_spec.rb                   |  3 ++-
 spec/models/key_spec.rb                              |  3 ++-
 spec/models/organizations/organization_spec.rb       |  3 ++-
 spec/models/route_spec.rb                            |  3 ++-
 spec/models/user_spec.rb                             |  3 ++-
 .../shared_examples/cells/unique_shared_examples.rb  | 12 ++++++++----
 9 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/app/models/organizations/organization.rb b/app/models/organizations/organization.rb
index c43020209a99c4..eb97b2a4d3bc4c 100644
--- a/app/models/organizations/organization.rb
+++ b/app/models/organizations/organization.rb
@@ -10,7 +10,7 @@ class Organization < ApplicationRecord
     include Organizations::Isolatable
     include Cells::Claimable
 
-    cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::UNSPECIFIED # TODO: PATH?
+    cells_claims_attribute :path, type: CLAIMS_BUCKET_TYPE::ORGANIZATION_PATH
 
     cells_claims_metadata subject_type: CLAIMS_SUBJECT_TYPE::ORGANIZATION
 
diff --git a/spec/models/email_spec.rb b/spec/models/email_spec.rb
index 8d58611df77750..5c92590905e84a 100644
--- a/spec/models/email_spec.rb
+++ b/spec/models/email_spec.rb
@@ -6,7 +6,8 @@
   let_it_be(:user) { create(:user) }
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::USER,
+    subject_key: :user_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_EMAILS,
     unique_attributes: [:email]
 
diff --git a/spec/models/gpg_key_spec.rb b/spec/models/gpg_key_spec.rb
index f3994ccc84cdaf..433508097c11b4 100644
--- a/spec/models/gpg_key_spec.rb
+++ b/spec/models/gpg_key_spec.rb
@@ -6,7 +6,8 @@
   subject { build(:gpg_key) }
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::USER,
+    subject_key: :user_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint, :primary_keyid]
 
diff --git a/spec/models/gpg_key_subkey_spec.rb b/spec/models/gpg_key_subkey_spec.rb
index 99e7f81063eb79..069206fbf29b93 100644
--- a/spec/models/gpg_key_subkey_spec.rb
+++ b/spec/models/gpg_key_subkey_spec.rb
@@ -16,7 +16,8 @@
   end
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_key: :gpg_key_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:fingerprint]
 end
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index 2e6b2bc0e95750..96cf58c3c33106 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -6,7 +6,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_key: :organization_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::UNSPECIFIED,
     unique_attributes: [:key, :fingerprint_sha256]
 
diff --git a/spec/models/organizations/organization_spec.rb b/spec/models/organizations/organization_spec.rb
index 4a3111f4438eac..7b343a5eca3f4b 100644
--- a/spec/models/organizations/organization_spec.rb
+++ b/spec/models/organizations/organization_spec.rb
@@ -6,7 +6,8 @@
   let_it_be_with_refind(:organization) { create(:organization) }
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::ORGANIZATION,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::ORGANIZATION,
+    subject_key: :organization_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ORGANIZATIONS,
     unique_attributes: [:path]
 
diff --git a/spec/models/route_spec.rb b/spec/models/route_spec.rb
index 2bff53c8bd91a7..308cde317e6930 100644
--- a/spec/models/route_spec.rb
+++ b/spec/models/route_spec.rb
@@ -7,7 +7,8 @@
   let(:route) { group.route }
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::UNSPECIFIED,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::GROUP,
+    subject_key: :namespace_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_ROUTES,
     unique_attributes: [:path]
 
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index 582edabb365271..928b9cdd1feef2 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -50,7 +50,8 @@
   it_behaves_like 'having unique enum values'
 
   it_behaves_like 'cells claimable model',
-    type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::USER,
+    subject_type: Cells::Claimable::CLAIMS_SUBJECT_TYPE::USER,
+    subject_key: :organization_id,
     source_type: Cells::Claimable::CLAIMS_SOURCE_TYPE::RAILS_TABLE_USERS,
     unique_attributes: [:username]
 
diff --git a/spec/support/shared_examples/cells/unique_shared_examples.rb b/spec/support/shared_examples/cells/unique_shared_examples.rb
index 883d184158ea23..0aa2beed0f260d 100644
--- a/spec/support/shared_examples/cells/unique_shared_examples.rb
+++ b/spec/support/shared_examples/cells/unique_shared_examples.rb
@@ -1,12 +1,16 @@
 # frozen_string_literal: true
 
-RSpec.shared_examples 'cells claimable model' do |type:, source_type:, unique_attributes:|
-  it 'has the expected owner_type' do
-    expect(described_class.cells_type).to eq(type)
+RSpec.shared_examples 'cells claimable model' do |subject_type:, subject_key:, source_type:, unique_attributes:|
+  it 'has the expected subject_type' do
+    expect(described_class.cells_claims_subject_type).to eq(subject_type)
+  end
+
+  it 'has the expected subject_key' do
+    expect(described_class.cells_claims_subject_key).to eq(subject_key)
   end
 
   it 'has the expected source_type' do
-    expect(described_class.cells_source_type).to eq(source_type)
+    expect(described_class.cells_claims_source_type).to eq(source_type)
   end
 
   it 'has the expected unique attributes' do
-- 
GitLab


From 2108931911cf04868f5cc5a7fd8b425c66379649 Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Thu, 16 Oct 2025 15:09:54 +0200
Subject: [PATCH 27/29] Add timeout to TS Claim Service

---
 lib/gitlab/topology_service_client/claim_service.rb        | 6 ++++++
 .../gitlab/topology_service_client/claim_service_spec.rb   | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/lib/gitlab/topology_service_client/claim_service.rb b/lib/gitlab/topology_service_client/claim_service.rb
index b248e63afdeaac..9dbb19078d6903 100644
--- a/lib/gitlab/topology_service_client/claim_service.rb
+++ b/lib/gitlab/topology_service_client/claim_service.rb
@@ -5,8 +5,14 @@ module TopologyServiceClient
     class ClaimService < BaseService
       include Singleton
 
+      CLAIM_SERVICE_TIMEOUT_IN_SECONDS = 0.2
+
       delegate :begin_update, :commit_update, :rollback_update, to: :client
 
+      def initialize
+        super(timeout: CLAIM_SERVICE_TIMEOUT_IN_SECONDS)
+      end
+
       private
 
       def service_class
diff --git a/spec/lib/gitlab/topology_service_client/claim_service_spec.rb b/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
index fd29d06080d4cb..e36adabaa61787 100644
--- a/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
+++ b/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
@@ -47,6 +47,13 @@
     end
   end
 
+  describe 'timeout configuration' do
+    it 'uses CLAIM_SERVICE_TIMEOUT_IN_SECONDS timeout' do
+      timeout_option = service.send(:options)[:timeout]
+      expect(timeout_option).to eq(described_class::CLAIM_SERVICE_TIMEOUT_IN_SECONDS)
+    end
+  end
+
   describe '#service_class' do
     it 'returns the correct gRPC stub class' do
       result = service.send(:service_class)
-- 
GitLab


From 480075ad408df6a82250023f235e3dc1f0da9ebd Mon Sep 17 00:00:00 2001
From: Bojan Marjanovic <bmarjanovic@gitlab.com>
Date: Mon, 20 Oct 2025 13:14:30 +0200
Subject: [PATCH 28/29] Replace timeout with deadline for transactions

---
 app/models/cells/outstanding_lease.rb         |  16 +-
 app/models/cells/transaction_record.rb        |  10 +-
 .../topology_service_client/claim_service.rb  |   6 -
 .../claim_service_spec.rb                     |   7 -
 spec/models/cells/outstanding_lease_spec.rb   | 127 ++++++++++++
 spec/models/cells/transaction_record_spec.rb  | 184 ++++++++++++++++++
 6 files changed, 324 insertions(+), 26 deletions(-)
 create mode 100644 spec/models/cells/outstanding_lease_spec.rb
 create mode 100644 spec/models/cells/transaction_record_spec.rb

diff --git a/app/models/cells/outstanding_lease.rb b/app/models/cells/outstanding_lease.rb
index 2bf3764b2d4f65..d9f3701668f53a 100644
--- a/app/models/cells/outstanding_lease.rb
+++ b/app/models/cells/outstanding_lease.rb
@@ -8,7 +8,7 @@ def self.claim_service
       ::Gitlab::TopologyServiceClient::ClaimService.instance # rubocop:disable CodeReuse/ServiceClass -- this is a gRPC client
     end
 
-    def self.create_from_request!(create_records:, destroy_records:)
+    def self.create_from_request!(create_records:, destroy_records:, deadline: nil)
       req = Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateRequest.new(
         create_records: create_records,
         destroy_records: destroy_records,
@@ -16,27 +16,25 @@ def self.create_from_request!(create_records:, destroy_records:)
       )
 
       pp("send_begin_update req #{req}")
-      res = claim_service.begin_update(req)
+      res = claim_service.begin_update(req, deadline: deadline)
 
       pp("send_begin_update res #{res.lease_uuid}")
       create!(uuid: res.lease_uuid.value)
     end
 
-    def send_commit_update!
+    def send_commit_update!(deadline: nil)
       req = Gitlab::Cells::TopologyService::Claims::V1::CommitUpdateRequest.new(
         lease_uuid: Gitlab::Cells::TopologyService::Types::V1::UUID.new(value: uuid),
         cell_id: self.class.claim_service.cell_id
       )
 
       pp("send_commit_update #{object_id} #{req}")
-      self.class.claim_service.commit_update(
-        req
-      )
+      self.class.claim_service.commit_update(req, deadline: deadline)
 
       pp("send_commit_update done #{object_id} #{req}")
     end
 
-    def send_rollback_update!
+    def send_rollback_update!(deadline: nil)
       return unless uuid
 
       req = Gitlab::Cells::TopologyService::Claims::V1::RollbackUpdateRequest.new(
@@ -46,9 +44,7 @@ def send_rollback_update!
 
       pp("send_rollback_update #{object_id} #{req}")
 
-      self.class.claim_service.rollback_update(
-        req
-      )
+      self.class.claim_service.rollback_update(req, deadline: deadline)
 
       pp("send_rollback_update done #{object_id} #{req}")
     end
diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
index 29e40710214f94..a96696d7312385 100644
--- a/app/models/cells/transaction_record.rb
+++ b/app/models/cells/transaction_record.rb
@@ -6,6 +6,8 @@ module TransactionExtension
       attr_accessor :cells_current_transaction_record
     end
 
+    TIMEOUT_IN_SECONDS = 0.2
+
     # TODO: hack to assign to current transaction
     # Figure out a better way to do this
     # Current? but what if nested transactions?
@@ -38,6 +40,7 @@ def initialize(connection, transaction)
       @destroy_records = []
       @outstanding_lease = nil
       @done = false
+      @deadline = GRPC::Core::TimeConsts.from_relative_time(TIMEOUT_IN_SECONDS)
     end
 
     def create_record(metadata)
@@ -67,7 +70,8 @@ def before_committed!
 
       @outstanding_lease = Cells::OutstandingLease.create_from_request!(
         create_records: @create_records,
-        destroy_records: @destroy_records
+        destroy_records: @destroy_records,
+        deadline: @deadline
       )
     end
 
@@ -78,7 +82,7 @@ def rolledback!(force_restore_state: false, should_run_callbacks: true) # ruboco
       # since the transaction might be rolledback prematurely
       return unless @outstanding_lease
 
-      @outstanding_lease.send_rollback_update!
+      @outstanding_lease.send_rollback_update!(deadline: @deadline)
       @outstanding_lease.destroy! # the lease is no longer needed
       @done = true
     end
@@ -87,7 +91,7 @@ def committed!(should_run_callbacks: true) # rubocop:disable Lint/UnusedMethodAr
       raise 'Already done' if @done
       raise 'No lease created' unless @outstanding_lease
 
-      @outstanding_lease.send_commit_update!
+      @outstanding_lease.send_commit_update!(deadline: @deadline)
       @outstanding_lease.destroy! # the lease is no longer needed
       @done = true
     end
diff --git a/lib/gitlab/topology_service_client/claim_service.rb b/lib/gitlab/topology_service_client/claim_service.rb
index 9dbb19078d6903..b248e63afdeaac 100644
--- a/lib/gitlab/topology_service_client/claim_service.rb
+++ b/lib/gitlab/topology_service_client/claim_service.rb
@@ -5,14 +5,8 @@ module TopologyServiceClient
     class ClaimService < BaseService
       include Singleton
 
-      CLAIM_SERVICE_TIMEOUT_IN_SECONDS = 0.2
-
       delegate :begin_update, :commit_update, :rollback_update, to: :client
 
-      def initialize
-        super(timeout: CLAIM_SERVICE_TIMEOUT_IN_SECONDS)
-      end
-
       private
 
       def service_class
diff --git a/spec/lib/gitlab/topology_service_client/claim_service_spec.rb b/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
index e36adabaa61787..fd29d06080d4cb 100644
--- a/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
+++ b/spec/lib/gitlab/topology_service_client/claim_service_spec.rb
@@ -47,13 +47,6 @@
     end
   end
 
-  describe 'timeout configuration' do
-    it 'uses CLAIM_SERVICE_TIMEOUT_IN_SECONDS timeout' do
-      timeout_option = service.send(:options)[:timeout]
-      expect(timeout_option).to eq(described_class::CLAIM_SERVICE_TIMEOUT_IN_SECONDS)
-    end
-  end
-
   describe '#service_class' do
     it 'returns the correct gRPC stub class' do
       result = service.send(:service_class)
diff --git a/spec/models/cells/outstanding_lease_spec.rb b/spec/models/cells/outstanding_lease_spec.rb
new file mode 100644
index 00000000000000..ac7b139a36596b
--- /dev/null
+++ b/spec/models/cells/outstanding_lease_spec.rb
@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Cells::OutstandingLease, feature_category: :cell do
+  let(:cell_id) { 1 }
+  let(:uuid_value) { SecureRandom.uuid }
+  let(:mock_service) { instance_double(::Gitlab::TopologyServiceClient::ClaimService) }
+
+  before do
+    allow(described_class).to receive(:claim_service).and_return(mock_service)
+    allow(mock_service).to receive(:cell_id).and_return(cell_id)
+  end
+
+  describe '.create_from_request!' do
+    let(:create_records) { [{}] }
+    let(:destroy_records) { [{}] }
+
+    let(:mock_response) do
+      Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateResponse.new(
+        lease_uuid: Gitlab::Cells::TopologyService::Types::V1::UUID.new(value: uuid_value)
+      )
+    end
+
+    it 'calls begin_update and creates an OutstandingLease' do
+      expect(mock_service).to receive(:begin_update).with(
+        an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateRequest),
+        deadline: nil
+      ).and_return(mock_response)
+
+      lease = described_class.create_from_request!(create_records: create_records, destroy_records: destroy_records)
+
+      expect(lease).to be_persisted
+      expect(lease.uuid).to eq(uuid_value)
+    end
+
+    context 'with deadline' do
+      let(:deadline) { GRPC::Core::TimeConsts.from_relative_time(5.0) }
+
+      it 'calls begin_update with deadline and creates an OutstandingLease' do
+        expect(mock_service).to receive(:begin_update).with(
+          an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::BeginUpdateRequest),
+          deadline: deadline
+        ).and_return(mock_response)
+
+        lease = described_class.create_from_request!(
+          create_records: create_records,
+          destroy_records: destroy_records,
+          deadline: deadline
+        )
+
+        expect(lease).to be_persisted
+        expect(lease.uuid).to eq(uuid_value)
+      end
+    end
+  end
+
+  describe '#send_commit_update!' do
+    let(:lease) { described_class.create!(uuid: uuid_value) }
+
+    it 'sends a commit_update request with the correct parameters' do
+      expect(mock_service).to receive(:commit_update).with(
+        an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::CommitUpdateRequest),
+        deadline: nil
+      )
+
+      lease.send_commit_update!
+    end
+
+    context 'with deadline' do
+      let(:deadline) { GRPC::Core::TimeConsts.from_relative_time(5.0) }
+
+      it 'sends a commit_update request with the correct parameters' do
+        expect(mock_service).to receive(:commit_update).with(
+          an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::CommitUpdateRequest),
+          deadline: deadline
+        )
+
+        lease.send_commit_update!(deadline: deadline)
+      end
+    end
+  end
+
+  describe '#send_rollback_update!' do
+    context 'when uuid is present' do
+      let(:lease) { described_class.create!(uuid: uuid_value) }
+
+      it 'sends a rollback_update request with the correct parameters' do
+        expect(mock_service).to receive(:rollback_update).with(
+          an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::RollbackUpdateRequest),
+          deadline: nil
+        )
+
+        lease.send_rollback_update!
+      end
+
+      context 'with deadline' do
+        let(:deadline) { GRPC::Core::TimeConsts.from_relative_time(5.0) }
+
+        it 'sends a rollback_update request with the correct parameters' do
+          expect(mock_service).to receive(:rollback_update).with(
+            an_instance_of(Gitlab::Cells::TopologyService::Claims::V1::RollbackUpdateRequest),
+            deadline: deadline
+          )
+
+          lease.send_rollback_update!(deadline: deadline)
+        end
+      end
+    end
+
+    context 'when uuid is nil' do
+      let(:lease) { described_class.new(uuid: nil) }
+
+      it 'does not send a rollback_update request' do
+        expect(mock_service).not_to receive(:rollback_update)
+
+        lease.send_rollback_update!
+      end
+    end
+  end
+
+  describe '.claim_service' do
+    it 'returns the ClaimService instance' do
+      expect(described_class.claim_service).to eq(mock_service)
+    end
+  end
+end
diff --git a/spec/models/cells/transaction_record_spec.rb b/spec/models/cells/transaction_record_spec.rb
new file mode 100644
index 00000000000000..aca2fab02d5946
--- /dev/null
+++ b/spec/models/cells/transaction_record_spec.rb
@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Cells::TransactionRecord, feature_category: :cell do
+  let(:connection) { ApplicationRecord.connection }
+  let(:transaction) do
+    instance_double(ActiveRecord::ConnectionAdapters::NullTransaction, closed?: false, add_record: nil)
+  end
+
+  let(:current_transaction) { transaction }
+
+  before do
+    allow(connection).to receive(:current_transaction).and_return(current_transaction)
+    allow(GRPC::Core::TimeConsts).to receive(:from_relative_time).and_return("fake-deadline")
+  end
+
+  describe ".current_transaction" do
+    subject(:transaction_record) { described_class.current_transaction(connection) }
+
+    before do
+      transaction.extend(Cells::TransactionRecord::TransactionExtension)
+
+      allow(Current).to receive(:cells_claims_leases?).and_return(true)
+    end
+
+    context "when Current.cells_claims_leases? is false" do
+      before do
+        allow(Current).to receive(:cells_claims_leases?).and_return(false)
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context "when transaction is closed" do
+      before do
+        allow(transaction).to receive(:closed?).and_return(true)
+      end
+
+      it "raises an error" do
+        expect { transaction_record }.to raise_error("The Cells::TransactionRecord requires transaction to be open")
+      end
+    end
+
+    context "when transaction already has a TransactionRecord" do
+      let(:existing_record) { instance_double(described_class) }
+
+      before do
+        transaction.cells_current_transaction_record = existing_record
+      end
+
+      it { is_expected.to eq(existing_record) }
+    end
+
+    context "when transaction does not have a TransactionRecord" do
+      it "creates and attaches a new TransactionRecord" do
+        new_record = transaction_record
+        expect(new_record).to be_a(described_class)
+        expect(transaction.cells_current_transaction_record).to eq(new_record)
+      end
+    end
+  end
+
+  describe "#create_record and #destroy_record" do
+    let(:record) { described_class.new(connection, transaction) }
+
+    it "stores create metadata" do
+      expect { record.create_record("meta1") }.to change {
+        record.instance_variable_get(:@create_records)
+      }.to include("meta1")
+    end
+
+    it "stores destroy metadata" do
+      expect { record.destroy_record("meta2") }.to change {
+        record.instance_variable_get(:@destroy_records)
+      }.to include("meta2")
+    end
+
+    context "when after lease is created" do
+      before do
+        record.instance_variable_set(:@outstanding_lease, instance_double(Cells::OutstandingLease))
+      end
+
+      it "raises if create_record is called" do
+        expect { record.create_record("meta") }.to raise_error("Lease already created")
+      end
+
+      it "raises if destroy_record is called" do
+        expect { record.destroy_record("meta") }.to raise_error("Lease already created")
+      end
+    end
+  end
+
+  describe "transaction lifecycle callbacks" do
+    let(:record) { described_class.new(connection, transaction) }
+    let(:lease) do
+      instance_double(Cells::OutstandingLease, send_commit_update!: nil, send_rollback_update!: nil, destroy!: nil)
+    end
+
+    before do
+      allow(Cells::OutstandingLease).to receive(:connection).and_return(connection)
+    end
+
+    describe "#before_committed!" do
+      it "creates a lease" do
+        expect(Cells::OutstandingLease).to receive(:create_from_request!).with(
+          create_records: [],
+          destroy_records: [],
+          deadline: "fake-deadline"
+        ).and_return(lease)
+
+        record.before_committed!
+        expect(record.instance_variable_get(:@outstanding_lease)).to eq(lease)
+      end
+
+      it "raises if already done" do
+        record.instance_variable_set(:@done, true)
+        expect { record.before_committed! }.to raise_error("Already done")
+      end
+
+      it "raises if lease already created" do
+        record.instance_variable_set(:@outstanding_lease, lease)
+        expect { record.before_committed! }.to raise_error("Already created lease")
+      end
+
+      it "raises if connection mismatch" do
+        allow(Cells::OutstandingLease)
+          .to receive(:connection)
+          .and_return(instance_double(Gitlab::Database::LoadBalancing::ConnectionProxy))
+
+        expect { record.before_committed! }.to raise_error("Cross-database lease")
+      end
+    end
+
+    describe "#committed!" do
+      before do
+        record.instance_variable_set(:@outstanding_lease, lease)
+      end
+
+      it "sends commit update and destroys lease" do
+        expect(lease).to receive(:send_commit_update!).with(deadline: "fake-deadline")
+        expect(lease).to receive(:destroy!)
+        record.committed!
+        expect(record.instance_variable_get(:@done)).to be true
+      end
+
+      it "raises if already done" do
+        record.instance_variable_set(:@done, true)
+        expect { record.committed! }.to raise_error("Already done")
+      end
+
+      it "raises if no lease created" do
+        record.instance_variable_set(:@outstanding_lease, nil)
+        expect { record.committed! }.to raise_error("No lease created")
+      end
+    end
+
+    describe "#rolledback!" do
+      context "with lease" do
+        before do
+          record.instance_variable_set(:@outstanding_lease, lease)
+        end
+
+        it "sends rollback update and destroys lease" do
+          expect(lease).to receive(:send_rollback_update!).with(deadline: "fake-deadline")
+          expect(lease).to receive(:destroy!)
+          record.rolledback!
+          expect(record.instance_variable_get(:@done)).to be true
+        end
+      end
+
+      context "without lease" do
+        it "does not raise" do
+          expect { record.rolledback! }.not_to raise_error
+        end
+      end
+
+      it "raises if already done" do
+        record.instance_variable_set(:@done, true)
+        expect { record.rolledback! }.to raise_error("Already done")
+      end
+    end
+  end
+end
-- 
GitLab


From 696e6644797a7b32fddd5c41410fe8cbb5f9d6a4 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Wed, 22 Oct 2025 01:48:04 +0800
Subject: [PATCH 29/29] Document some contracts we need

---
 app/models/cells/transaction_record.rb        | 29 +++++++++++++++++--
 ..._record_connection_adapters_transaction.rb |  4 +++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/app/models/cells/transaction_record.rb b/app/models/cells/transaction_record.rb
index a96696d7312385..ce052761c3965a 100644
--- a/app/models/cells/transaction_record.rb
+++ b/app/models/cells/transaction_record.rb
@@ -1,6 +1,27 @@
 # frozen_string_literal: true
 
 module Cells
+  # This class is a fake ActiveRecord::Base which implements the interface
+  # partially so that it can be used with
+  # ActiveRecord::ConnectionAdapters::Transaction#add_record
+  # https://github.com/rails/rails/blob/v7.1.5.2/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb#L147-L155
+  # From there, several methods are required to be implemented:
+  # * #before_committed!
+  #   * Normally implemented by ActiveRecord::Transactions
+  # * #committed!
+  #   * Normally implemented by ActiveRecord::Transactions
+  # * #rolledback!
+  #   * Normally implemented by ActiveRecord::Transactions
+  # * #trigger_transactional_callbacks?
+  #   * Normally implemented by ActiveRecord::Transactions
+  #   * Overridden by Cells::TransactionRecord to avoid more interfaces needed
+  # * #destroyed?
+  #   * Normally implemented by ActiveRecord::Persistence but we ignore it
+  # * #_new_record_before_last_commit
+  #   * Normally implemented by ActiveRecord::Transactions
+  # We ignore some methods because they won't be called under specific
+  # conditions, specifically when the following both return true:
+  # * #trigger_transactional_callbacks?
   class TransactionRecord
     module TransactionExtension
       attr_accessor :cells_current_transaction_record
@@ -29,6 +50,8 @@ def self.current_transaction(connection)
 
       TransactionRecord.new(connection, current_transaction).tap do |instance|
         current_transaction.cells_current_transaction_record = instance
+        # https://api.rubyonrails.org/v7.1.5.2/classes/ActiveRecord/ConnectionAdapters/DatabaseStatements.html#method-i-add_transaction_record
+        # https://github.com/rails/rails/blob/v7.1.5.2/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb#L147-L155
         current_transaction.add_record(instance)
       end
     end
@@ -57,8 +80,10 @@ def destroy_record(metadata)
       pp("destroy_record #{object_id} #{metadata}")
     end
 
-    # Callbacks required by the ActiveRecord::ConnectionAdapters::TransactionState
-
+    # Always trigger callbacks. See:
+    # ActiveRecord::ConnectionAdapters::Transaction#
+    #   prepare_instances_to_run_callbacks_on
+    # https://github.com/rails/rails/blob/v7.1.5.2/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb#L269
     def trigger_transactional_callbacks?
       true
     end
diff --git a/config/initializers/active_record_connection_adapters_transaction.rb b/config/initializers/active_record_connection_adapters_transaction.rb
index eb4f509fa4de96..049a5103e357a3 100644
--- a/config/initializers/active_record_connection_adapters_transaction.rb
+++ b/config/initializers/active_record_connection_adapters_transaction.rb
@@ -1,4 +1,8 @@
 # frozen_string_literal: true
 
+if Rails.gem_version >= Gem::Version.new('7.2')
+  raise 'Given that we are relying on a few Rails internal in `app/models/cells/transaction_record.rb`, we should verify if the contract still holds. If it does, please bump the version here.'
+end
+
 ActiveRecord::ConnectionAdapters::Transaction
   .prepend(Cells::TransactionRecord::TransactionExtension)
-- 
GitLab