YML based loose foreign key definition

This change introduces YML based loose foreign key definition and docs
update.

app/models/chat_name.rb

 # frozen_string_literal: true
 
 class ChatName < ApplicationRecord
-  include LooseForeignKey
-
   LAST_USED_AT_INTERVAL = 1.hour
 
   belongs_to :integration, foreign_key: :service_id
@@ -16,8 +14,6 @@ class ChatName < ApplicationRecord
   validates :user_id, uniqueness: { scope: [:service_id] }
   validates :chat_id, uniqueness: { scope: [:service_id, :team_id] }
 
-  loose_foreign_key :ci_pipeline_chat_data, :chat_name_id, on_delete: :async_delete
-
   # Updates the "last_used_timestamp" but only if it wasn't already updated
   # recently.
   #

app/models/ci/runner.rb

@@ -12,7 +12,6 @@ module Ci
     include Gitlab::Utils::StrongMemoize
     include TaggableQueries
     include Presentable
-    include LooseForeignKey
 
     add_authentication_token_field :token, encrypted: :optional
 
@@ -180,8 +179,6 @@ module Ci
 
     validates :config, json_schema: { filename: 'ci_runner_config' }
 
-    loose_foreign_key :clusters_applications_runners, :runner_id, on_delete: :async_nullify
-
     # Searches for runners matching the given query.
     #
     # This method uses ILIKE on PostgreSQL for the description field and performs a full match on tokens.

app/models/concerns/loose_foreign_key.rb

-# frozen_string_literal: true
-
-module LooseForeignKey
-  extend ActiveSupport::Concern
-
-  # This concern adds loose foreign key support to ActiveRecord models.
-  # Loose foreign keys allow delayed processing of associated database records
-  # with similar guarantees than a database foreign key.
-  #
-  # Prerequisites:
-  #
-  # To start using the concern, you'll need to install a database trigger to the parent
-  # table in a standard DB migration (not post-migration).
-  #
-  # > track_record_deletions(:projects)
-  #
-  # Usage:
-  #
-  # > class Ci::Build < ApplicationRecord
-  # >
-  # >   loose_foreign_key :security_scans, :build_id, on_delete: :async_delete
-  # >
-  # >   # associations can be still defined, the dependent options is no longer necessary:
-  # >   has_many :security_scans, class_name: 'Security::Scan'
-  # >
-  # > end
-  #
-  # Options for on_delete:
-  #
-  # - :async_delete - deletes the children rows via an asynchronous process.
-  # - :async_nullify - sets the foreign key column to null via an asynchronous process.
-  #
-  # How it works:
-  #
-  # When adding loose foreign key support to the table, a DELETE trigger is installed
-  # which tracks the record deletions (stores primary key value of the deleted row) in
-  # a database table.
-  #
-  # These deletion records are processed asynchronously and records are cleaned up
-  # according to the loose foreign key definitions described in the model.
-  #
-  # The cleanup happens in batches, which reduces the likelyhood of statement timeouts.
-  #
-  # When all associations related to the deleted record are cleaned up, the record itself
-  # is deleted.
-  included do
-    class_attribute :loose_foreign_key_definitions, default: []
-  end
-
-  class_methods do
-    def loose_foreign_key(to_table, column, options)
-      symbolized_options = options.symbolize_keys
-
-      unless base_class?
-        raise <<~MSG
-        loose_foreign_key can be only used on base classes, inherited classes are not supported.
-        Please define the loose_foreign_key on the #{base_class.name} class.
-        MSG
-      end
-
-      on_delete_options = %i[async_delete async_nullify]
-
-      unless on_delete_options.include?(symbolized_options[:on_delete]&.to_sym)
-        raise "Invalid on_delete option given: #{symbolized_options[:on_delete]}. Valid options: #{on_delete_options.join(', ')}"
-      end
-
-      definition = ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
-        table_name.to_s,
-        to_table.to_s,
-        {
-          column: column.to_s,
-          on_delete: symbolized_options[:on_delete].to_sym
-        }
-      )
-
-      self.loose_foreign_key_definitions += [definition]
-    end
-  end
-end

app/services/loose_foreign_keys/batch_cleaner_service.rb

@@ -2,11 +2,11 @@
 
 module LooseForeignKeys
   class BatchCleanerService
-    def initialize(parent_klass:, deleted_parent_records:, modification_tracker: LooseForeignKeys::ModificationTracker.new, models_by_table_name:)
-      @parent_klass = parent_klass
+    def initialize(parent_table:, loose_foreign_key_definitions:, deleted_parent_records:, modification_tracker: LooseForeignKeys::ModificationTracker.new)
+      @parent_table = parent_table
+      @loose_foreign_key_definitions = loose_foreign_key_definitions
       @deleted_parent_records = deleted_parent_records
       @modification_tracker = modification_tracker
-      @models_by_table_name = models_by_table_name
       @deleted_records_counter = Gitlab::Metrics.counter(
         :loose_foreign_key_processed_deleted_records,
         'The number of processed loose foreign key deleted records'
@@ -14,11 +14,11 @@ module LooseForeignKeys
     end
 
     def execute
-      parent_klass.loose_foreign_key_definitions.each do |foreign_key_definition|
-        run_cleaner_service(foreign_key_definition, with_skip_locked: true)
+      loose_foreign_key_definitions.each do |loose_foreign_key_definition|
+        run_cleaner_service(loose_foreign_key_definition, with_skip_locked: true)
         break if modification_tracker.over_limit?
 
-        run_cleaner_service(foreign_key_definition, with_skip_locked: false)
+        run_cleaner_service(loose_foreign_key_definition, with_skip_locked: false)
         break if modification_tracker.over_limit?
       end
 
@@ -27,12 +27,12 @@ module LooseForeignKeys
       # At this point, all associations are cleaned up, we can update the status of the parent records
       update_count = LooseForeignKeys::DeletedRecord.mark_records_processed(deleted_parent_records)
 
-      deleted_records_counter.increment({ table: parent_klass.table_name, db_config_name: LooseForeignKeys::DeletedRecord.connection.pool.db_config.name }, update_count)
+      deleted_records_counter.increment({ table: parent_table, db_config_name: LooseForeignKeys::DeletedRecord.connection.pool.db_config.name }, update_count)
     end
 
     private
 
-    attr_reader :parent_klass, :deleted_parent_records, :modification_tracker, :models_by_table_name, :deleted_records_counter
+    attr_reader :parent_table, :loose_foreign_key_definitions, :deleted_parent_records, :modification_tracker, :deleted_records_counter
 
     def record_result(cleaner, result)
       if cleaner.async_delete?
@@ -42,10 +42,12 @@ module LooseForeignKeys
       end
     end
 
-    def run_cleaner_service(foreign_key_definition, with_skip_locked:)
+    def run_cleaner_service(loose_foreign_key_definition, with_skip_locked:)
+      base_models_for_gitlab_schema = Gitlab::Database.schemas_to_base_models.fetch(loose_foreign_key_definition.options[:gitlab_schema])
+      base_models_for_gitlab_schema.each do |base_model|
         cleaner = CleanerService.new(
-        model: models_by_table_name.fetch(foreign_key_definition.to_table),
-        foreign_key_definition: foreign_key_definition,
+          loose_foreign_key_definition: loose_foreign_key_definition,
+          connection: base_model.connection,
           deleted_parent_records: deleted_parent_records,
           with_skip_locked: with_skip_locked
         )
@@ -58,4 +60,5 @@ module LooseForeignKeys
         end
       end
     end
+  end
 end

app/services/loose_foreign_keys/cleaner_service.rb

@@ -6,11 +6,9 @@ module LooseForeignKeys
     DELETE_LIMIT = 1000
     UPDATE_LIMIT = 500
 
-    delegate :connection, to: :model
-
-    def initialize(model:, foreign_key_definition:, deleted_parent_records:, with_skip_locked: false)
-      @model = model
-      @foreign_key_definition = foreign_key_definition
+    def initialize(loose_foreign_key_definition:, connection:, deleted_parent_records:, with_skip_locked: false)
+      @loose_foreign_key_definition = loose_foreign_key_definition
+      @connection = connection
       @deleted_parent_records = deleted_parent_records
       @with_skip_locked = with_skip_locked
     end
@@ -18,20 +16,20 @@ module LooseForeignKeys
     def execute
       result = connection.execute(build_query)
 
-      { affected_rows: result.cmd_tuples, table: foreign_key_definition.to_table }
+      { affected_rows: result.cmd_tuples, table: loose_foreign_key_definition.to_table }
     end
 
     def async_delete?
-      foreign_key_definition.on_delete == :async_delete
+      loose_foreign_key_definition.on_delete == :async_delete
     end
 
     def async_nullify?
-      foreign_key_definition.on_delete == :async_nullify
+      loose_foreign_key_definition.on_delete == :async_nullify
     end
 
     private
 
-    attr_reader :model, :foreign_key_definition, :deleted_parent_records, :with_skip_locked
+    attr_reader :loose_foreign_key_definition, :connection, :deleted_parent_records, :with_skip_locked
 
     def build_query
       query = if async_delete?
@@ -39,10 +37,10 @@ module LooseForeignKeys
               elsif async_nullify?
                 update_query
               else
-                raise "Invalid on_delete argument: #{foreign_key_definition.on_delete}"
+                raise "Invalid on_delete argument: #{loose_foreign_key_definition.on_delete}"
               end
 
-      unless query.include?(%{"#{foreign_key_definition.column}" IN (})
+      unless query.include?(%{"#{loose_foreign_key_definition.column}" IN (})
         raise("FATAL: foreign key condition is missing from the generated query: #{query}")
       end
 
@@ -50,15 +48,15 @@ module LooseForeignKeys
     end
 
     def arel_table
-      @arel_table ||= model.arel_table
+      @arel_table ||= Arel::Table.new(loose_foreign_key_definition.to_table)
     end
 
     def primary_keys
-      @primary_keys ||= connection.primary_keys(model.table_name).map { |key| arel_table[key] }
+      @primary_keys ||= connection.primary_keys(loose_foreign_key_definition.to_table).map { |key| arel_table[key] }
     end
 
     def quoted_table_name
-      @quoted_table_name ||= Arel.sql(connection.quote_table_name(model.table_name))
+      @quoted_table_name ||= Arel.sql(connection.quote_table_name(loose_foreign_key_definition.to_table))
     end
 
     def delete_query
@@ -71,7 +69,7 @@ module LooseForeignKeys
     def update_query
       query = Arel::UpdateManager.new
       query.table(quoted_table_name)
-      query.set([[arel_table[foreign_key_definition.column], nil]])
+      query.set([[arel_table[loose_foreign_key_definition.column], nil]])
 
       add_in_query_with_limit(query, UPDATE_LIMIT)
     end
@@ -88,7 +86,7 @@ module LooseForeignKeys
     def in_query_with_limit(limit)
       in_query = Arel::SelectManager.new
       in_query.from(quoted_table_name)
-      in_query.where(arel_table[foreign_key_definition.column].in(deleted_parent_records.map(&:primary_key_value)))
+      in_query.where(arel_table[loose_foreign_key_definition.column].in(deleted_parent_records.map(&:primary_key_value)))
       in_query.projections = primary_keys
       in_query.take(limit)
       in_query.lock(Arel.sql('FOR UPDATE SKIP LOCKED')) if with_skip_locked

app/services/loose_foreign_keys/process_deleted_records_service.rb

@@ -21,13 +21,16 @@ module LooseForeignKeys
 
         break if modification_tracker.over_limit?
 
-        model = find_parent_model!(table)
+        loose_foreign_key_definitions = Gitlab::Database::LooseForeignKeys.definitions_by_table[table]
+
+        next if loose_foreign_key_definitions.empty?
 
         LooseForeignKeys::BatchCleanerService
-          .new(parent_klass: model,
+          .new(
+            parent_table: table,
+            loose_foreign_key_definitions: loose_foreign_key_definitions,
             deleted_parent_records: records,
-               modification_tracker: modification_tracker,
-               models_by_table_name: models_by_table_name)
+            modification_tracker: modification_tracker)
           .execute
 
         break if modification_tracker.over_limit?
@@ -45,30 +48,12 @@ module LooseForeignKeys
       LooseForeignKeys::DeletedRecord.load_batch_for_table(fully_qualified_table_name, BATCH_SIZE)
     end
 
-    def find_parent_model!(table)
-      models_by_table_name.fetch(table)
-    end
-
     def current_schema
       @current_schema = connection.current_schema
     end
 
     def tracked_tables
-      @tracked_tables ||= models_by_table_name
-        .select { |table_name, model| model.respond_to?(:loose_foreign_key_definitions) }
-        .keys
-    end
-
-    def models_by_table_name
-      @models_by_table_name ||= begin
-        all_models
-          .select(&:base_class?)
-          .index_by(&:table_name)
-      end
-    end
-
-    def all_models
-      ApplicationRecord.descendants
+      @tracked_tables ||= Gitlab::Database::LooseForeignKeys.definitions_by_table.keys
     end
   end
 end

doc/development/database/loose_foreign_keys.md

@@ -52,25 +52,40 @@ For this procedure to work, we must register which tables to clean up asynchrono
 
 ## Example migration and configuration
 
-### Configure the model
+### Configure the loose foreign key
 
-First, tell the application that the `projects` table has a new loose foreign key.
-You can do this in the `Project` model:
+Loose foreign keys are defined in a YAML file. The configuration requires the
+following information:
 
-```ruby
-class Project < ApplicationRecord
-  # ...
+- Parent table name (`projects`)
+- Child table name (`ci_pipelines`)
+- The data cleanup method (`async_delete` or `async_nullify`)
 
-  include LooseForeignKey
+The YAML file is located at `lib/gitlab/database/gitlab_loose_foreign_keys.yml`. The file groups
+foreign key definitions by the name of the parent table. The parent table can have multiple loose
+foreign key definitions, therefore we store them as an array.
 
-  loose_foreign_key :ci_pipelines, :project_id, on_delete: :async_delete # or async_nullify
+Example definition:
 
-  # ...
-end
+```yaml
+projects:
+  - to_table: ci_pipelines
+    column: project_id
+    on_delete: async_delete
 ```
 
-This instruction ensures the asynchronous cleanup process knows about the association, and the
-how to do the cleanup. In this case, the associated `ci_pipelines` records are deleted.
+If the `projects` key is already present in the YAML file, then a new entry can be added
+to the array:
+
+```yaml
+projects:
+  - to_table: ci_pipelines
+    column: project_id
+    on_delete: async_delete
+  - to_table: another_table
+    column: project_id
+    on_delete: :async_nullify
+```
 
 ### Track record changes
 
@@ -127,6 +142,19 @@ end
 At this point, the setup phase is concluded. The deleted `projects` records should be automatically
 picked up by the scheduled cleanup worker job.
 
+## Testing
+
+The "`it has loose foreign keys`" shared example can be used to test the presence of the `ON DELETE` trigger and the
+loose foreign key definitions.
+
+Simply add to the model test file:
+
+```ruby
+it_behaves_like 'it has loose foreign keys' do
+  let(:factory_name) { :project }
+end
+```
+
 ## Caveats of loose foreign keys
 
 ### Record creation

lib/gitlab/database.rb

@@ -63,6 +63,15 @@ module Gitlab
       }.compact.with_indifferent_access.freeze
     end
 
+    # This returns a list of base models with connection associated for a given gitlab_schema
+    def self.schemas_to_base_models
+      @schemas_to_base_models ||= {
+        gitlab_main: [self.database_base_models.fetch(:main)],
+        gitlab_ci: [self.database_base_models[:ci] || self.database_base_models.fetch(:main)], # use CI or fallback to main
+        gitlab_shared: self.database_base_models.values # all models
+      }.with_indifferent_access.freeze
+    end
+
     # We configure the database connection pool size automatically based on the
     # configured concurrency. We also add some headroom, to make sure we don't
     # run out of connections when more threads besides the 'user-facing' ones

lib/gitlab/database/gitlab_loose_foreign_keys.yml

+chat_names:
+  - to_table: ci_pipeline_chat_data
+    column: chat_name_id
+    on_delete: async_delete
+ci_runners:
+  - to_table: clusters_applications_runners
+    column: runner_id
+    on_delete: async_nullify

lib/gitlab/database/loose_foreign_keys.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LooseForeignKeys
+      def self.definitions_by_table
+        @definitions_by_table ||= definitions.group_by(&:from_table).with_indifferent_access.freeze
+      end
+
+      def self.definitions
+        @definitions ||= loose_foreign_keys_yaml.flat_map do |parent_table_name, configs|
+          configs.map { |config| build_definition(parent_table_name, config) }
+        end.freeze
+      end
+
+      def self.build_definition(parent_table_name, config)
+        to_table = config.fetch('to_table')
+
+        ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+          parent_table_name,
+          to_table,
+          {
+            column: config.fetch('column'),
+            on_delete: config.fetch('on_delete').to_sym,
+            gitlab_schema: GitlabSchema.table_schema(to_table)
+          }
+        )
+      end
+
+      def self.loose_foreign_keys_yaml
+        @loose_foreign_keys_yaml ||= YAML.load_file(Rails.root.join('lib/gitlab/database/gitlab_loose_foreign_keys.yml'))
+      end
+
+      private_class_method :build_definition
+      private_class_method :loose_foreign_keys_yaml
+    end
+  end
+end

spec/lib/gitlab/database/loose_foreign_keys_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LooseForeignKeys do
+  describe 'verify all definitions' do
+    subject(:definitions) { described_class.definitions }
+
+    it 'all definitions have assigned a known gitlab_schema and on_delete' do
+      is_expected.to all(have_attributes(
+                           options: a_hash_including(
+                             column: be_a(String),
+                             gitlab_schema: be_in(Gitlab::Database.schemas_to_base_models.symbolize_keys.keys),
+                             on_delete: be_in([:async_delete, :async_nullify])
+                           ),
+                           from_table: be_a(String),
+                           to_table: be_a(String)
+                         ))
+    end
+
+    describe 'ensuring database integrity' do
+      def base_models_for(table)
+        parent_table_schema = Gitlab::Database::GitlabSchema.table_schema(table)
+        Gitlab::Database.schemas_to_base_models.fetch(parent_table_schema)
+      end
+
+      it 'all `from_table` tables are present' do
+        definitions.each do |definition|
+          base_models_for(definition.from_table).each do |model|
+            expect(model.connection).to be_table_exist(definition.from_table)
+          end
+        end
+      end
+
+      it 'all `to_table` tables are present' do
+        definitions.each do |definition|
+          base_models_for(definition.to_table).each do |model|
+            expect(model.connection).to be_table_exist(definition.to_table)
+            expect(model.connection).to be_column_exist(definition.to_table, definition.column)
+          end
+        end
+      end
+    end
+  end
+end

spec/models/chat_name_spec.rb

@@ -46,9 +46,5 @@ RSpec.describe ChatName do
 
   it_behaves_like 'it has loose foreign keys' do
     let(:factory_name) { :chat_name }
-
-    before do
-      Ci::PipelineChatData # ensure that the referenced model is loaded
-    end
   end
 end

spec/models/ci/runner_spec.rb

@@ -7,10 +7,6 @@ RSpec.describe Ci::Runner do
 
   it_behaves_like 'it has loose foreign keys' do
     let(:factory_name) { :ci_runner }
-
-    before do
-      Clusters::Applications::Runner # ensure that the referenced model is loaded
-    end
   end
 
   describe 'groups association' do

spec/models/concerns/loose_foreign_key_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe LooseForeignKey do
-  let(:project_klass) do
-    Class.new(ApplicationRecord) do
-      include LooseForeignKey
-
-      self.table_name = 'projects'
-
-      loose_foreign_key :issues, :project_id, on_delete: :async_delete
-      loose_foreign_key 'merge_requests', 'project_id', 'on_delete' => 'async_nullify'
-    end
-  end
-
-  it 'exposes the loose foreign key definitions' do
-    definitions = project_klass.loose_foreign_key_definitions
-
-    tables = definitions.map(&:to_table)
-    expect(tables).to eq(%w[issues merge_requests])
-  end
-
-  it 'casts strings to symbol' do
-    definition = project_klass.loose_foreign_key_definitions.last
-
-    expect(definition.from_table).to eq('projects')
-    expect(definition.to_table).to eq('merge_requests')
-    expect(definition.column).to eq('project_id')
-    expect(definition.on_delete).to eq(:async_nullify)
-  end
-
-  context 'validation' do
-    context 'on_delete validation' do
-      let(:invalid_class) do
-        Class.new(ApplicationRecord) do
-          include LooseForeignKey
-
-          self.table_name = 'projects'
-
-          loose_foreign_key :issues, :project_id, on_delete: :async_delete
-          loose_foreign_key :merge_requests, :project_id, on_delete: :async_nullify
-          loose_foreign_key :merge_requests, :project_id, on_delete: :destroy
-        end
-      end
-
-      it 'raises error when invalid `on_delete` option was given' do
-        expect { invalid_class }.to raise_error /Invalid on_delete option given: destroy/
-      end
-    end
-
-    context 'inheritance validation' do
-      let(:inherited_project_class) do
-        Class.new(Project) do
-          include LooseForeignKey
-
-          loose_foreign_key :issues, :project_id, on_delete: :async_delete
-        end
-      end
-
-      it 'raises error when loose_foreign_key is defined in a child ActiveRecord model' do
-        expect { inherited_project_class }.to raise_error /Please define the loose_foreign_key on the Project class/
-      end
-    end
-  end
-end

spec/services/loose_foreign_keys/batch_cleaner_service_spec.rb

@@ -21,33 +21,34 @@ RSpec.describe LooseForeignKeys::BatchCleanerService do
     migration.track_record_deletions(:_test_loose_fk_parent_table)
   end
 
-  let(:parent_model) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_parent_table'
-
-      include LooseForeignKey
-
-      loose_foreign_key :_test_loose_fk_child_table_1, :parent_id, on_delete: :async_delete
-      loose_foreign_key :_test_loose_fk_child_table_2, :parent_id_with_different_column, on_delete: :async_nullify
-    end
-  end
-
-  let(:child_model_1) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_child_table_1'
-    end
-  end
-
-  let(:child_model_2) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_child_table_2'
-    end
-  end
-
+  let(:loose_foreign_key_definitions) do
+    [
+      ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+        '_test_loose_fk_parent_table',
+        '_test_loose_fk_child_table_1',
+        {
+          column: 'parent_id',
+          on_delete: :async_delete,
+          gitlab_schema: :gitlab_main
+        }
+      ),
+      ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+        '_test_loose_fk_parent_table',
+        '_test_loose_fk_child_table_2',
+        {
+          column: 'parent_id_with_different_column',
+          on_delete: :async_nullify,
+          gitlab_schema: :gitlab_main
+        }
+      )
+    ]
+  end
+
+  let(:loose_fk_parent_table) { table(:_test_loose_fk_parent_table) }
   let(:loose_fk_child_table_1) { table(:_test_loose_fk_child_table_1) }
   let(:loose_fk_child_table_2) { table(:_test_loose_fk_child_table_2) }
-  let(:parent_record_1) { parent_model.create! }
-  let(:other_parent_record) { parent_model.create! }
+  let(:parent_record_1) { loose_fk_parent_table.create! }
+  let(:other_parent_record) { loose_fk_parent_table.create! }
 
   before(:all) do
     create_table_structure
@@ -87,12 +88,10 @@ RSpec.describe LooseForeignKeys::BatchCleanerService do
       expect(loose_fk_child_table_1.count).to eq(4)
       expect(loose_fk_child_table_2.count).to eq(4)
 
-      described_class.new(parent_klass: parent_model,
-                          deleted_parent_records: LooseForeignKeys::DeletedRecord.status_pending.all,
-                          models_by_table_name: {
-                            '_test_loose_fk_child_table_1' => child_model_1,
-                            '_test_loose_fk_child_table_2' => child_model_2
-                          }).execute
+      described_class.new(parent_table: '_test_loose_fk_parent_table',
+                          loose_foreign_key_definitions: loose_foreign_key_definitions,
+                          deleted_parent_records: LooseForeignKeys::DeletedRecord.status_pending.all
+                         ).execute
     end
 
     it 'cleans up the child records' do
@@ -108,7 +107,7 @@ RSpec.describe LooseForeignKeys::BatchCleanerService do
     it 'records the DeletedRecord status updates', :prometheus do
       counter = Gitlab::Metrics.registry.get(:loose_foreign_key_processed_deleted_records)
 
-      expect(counter.get(table: parent_model.table_name, db_config_name: 'main')).to eq(1)
+      expect(counter.get(table: loose_fk_parent_table.table_name, db_config_name: 'main')).to eq(1)
     end
 
     it 'does not delete unrelated records' do

spec/services/loose_foreign_keys/cleaner_service_spec.rb

@@ -17,17 +17,17 @@ RSpec.describe LooseForeignKeys::CleanerService do
       'issues',
       {
         column: 'project_id',
-        on_delete: :async_nullify
+        on_delete: :async_nullify,
+        gitlab_schema: :gitlab_main
       }
     )
   end
 
   subject(:cleaner_service) do
     described_class.new(
-      model: Issue,
-      foreign_key_definition: loose_fk_definition,
-      deleted_parent_records: deleted_records
-    )
+      loose_foreign_key_definition: loose_fk_definition,
+      connection: ApplicationRecord.connection,
+      deleted_parent_records: deleted_records)
   end
 
   context 'when invalid foreign key definition is passed' do
@@ -84,7 +84,8 @@ RSpec.describe LooseForeignKeys::CleanerService do
           'project_authorizations',
           {
             column: 'user_id',
-            on_delete: :async_delete
+            on_delete: :async_delete,
+            gitlab_schema: :gitlab_main
           }
         )
       end
@@ -97,8 +98,8 @@ RSpec.describe LooseForeignKeys::CleanerService do
 
       subject(:cleaner_service) do
         described_class.new(
-          model: ProjectAuthorization,
-          foreign_key_definition: loose_fk_definition,
+          loose_foreign_key_definition: loose_fk_definition,
+          connection: ApplicationRecord.connection,
           deleted_parent_records: deleted_records
         )
       end
@@ -130,8 +131,8 @@ RSpec.describe LooseForeignKeys::CleanerService do
     context 'when with_skip_locked parameter is true' do
       subject(:cleaner_service) do
         described_class.new(
-          model: Issue,
-          foreign_key_definition: loose_fk_definition,
+          loose_foreign_key_definition: loose_fk_definition,
+          connection: ApplicationRecord.connection,
           deleted_parent_records: deleted_records,
           with_skip_locked: true
         )

spec/support/shared_examples/loose_foreign_keys/have_loose_foreign_key.rb

@@ -5,16 +5,9 @@ RSpec.shared_examples 'it has loose foreign keys' do
   let(:table_name) { described_class.table_name }
   let(:connection) { described_class.connection }
 
-  it 'includes the LooseForeignKey module' do
-    expect(described_class.ancestors).to include(LooseForeignKey)
-  end
-
-  it 'responds to #loose_foreign_key_definitions' do
-    expect(described_class).to respond_to(:loose_foreign_key_definitions)
-  end
-
   it 'has at least one loose foreign key definition' do
-    expect(described_class.loose_foreign_key_definitions.size).to be > 0
+    definitions = Gitlab::Database::LooseForeignKeys.definitions_by_table[table_name]
+    expect(definitions.size).to be > 0
   end
 
   it 'has the deletion trigger present' do

spec/workers/loose_foreign_keys/cleanup_worker_spec.rb

@@ -27,43 +27,40 @@ RSpec.describe LooseForeignKeys::CleanupWorker do
     migration.track_record_deletions(:_test_loose_fk_parent_table_2)
   end
 
-  let!(:parent_model_1) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_parent_table_1'
-
-      include LooseForeignKey
-
-      loose_foreign_key :_test_loose_fk_child_table_1_1, :parent_id, on_delete: :async_delete
-      loose_foreign_key :_test_loose_fk_child_table_1_2, :parent_id_with_different_column, on_delete: :async_nullify
-    end
-  end
-
-  let!(:parent_model_2) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_parent_table_2'
-
-      include LooseForeignKey
-
-      loose_foreign_key :_test_loose_fk_child_table_2_1, :parent_id, on_delete: :async_delete
-    end
-  end
-
-  let!(:child_model_1) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_child_table_1_1'
-    end
-  end
-
-  let!(:child_model_2) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_child_table_1_2'
-    end
-  end
-
-  let!(:child_model_3) do
-    Class.new(ApplicationRecord) do
-      self.table_name = '_test_loose_fk_child_table_2_1'
-    end
+  let(:all_loose_foreign_key_definitions) do
+    {
+      '_test_loose_fk_parent_table_1' => [
+        ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+          '_test_loose_fk_parent_table_1',
+          '_test_loose_fk_child_table_1_1',
+          {
+            column: 'parent_id',
+            on_delete: :async_delete,
+            gitlab_schema: :gitlab_main
+          }
+        ),
+        ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+          '_test_loose_fk_parent_table_1',
+          '_test_loose_fk_child_table_1_2',
+          {
+            column: 'parent_id_with_different_column',
+            on_delete: :async_nullify,
+            gitlab_schema: :gitlab_main
+          }
+        )
+      ],
+      '_test_loose_fk_parent_table_2' => [
+        ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(
+          '_test_loose_fk_parent_table_2',
+          '_test_loose_fk_child_table_2_1',
+          {
+            column: 'parent_id',
+            on_delete: :async_delete,
+            gitlab_schema: :gitlab_main
+          }
+        )
+      ]
+    }
   end
 
   let(:loose_fk_parent_table_1) { table(:_test_loose_fk_parent_table_1) }
@@ -87,6 +84,8 @@ RSpec.describe LooseForeignKeys::CleanupWorker do
   end
 
   before do
+    allow(Gitlab::Database::LooseForeignKeys).to receive(:definitions_by_table).and_return(all_loose_foreign_key_definitions)
+
     parent_record_1 = loose_fk_parent_table_1.create!
     loose_fk_child_table_1_1.create!(parent_id: parent_record_1.id)
     loose_fk_child_table_1_2.create!(parent_id_with_different_column: parent_record_1.id)
@@ -98,8 +97,8 @@ RSpec.describe LooseForeignKeys::CleanupWorker do
     parent_record_3 = loose_fk_parent_table_2.create!
     5.times { loose_fk_child_table_2_1.create!(parent_id: parent_record_3.id) }
 
-    parent_model_1.delete_all
-    parent_model_2.delete_all
+    loose_fk_parent_table_1.delete_all
+    loose_fk_parent_table_2.delete_all
   end
 
   it 'cleans up all rows' do