Add latest changes from gitlab-org/gitlab@14-7-stable-eev14.7.0-rc42

14 files changed, 560 insertions, 628 deletions

diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
index 49714cfc4dd..01d61a525e6 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -336,8 +336,8 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
   end
 
   describe '#smoothed_time_efficiency' do
-    let(:migration) { create(:batched_background_migration, interval: 120.seconds) }
-    let(:end_time) { Time.zone.now }
+    let_it_be(:migration) { create(:batched_background_migration, interval: 120.seconds) }
+    let_it_be(:end_time) { Time.zone.now }
 
     around do |example|
       freeze_time do
@@ -345,7 +345,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
       end
     end
 
-    let(:common_attrs) do
+    let_it_be(:common_attrs) do
       {
         status: :succeeded,
         batched_migration: migration,
@@ -364,13 +364,14 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     end
 
     context 'when there are enough jobs' do
-      subject { migration.smoothed_time_efficiency(number_of_jobs: number_of_jobs) }
+      let_it_be(:number_of_jobs) { 10 }
+      let_it_be(:jobs) { create_list(:batched_background_migration_job, number_of_jobs, **common_attrs.merge(batched_migration: migration)) }
 
-      let!(:jobs) { create_list(:batched_background_migration_job, number_of_jobs, **common_attrs.merge(batched_migration: migration)) }
-      let(:number_of_jobs) { 10 }
+      subject { migration.smoothed_time_efficiency(number_of_jobs: number_of_jobs) }
 
       before do
-        expect(migration).to receive_message_chain(:batched_jobs, :successful_in_execution_order, :reverse_order, :limit).with(no_args).with(no_args).with(number_of_jobs).and_return(jobs)
+        expect(migration).to receive_message_chain(:batched_jobs, :successful_in_execution_order, :reverse_order, :limit, :with_preloads)
+                               .and_return(jobs)
       end
 
       def mock_efficiencies(*effs)
@@ -411,6 +412,18 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
         end
       end
     end
+
+    context 'with preloaded batched migration' do
+      it 'avoids N+1' do
+        create_list(:batched_background_migration_job, 11, **common_attrs.merge(started_at: end_time - 10.seconds))
+
+        control = ActiveRecord::QueryRecorder.new do
+          migration.smoothed_time_efficiency(number_of_jobs: 10)
+        end
+
+        expect { migration.smoothed_time_efficiency(number_of_jobs: 11) }.not_to exceed_query_limit(control)
+      end
+    end
   end
 
   describe '#optimize!' do
diff --git a/spec/lib/gitlab/database/background_migration_job_spec.rb b/spec/lib/gitlab/database/background_migration_job_spec.rb
index 42695925a1c..1117c17c84a 100644
--- a/spec/lib/gitlab/database/background_migration_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration_job_spec.rb
@@ -5,6 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::BackgroundMigrationJob do
   it_behaves_like 'having unique enum values'
 
+  it { is_expected.to be_a Gitlab::Database::SharedModel }
+
   describe '.for_migration_execution' do
     let!(:job1) { create(:background_migration_job) }
     let!(:job2) { create(:background_migration_job, arguments: ['hi', 2]) }
diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index 9831510f014..028bdce852e 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -270,8 +270,6 @@ RSpec.describe Gitlab::Database::BatchCount do
     end
 
     it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_DISTINCT_BATCH_SIZE}" do
-      stub_feature_flags(loose_index_scan_for_distinct_values: false)
-
       min_id = model.minimum(:id)
       relation = instance_double(ActiveRecord::Relation)
       allow(model).to receive_message_chain(:select, public_send: relation)
@@ -317,85 +315,13 @@ RSpec.describe Gitlab::Database::BatchCount do
       end
     end
 
-    context 'when the loose_index_scan_for_distinct_values feature flag is off' do
-      it_behaves_like 'when batch fetch query is canceled' do
-        let(:mode) { :distinct }
-        let(:operation) { :count }
-        let(:operation_args) { nil }
-        let(:column) { nil }
-
-        subject { described_class.method(:batch_distinct_count) }
-
-        before do
-          stub_feature_flags(loose_index_scan_for_distinct_values: false)
-        end
-      end
-    end
-
-    context 'when the loose_index_scan_for_distinct_values feature flag is on' do
+    it_behaves_like 'when batch fetch query is canceled' do
       let(:mode) { :distinct }
       let(:operation) { :count }
       let(:operation_args) { nil }
       let(:column) { nil }
 
-      let(:batch_size) { 10_000 }
-
       subject { described_class.method(:batch_distinct_count) }
-
-      before do
-        stub_feature_flags(loose_index_scan_for_distinct_values: true)
-      end
-
-      it 'reduces batch size by half and retry fetch' do
-        too_big_batch_relation_mock = instance_double(ActiveRecord::Relation)
-
-        count_method = double(send: 1)
-
-        allow(too_big_batch_relation_mock).to receive(:send).and_raise(ActiveRecord::QueryCanceled)
-        allow(Gitlab::Database::LooseIndexScanDistinctCount).to receive_message_chain(:new, :build_query).with(from: 0, to: batch_size).and_return(too_big_batch_relation_mock)
-        allow(Gitlab::Database::LooseIndexScanDistinctCount).to receive_message_chain(:new, :build_query).with(from: 0, to: batch_size / 2).and_return(count_method)
-        allow(Gitlab::Database::LooseIndexScanDistinctCount).to receive_message_chain(:new, :build_query).with(from: batch_size / 2, to: batch_size).and_return(count_method)
-
-        subject.call(model, column, batch_size: batch_size, start: 0, finish: batch_size - 1)
-      end
-
-      context 'when all retries fail' do
-        let(:batch_count_query) { 'SELECT COUNT(id) FROM relation WHERE id BETWEEN 0 and 1' }
-
-        before do
-          relation = instance_double(ActiveRecord::Relation)
-          allow(Gitlab::Database::LooseIndexScanDistinctCount).to receive_message_chain(:new, :build_query).and_return(relation)
-          allow(relation).to receive(:send).and_raise(ActiveRecord::QueryCanceled.new('query timed out'))
-          allow(relation).to receive(:to_sql).and_return(batch_count_query)
-        end
-
-        it 'logs failing query' do
-          expect(Gitlab::AppJsonLogger).to receive(:error).with(
-            event: 'batch_count',
-            relation: model.table_name,
-            operation: operation,
-            operation_args: operation_args,
-            start: 0,
-            mode: mode,
-            query: batch_count_query,
-            message: 'Query has been canceled with message: query timed out'
-          )
-          expect(subject.call(model, column, batch_size: batch_size, start: 0)).to eq(-1)
-        end
-      end
-
-      context 'when LooseIndexScanDistinctCount raises error' do
-        let(:column) { :creator_id }
-        let(:error_class) { Gitlab::Database::LooseIndexScanDistinctCount::ColumnConfigurationError }
-
-        it 'rescues ColumnConfigurationError' do
-          allow(Gitlab::Database::LooseIndexScanDistinctCount).to receive(:new).and_raise(error_class.new('error message'))
-
-          expect(Gitlab::AppJsonLogger).to receive(:error).with(a_hash_including(message: 'LooseIndexScanDistinctCount column error: error message'))
-
-          expect(subject.call(Project, column, batch_size: 10_000, start: 0)).to eq(-1)
-        end
-      end
     end
   end
 
diff --git a/spec/lib/gitlab/database/bulk_update_spec.rb b/spec/lib/gitlab/database/bulk_update_spec.rb
index 9a6463c99fa..08b4d50f83b 100644
--- a/spec/lib/gitlab/database/bulk_update_spec.rb
+++ b/spec/lib/gitlab/database/bulk_update_spec.rb
@@ -101,7 +101,7 @@ RSpec.describe Gitlab::Database::BulkUpdate do
     before do
       configuration_hash = ActiveRecord::Base.connection_db_config.configuration_hash
 
-      ActiveRecord::Base.establish_connection(
+      ActiveRecord::Base.establish_connection( # rubocop: disable Database/EstablishConnection
         configuration_hash.merge(prepared_statements: prepared_statements)
       )
     end
diff --git a/spec/lib/gitlab/database/loose_index_scan_distinct_count_spec.rb b/spec/lib/gitlab/database/loose_index_scan_distinct_count_spec.rb
deleted file mode 100644
index e0eac26e4d9..00000000000
--- a/spec/lib/gitlab/database/loose_index_scan_distinct_count_spec.rb
+++ /dev/null
@@ -1,71 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::LooseIndexScanDistinctCount do
-  context 'counting distinct users' do
-    let_it_be(:user) { create(:user) }
-    let_it_be(:other_user) { create(:user) }
-
-    let(:column) { :creator_id }
-
-    before_all do
-      create_list(:project, 3, creator: user)
-      create_list(:project, 1, creator: other_user)
-    end
-
-    subject(:count) { described_class.new(Project, :creator_id).count(from: Project.minimum(:creator_id), to: Project.maximum(:creator_id) + 1) }
-
-    it { is_expected.to eq(2) }
-
-    context 'when STI model is queried' do
-      it 'does not raise error' do
-        expect { described_class.new(Group, :owner_id).count(from: 0, to: 1) }.not_to raise_error
-      end
-    end
-
-    context 'when model with default_scope is queried' do
-      it 'does not raise error' do
-        expect { described_class.new(GroupMember, :id).count(from: 0, to: 1) }.not_to raise_error
-      end
-    end
-
-    context 'when the fully qualified column is given' do
-      let(:column) { 'projects.creator_id' }
-
-      it { is_expected.to eq(2) }
-    end
-
-    context 'when AR attribute is given' do
-      let(:column) { Project.arel_table[:creator_id] }
-
-      it { is_expected.to eq(2) }
-    end
-
-    context 'when invalid value is given for the column' do
-      let(:column) { Class.new }
-
-      it { expect { described_class.new(Group, column) }.to raise_error(Gitlab::Database::LooseIndexScanDistinctCount::ColumnConfigurationError) }
-    end
-
-    context 'when null values are present' do
-      before do
-        create_list(:project, 2).each { |p| p.update_column(:creator_id, nil) }
-      end
-
-      it { is_expected.to eq(2) }
-    end
-  end
-
-  context 'counting STI models' do
-    let!(:groups) { create_list(:group, 3) }
-    let!(:namespaces) { create_list(:namespace, 2) }
-
-    let(:max_id) { Namespace.maximum(:id) + 1 }
-
-    it 'counts groups' do
-      count = described_class.new(Group, :id).count(from: 0, to: max_id)
-      expect(count).to eq(3)
-    end
-  end
-end
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 7f80bed04a4..7e3de32b965 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -1752,116 +1752,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
-  describe '#change_column_type_using_background_migration' do
-    let!(:issue) { create(:issue, :closed, closed_at: Time.zone.now) }
-
-    let(:issue_model) do
-      Class.new(ActiveRecord::Base) do
-        self.table_name = 'issues'
-        include EachBatch
-      end
-    end
-
-    it 'changes the type of a column using a background migration' do
-      expect(model)
-        .to receive(:add_column)
-        .with('issues', 'closed_at_for_type_change', :datetime_with_timezone)
-
-      expect(model)
-        .to receive(:install_rename_triggers)
-        .with('issues', :closed_at, 'closed_at_for_type_change')
-
-      expect(BackgroundMigrationWorker)
-        .to receive(:perform_in)
-        .ordered
-        .with(
-          10.minutes,
-          'CopyColumn',
-          ['issues', :closed_at, 'closed_at_for_type_change', issue.id, issue.id]
-        )
-
-      expect(BackgroundMigrationWorker)
-        .to receive(:perform_in)
-        .ordered
-        .with(
-          1.hour + 10.minutes,
-          'CleanupConcurrentTypeChange',
-          ['issues', :closed_at, 'closed_at_for_type_change']
-        )
-
-      expect(Gitlab::BackgroundMigration)
-        .to receive(:steal)
-        .ordered
-        .with('CopyColumn')
-
-      expect(Gitlab::BackgroundMigration)
-        .to receive(:steal)
-        .ordered
-        .with('CleanupConcurrentTypeChange')
-
-      model.change_column_type_using_background_migration(
-        issue_model.all,
-        :closed_at,
-        :datetime_with_timezone
-      )
-    end
-  end
-
-  describe '#rename_column_using_background_migration' do
-    let!(:issue) { create(:issue, :closed, closed_at: Time.zone.now) }
-
-    it 'renames a column using a background migration' do
-      expect(model)
-        .to receive(:add_column)
-        .with(
-          'issues',
-          :closed_at_timestamp,
-          :datetime_with_timezone,
-          limit: anything,
-          precision: anything,
-          scale: anything
-        )
-
-      expect(model)
-        .to receive(:install_rename_triggers)
-        .with('issues', :closed_at, :closed_at_timestamp)
-
-      expect(BackgroundMigrationWorker)
-        .to receive(:perform_in)
-        .ordered
-        .with(
-          10.minutes,
-          'CopyColumn',
-          ['issues', :closed_at, :closed_at_timestamp, issue.id, issue.id]
-        )
-
-      expect(BackgroundMigrationWorker)
-        .to receive(:perform_in)
-        .ordered
-        .with(
-          1.hour + 10.minutes,
-          'CleanupConcurrentRename',
-          ['issues', :closed_at, :closed_at_timestamp]
-        )
-
-      expect(Gitlab::BackgroundMigration)
-        .to receive(:steal)
-        .ordered
-        .with('CopyColumn')
-
-      expect(Gitlab::BackgroundMigration)
-        .to receive(:steal)
-        .ordered
-        .with('CleanupConcurrentRename')
-
-      model.rename_column_using_background_migration(
-        'issues',
-        :closed_at,
-        :closed_at_timestamp
-      )
-    end
-  end
-
   describe '#convert_to_bigint_column' do
     it 'returns the name of the temporary column used to convert to bigint' do
       expect(model.convert_to_bigint_column(:id)).to eq('id_convert_to_bigint')
@@ -2065,8 +1955,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         t.integer :other_id
         t.timestamps
       end
-
-      allow(model).to receive(:perform_background_migration_inline?).and_return(false)
     end
 
     context 'when the target table does not exist' do
diff --git a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
index 99c7d70724c..0abb76b9f8a 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -7,249 +7,208 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     ActiveRecord::Migration.new.extend(described_class)
   end
 
-  describe '#queue_background_migration_jobs_by_range_at_intervals' do
-    context 'when the model has an ID column' do
-      let!(:id1) { create(:user).id }
-      let!(:id2) { create(:user).id }
-      let!(:id3) { create(:user).id }
-
-      around do |example|
-        freeze_time { example.run }
-      end
-
-      before do
-        User.class_eval do
-          include EachBatch
-        end
-      end
+  shared_examples_for 'helpers that enqueue background migrations' do |worker_class, tracking_database|
+    before do
+      allow(model).to receive(:tracking_database).and_return(tracking_database)
+    end
 
-      it 'returns the final expected delay' do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
+    describe '#queue_background_migration_jobs_by_range_at_intervals' do
+      context 'when the model has an ID column' do
+        let!(:id1) { create(:user).id }
+        let!(:id2) { create(:user).id }
+        let!(:id3) { create(:user).id }
 
-          expect(final_delay.to_f).to eq(20.minutes.to_f)
+        around do |example|
+          freeze_time { example.run }
         end
-      end
-
-      it 'returns zero when nothing gets queued' do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User.none, 'FooJob', 10.minutes)
 
-          expect(final_delay).to eq(0)
+        before do
+          User.class_eval do
+            include EachBatch
+          end
         end
-      end
 
-      context 'with batch_size option' do
-        it 'queues jobs correctly' do
+        it 'returns the final expected delay' do
           Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
+            final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-            expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
+            expect(final_delay.to_f).to eq(20.minutes.to_f)
           end
         end
-      end
 
-      context 'without batch_size option' do
-        it 'queues jobs correctly' do
+        it 'returns zero when nothing gets queued' do
           Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes)
+            final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User.none, 'FooJob', 10.minutes)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            expect(final_delay).to eq(0)
           end
         end
-      end
 
-      context 'with other_job_arguments option' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
+        context 'when the delay_interval is smaller than the minimum' do
+          it 'sets the delay_interval to the minimum value' do
+            Sidekiq::Testing.fake! do
+              final_delay = model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 1.minute, batch_size: 2)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+              expect(worker_class.jobs[0]['at']).to eq(2.minutes.from_now.to_f)
+              expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+              expect(worker_class.jobs[1]['at']).to eq(4.minutes.from_now.to_f)
+
+              expect(final_delay.to_f).to eq(4.minutes.to_f)
+            end
           end
         end
-      end
 
-      context 'with initial_delay option' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2], initial_delay: 10.minutes)
+        context 'with batch_size option' do
+          it 'queues jobs correctly' do
+            Sidekiq::Testing.fake! do
+              model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, batch_size: 2)
 
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
-            expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(20.minutes.from_now.to_f)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+              expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+              expect(worker_class.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
+            end
           end
         end
-      end
-
-      context 'with track_jobs option' do
-        it 'creates a record for each job in the database' do
-          Sidekiq::Testing.fake! do
-            expect do
-              model.queue_background_migration_jobs_by_range_at_intervals(User, '::FooJob', 10.minutes,
-                other_job_arguments: [1, 2], track_jobs: true)
-            end.to change { Gitlab::Database::BackgroundMigrationJob.count }.from(0).to(1)
-
-            expect(BackgroundMigrationWorker.jobs.size).to eq(1)
 
-            tracked_job = Gitlab::Database::BackgroundMigrationJob.first
+        context 'without batch_size option' do
+          it 'queues jobs correctly' do
+            Sidekiq::Testing.fake! do
+              model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes)
 
-            expect(tracked_job.class_name).to eq('FooJob')
-            expect(tracked_job.arguments).to eq([id1, id3, 1, 2])
-            expect(tracked_job).to be_pending
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            end
           end
         end
-      end
 
-      context 'without track_jobs option' do
-        it 'does not create records in the database' do
-          Sidekiq::Testing.fake! do
-            expect do
+        context 'with other_job_arguments option' do
+          it 'queues jobs correctly' do
+            Sidekiq::Testing.fake! do
               model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
-            end.not_to change { Gitlab::Database::BackgroundMigrationJob.count }
 
-            expect(BackgroundMigrationWorker.jobs.size).to eq(1)
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
+              expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            end
           end
         end
-      end
-    end
-
-    context 'when the model specifies a primary_column_name' do
-      let!(:id1) { create(:container_expiration_policy).id }
-      let!(:id2) { create(:container_expiration_policy).id }
-      let!(:id3) { create(:container_expiration_policy).id }
 
-      around do |example|
-        freeze_time { example.run }
-      end
+        context 'with initial_delay option' do
+          it 'queues jobs correctly' do
+            Sidekiq::Testing.fake! do
+              model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2], initial_delay: 10.minutes)
 
-      before do
-        ContainerExpirationPolicy.class_eval do
-          include EachBatch
+              expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id3, 1, 2]])
+              expect(worker_class.jobs[0]['at']).to eq(20.minutes.from_now.to_f)
+            end
+          end
         end
-      end
 
-      it 'returns the final expected delay', :aggregate_failures do
-        Sidekiq::Testing.fake! do
-          final_delay = model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, batch_size: 2, primary_column_name: :project_id)
+        context 'with track_jobs option' do
+          it 'creates a record for each job in the database' do
+            Sidekiq::Testing.fake! do
+              expect do
+                model.queue_background_migration_jobs_by_range_at_intervals(User, '::FooJob', 10.minutes,
+                  other_job_arguments: [1, 2], track_jobs: true)
+              end.to change { Gitlab::Database::BackgroundMigrationJob.count }.from(0).to(1)
 
-          expect(final_delay.to_f).to eq(20.minutes.to_f)
-          expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-          expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
-          expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
-        end
-      end
+              expect(worker_class.jobs.size).to eq(1)
 
-      context "when the primary_column_name is not an integer" do
-        it 'raises error' do
-          expect do
-            model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :enabled)
-          end.to raise_error(StandardError, /is not an integer column/)
-        end
-      end
+              tracked_job = Gitlab::Database::BackgroundMigrationJob.first
 
-      context "when the primary_column_name does not exist" do
-        it 'raises error' do
-          expect do
-            model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :foo)
-          end.to raise_error(StandardError, /does not have an ID column of foo/)
+              expect(tracked_job.class_name).to eq('FooJob')
+              expect(tracked_job.arguments).to eq([id1, id3, 1, 2])
+              expect(tracked_job).to be_pending
+            end
+          end
         end
-      end
-    end
-
-    context "when the model doesn't have an ID or primary_column_name column" do
-      it 'raises error (for now)' do
-        expect do
-          model.queue_background_migration_jobs_by_range_at_intervals(ProjectAuthorization, 'FooJob', 10.seconds)
-        end.to raise_error(StandardError, /does not have an ID/)
-      end
-    end
-  end
 
-  describe '#requeue_background_migration_jobs_by_range_at_intervals' do
-    let!(:job_class_name) { 'TestJob' }
-    let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) }
-    let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) }
-    let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) }
-    let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) }
+        context 'without track_jobs option' do
+          it 'does not create records in the database' do
+            Sidekiq::Testing.fake! do
+              expect do
+                model.queue_background_migration_jobs_by_range_at_intervals(User, 'FooJob', 10.minutes, other_job_arguments: [1, 2])
+              end.not_to change { Gitlab::Database::BackgroundMigrationJob.count }
 
-    around do |example|
-      freeze_time do
-        Sidekiq::Testing.fake! do
-          example.run
+              expect(worker_class.jobs.size).to eq(1)
+            end
+          end
         end
       end
-    end
-
-    subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) }
-
-    it 'returns the expected duration' do
-      expect(subject).to eq(20.minutes)
-    end
 
-    context 'when nothing is queued' do
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
+      context 'when the model specifies a primary_column_name' do
+        let!(:id1) { create(:container_expiration_policy).id }
+        let!(:id2) { create(:container_expiration_policy).id }
+        let!(:id3) { create(:container_expiration_policy).id }
 
-      it 'returns expected duration of zero when nothing gets queued' do
-        expect(subject).to eq(0)
-      end
-    end
-
-    it 'queues pending jobs' do
-      subject
+        around do |example|
+          freeze_time { example.run }
+        end
 
-      expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-      expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
-      expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-      expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
-    end
+        before do
+          ContainerExpirationPolicy.class_eval do
+            include EachBatch
+          end
+        end
 
-    context 'with batch_size option' do
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) }
+        it 'returns the final expected delay', :aggregate_failures do
+          Sidekiq::Testing.fake! do
+            final_delay = model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, batch_size: 2, primary_column_name: :project_id)
 
-      it 'returns the expected duration' do
-        expect(subject).to eq(20.minutes)
-      end
+            expect(final_delay.to_f).to eq(20.minutes.to_f)
+            expect(worker_class.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
+            expect(worker_class.jobs[0]['at']).to eq(10.minutes.from_now.to_f)
+            expect(worker_class.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
+            expect(worker_class.jobs[1]['at']).to eq(20.minutes.from_now.to_f)
+          end
+        end
 
-      it 'queues pending jobs' do
-        subject
+        context "when the primary_column_name is not an integer" do
+          it 'raises error' do
+            expect do
+              model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :enabled)
+            end.to raise_error(StandardError, /is not an integer column/)
+          end
+        end
 
-        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-        expect(BackgroundMigrationWorker.jobs[0]['at']).to be_nil
-        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+        context "when the primary_column_name does not exist" do
+          it 'raises error' do
+            expect do
+              model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :foo)
+            end.to raise_error(StandardError, /does not have an ID column of foo/)
+          end
+        end
       end
 
-      it 'retrieve jobs in batches' do
-        jobs = double('jobs')
-        expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs }
-        allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs }
-        expect(jobs).to receive(:each_batch).with(of: 1)
-
-        subject
+      context "when the model doesn't have an ID or primary_column_name column" do
+        it 'raises error (for now)' do
+          expect do
+            model.queue_background_migration_jobs_by_range_at_intervals(ProjectAuthorization, 'FooJob', 10.seconds)
+          end.to raise_error(StandardError, /does not have an ID/)
+        end
       end
     end
 
-    context 'with initial_delay option' do
-      let_it_be(:initial_delay) { 3.minutes }
+    describe '#requeue_background_migration_jobs_by_range_at_intervals' do
+      let!(:job_class_name) { 'TestJob' }
+      let!(:pending_job_1) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1, 2]) }
+      let!(:pending_job_2) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [3, 4]) }
+      let!(:successful_job_1) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [5, 6]) }
+      let!(:successful_job_2) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [7, 8]) }
 
-      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) }
-
-      it 'returns the expected duration' do
-        expect(subject).to eq(23.minutes)
+      around do |example|
+        freeze_time do
+          Sidekiq::Testing.fake! do
+            example.run
+          end
+        end
       end
 
-      it 'queues pending jobs' do
-        subject
+      subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes) }
 
-        expect(BackgroundMigrationWorker.jobs[0]['args']).to eq([job_class_name, [1, 2]])
-        expect(BackgroundMigrationWorker.jobs[0]['at']).to eq(3.minutes.from_now.to_f)
-        expect(BackgroundMigrationWorker.jobs[1]['args']).to eq([job_class_name, [3, 4]])
-        expect(BackgroundMigrationWorker.jobs[1]['at']).to eq(13.minutes.from_now.to_f)
+      it 'returns the expected duration' do
+        expect(subject).to eq(20.minutes)
       end
 
       context 'when nothing is queued' do
@@ -259,195 +218,226 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
           expect(subject).to eq(0)
         end
       end
-    end
-  end
 
-  describe '#perform_background_migration_inline?' do
-    it 'returns true in a test environment' do
-      stub_rails_env('test')
+      it 'queues pending jobs' do
+        subject
 
-      expect(model.perform_background_migration_inline?).to eq(true)
-    end
+        expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+        expect(worker_class.jobs[0]['at']).to be_nil
+        expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+        expect(worker_class.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+      end
 
-    it 'returns true in a development environment' do
-      stub_rails_env('development')
+      context 'with batch_size option' do
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, batch_size: 1) }
 
-      expect(model.perform_background_migration_inline?).to eq(true)
-    end
+        it 'returns the expected duration' do
+          expect(subject).to eq(20.minutes)
+        end
 
-    it 'returns false in a production environment' do
-      stub_rails_env('production')
+        it 'queues pending jobs' do
+          subject
 
-      expect(model.perform_background_migration_inline?).to eq(false)
-    end
-  end
+          expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+          expect(worker_class.jobs[0]['at']).to be_nil
+          expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+          expect(worker_class.jobs[1]['at']).to eq(10.minutes.from_now.to_f)
+        end
 
-  describe '#migrate_async' do
-    it 'calls BackgroundMigrationWorker.perform_async' do
-      expect(BackgroundMigrationWorker).to receive(:perform_async).with("Class", "hello", "world")
+        it 'retrieve jobs in batches' do
+          jobs = double('jobs')
+          expect(Gitlab::Database::BackgroundMigrationJob).to receive(:pending) { jobs }
+          allow(jobs).to receive(:where).with(class_name: job_class_name) { jobs }
+          expect(jobs).to receive(:each_batch).with(of: 1)
 
-      model.migrate_async("Class", "hello", "world")
-    end
+          subject
+        end
+      end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+      context 'with initial_delay option' do
+        let_it_be(:initial_delay) { 3.minutes }
 
-      model.migrate_async('Class', 'hello', 'world')
-    end
-  end
+        subject { model.requeue_background_migration_jobs_by_range_at_intervals(job_class_name, 10.minutes, initial_delay: initial_delay) }
 
-  describe '#migrate_in' do
-    it 'calls BackgroundMigrationWorker.perform_in' do
-      expect(BackgroundMigrationWorker).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
+        it 'returns the expected duration' do
+          expect(subject).to eq(23.minutes)
+        end
 
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
-    end
+        it 'queues pending jobs' do
+          subject
+
+          expect(worker_class.jobs[0]['args']).to eq([job_class_name, [1, 2]])
+          expect(worker_class.jobs[0]['at']).to eq(3.minutes.from_now.to_f)
+          expect(worker_class.jobs[1]['args']).to eq([job_class_name, [3, 4]])
+          expect(worker_class.jobs[1]['at']).to eq(13.minutes.from_now.to_f)
+        end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+        context 'when nothing is queued' do
+          subject { model.requeue_background_migration_jobs_by_range_at_intervals('FakeJob', 10.minutes) }
 
-      model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
+          it 'returns expected duration of zero when nothing gets queued' do
+            expect(subject).to eq(0)
+          end
+        end
+      end
     end
-  end
 
-  describe '#bulk_migrate_async' do
-    it 'calls BackgroundMigrationWorker.bulk_perform_async' do
-      expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([%w(Class hello world)])
+    describe '#finalized_background_migration' do
+      let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(worker_class) }
 
-      model.bulk_migrate_async([%w(Class hello world)])
-    end
+      let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
+      let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
+      let!(:job_class_name) { 'TestJob' }
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+      let!(:job_class) do
+        Class.new do
+          def perform(*arguments)
+            Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
+          end
+        end
+      end
 
-      model.bulk_migrate_async([%w(Class hello world)])
-    end
-  end
+      before do
+        allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
+          .with('main').and_return(coordinator)
 
-  describe '#bulk_migrate_in' do
-    it 'calls BackgroundMigrationWorker.bulk_perform_in_' do
-      expect(BackgroundMigrationWorker).to receive(:bulk_perform_in).with(10.minutes, [%w(Class hello world)])
+        expect(coordinator).to receive(:migration_class_for)
+          .with(job_class_name).at_least(:once) { job_class }
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
-    end
+        Sidekiq::Testing.disable! do
+          worker_class.perform_async(job_class_name, [1, 2])
+          worker_class.perform_async(job_class_name, [3, 4])
+          worker_class.perform_in(10, job_class_name, [5, 6])
+          worker_class.perform_in(20, job_class_name, [7, 8])
+        end
+      end
 
-    it 'pushes a context with the current class name as caller_id' do
-      expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
+      it_behaves_like 'finalized tracked background migration', worker_class do
+        before do
+          model.finalize_background_migration(job_class_name)
+        end
+      end
 
-      model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
-    end
-  end
+      context 'when removing all tracked job records' do
+        let!(:job_class) do
+          Class.new do
+            def perform(*arguments)
+              # Force pending jobs to remain pending
+            end
+          end
+        end
 
-  describe '#delete_queued_jobs' do
-    let(:job1) { double }
-    let(:job2) { double }
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
+        end
 
-    it 'deletes all queued jobs for the given background migration' do
-      expect(Gitlab::BackgroundMigration).to receive(:steal).with('BackgroundMigrationClassName') do |&block|
-        expect(block.call(job1)).to be(false)
-        expect(block.call(job2)).to be(false)
+        it_behaves_like 'finalized tracked background migration', worker_class
+        it_behaves_like 'removed tracked jobs', 'pending'
+        it_behaves_like 'removed tracked jobs', 'succeeded'
       end
 
-      expect(job1).to receive(:delete)
-      expect(job2).to receive(:delete)
+      context 'when retaining all tracked job records' do
+        before do
+          model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
+        end
 
-      model.delete_queued_jobs('BackgroundMigrationClassName')
-    end
-  end
+        it_behaves_like 'finalized background migration', worker_class
+        include_examples 'retained tracked jobs', 'succeeded'
+      end
 
-  describe '#finalized_background_migration' do
-    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(BackgroundMigrationWorker) }
+      context 'during retry race condition' do
+        let!(:job_class) do
+          Class.new do
+            class << self
+              attr_accessor :worker_class
 
-    let!(:job_class_name) { 'TestJob' }
-    let!(:job_class) { Class.new }
-    let!(:job_perform_method) do
-      ->(*arguments) do
-        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
-          # Value is 'TestJob' defined by :job_class_name in the let! above.
-          # Scoping prohibits us from directly referencing job_class_name.
-          RSpec.current_example.example_group_instance.job_class_name,
-          arguments
-        )
-      end
-    end
+              def queue_items_added
+                @queue_items_added ||= []
+              end
+            end
 
-    let!(:tracked_pending_job) { create(:background_migration_job, class_name: job_class_name, status: :pending, arguments: [1]) }
-    let!(:tracked_successful_job) { create(:background_migration_job, class_name: job_class_name, status: :succeeded, arguments: [2]) }
+            def worker_class
+              self.class.worker_class
+            end
 
-    before do
-      job_class.define_method(:perform, job_perform_method)
+            def queue_items_added
+              self.class.queue_items_added
+            end
 
-      allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-        .with('main').and_return(job_coordinator)
+            def perform(*arguments)
+              Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded('TestJob', arguments)
 
-      expect(job_coordinator).to receive(:migration_class_for)
-        .with(job_class_name).at_least(:once) { job_class }
+              # Mock another process pushing queue jobs.
+              if self.class.queue_items_added.count < 10
+                Sidekiq::Testing.disable! do
+                  queue_items_added << worker_class.perform_async('TestJob', [Time.current])
+                  queue_items_added << worker_class.perform_in(10, 'TestJob', [Time.current])
+                end
+              end
+            end
+          end
+        end
 
-      Sidekiq::Testing.disable! do
-        BackgroundMigrationWorker.perform_async(job_class_name, [1, 2])
-        BackgroundMigrationWorker.perform_async(job_class_name, [3, 4])
-        BackgroundMigrationWorker.perform_in(10, job_class_name, [5, 6])
-        BackgroundMigrationWorker.perform_in(20, job_class_name, [7, 8])
-      end
-    end
+        it_behaves_like 'finalized tracked background migration', worker_class do
+          before do
+            # deliberately set the worker class on our test job since it won't be pulled from the surrounding scope
+            job_class.worker_class = worker_class
 
-    it_behaves_like 'finalized tracked background migration' do
-      before do
-        model.finalize_background_migration(job_class_name)
+            model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+          end
+        end
       end
     end
 
-    context 'when removing all tracked job records' do
-      # Force pending jobs to remain pending.
-      let!(:job_perform_method) { ->(*arguments) { } }
+    describe '#migrate_in' do
+      it 'calls perform_in for the correct worker' do
+        expect(worker_class).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
 
-      before do
-        model.finalize_background_migration(job_class_name, delete_tracking_jobs: %w[pending succeeded])
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
       end
 
-      it_behaves_like 'finalized tracked background migration'
-      it_behaves_like 'removed tracked jobs', 'pending'
-      it_behaves_like 'removed tracked jobs', 'succeeded'
-    end
+      it 'pushes a context with the current class name as caller_id' do
+        expect(Gitlab::ApplicationContext).to receive(:with_context).with(caller_id: model.class.to_s)
 
-    context 'when retaining all tracked job records' do
-      before do
-        model.finalize_background_migration(job_class_name, delete_tracking_jobs: false)
+        model.migrate_in(10.minutes, 'Class', 'Hello', 'World')
       end
 
-      it_behaves_like 'finalized background migration'
-      include_examples 'retained tracked jobs', 'succeeded'
-    end
+      context 'when a specific coordinator is given' do
+        let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.for_tracking_database('main') }
 
-    context 'during retry race condition' do
-      let(:queue_items_added) { [] }
-      let!(:job_perform_method) do
-        ->(*arguments) do
-          Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
-            RSpec.current_example.example_group_instance.job_class_name,
-            arguments
-          )
-
-          # Mock another process pushing queue jobs.
-          queue_items_added = RSpec.current_example.example_group_instance.queue_items_added
-          if queue_items_added.count < 10
-            Sidekiq::Testing.disable! do
-              job_class_name = RSpec.current_example.example_group_instance.job_class_name
-              queue_items_added << BackgroundMigrationWorker.perform_async(job_class_name, [Time.current])
-              queue_items_added << BackgroundMigrationWorker.perform_in(10, job_class_name, [Time.current])
-            end
-          end
+        it 'uses that coordinator' do
+          expect(coordinator).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World').and_call_original
+          expect(worker_class).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World')
+
+          model.migrate_in(10.minutes, 'Class', 'Hello', 'World', coordinator: coordinator)
         end
       end
+    end
 
-      it_behaves_like 'finalized tracked background migration' do
-        before do
-          model.finalize_background_migration(job_class_name, delete_tracking_jobs: ['succeeded'])
+    describe '#delete_queued_jobs' do
+      let(:job1) { double }
+      let(:job2) { double }
+
+      it 'deletes all queued jobs for the given background migration' do
+        expect_next_instance_of(Gitlab::BackgroundMigration::JobCoordinator) do |coordinator|
+          expect(coordinator).to receive(:steal).with('BackgroundMigrationClassName') do |&block|
+            expect(block.call(job1)).to be(false)
+            expect(block.call(job2)).to be(false)
+          end
         end
+
+        expect(job1).to receive(:delete)
+        expect(job2).to receive(:delete)
+
+        model.delete_queued_jobs('BackgroundMigrationClassName')
       end
     end
   end
 
+  context 'when the migration is running against the main database' do
+    it_behaves_like 'helpers that enqueue background migrations', BackgroundMigrationWorker, 'main'
+  end
+
   describe '#delete_job_tracking' do
     let!(:job_class_name) { 'TestJob' }
 
diff --git a/spec/lib/gitlab/database/migrations/runner_spec.rb b/spec/lib/gitlab/database/migrations/runner_spec.rb
index 4616bd6941e..7dc965c84fa 100644
--- a/spec/lib/gitlab/database/migrations/runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/runner_spec.rb
@@ -28,7 +28,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
     allow(ActiveRecord::Migrator).to receive(:new) do |dir, _all_migrations, _schema_migration_class, version_to_migrate|
       migrator = double(ActiveRecord::Migrator)
       expect(migrator).to receive(:run) do
-        migration_runs << OpenStruct.new(dir: dir, version_to_migrate: version_to_migrate)
+        migration_runs << double('migrator', dir: dir, version_to_migrate: version_to_migrate)
       end
       migrator
     end
diff --git a/spec/lib/gitlab/database/no_cross_db_foreign_keys_spec.rb b/spec/lib/gitlab/database/no_cross_db_foreign_keys_spec.rb
new file mode 100644
index 00000000000..e5a8143fcc3
--- /dev/null
+++ b/spec/lib/gitlab/database/no_cross_db_foreign_keys_spec.rb
@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'cross-database foreign keys' do
+  # TODO: We are trying to empty out this list in
+  # https://gitlab.com/groups/gitlab-org/-/epics/7249 . Once we are done we can
+  # keep this test and assert that there are no cross-db foreign keys. We
+  # should not be adding anything to this list but should instead only add new
+  # loose foreign keys
+  # https://docs.gitlab.com/ee/development/database/loose_foreign_keys.html .
+  let(:allowed_cross_database_foreign_keys) do
+    %w(
+      ci_build_report_results.project_id
+      ci_builds.project_id
+      ci_builds_metadata.project_id
+      ci_daily_build_group_report_results.group_id
+      ci_daily_build_group_report_results.project_id
+      ci_freeze_periods.project_id
+      ci_job_artifacts.project_id
+      ci_job_token_project_scope_links.added_by_id
+      ci_job_token_project_scope_links.source_project_id
+      ci_job_token_project_scope_links.target_project_id
+      ci_pending_builds.namespace_id
+      ci_pending_builds.project_id
+      ci_pipeline_schedules.owner_id
+      ci_pipeline_schedules.project_id
+      ci_pipelines.merge_request_id
+      ci_pipelines.project_id
+      ci_project_monthly_usages.project_id
+      ci_refs.project_id
+      ci_resource_groups.project_id
+      ci_runner_namespaces.namespace_id
+      ci_runner_projects.project_id
+      ci_running_builds.project_id
+      ci_sources_pipelines.project_id
+      ci_sources_pipelines.source_project_id
+      ci_sources_projects.source_project_id
+      ci_stages.project_id
+      ci_subscriptions_projects.downstream_project_id
+      ci_subscriptions_projects.upstream_project_id
+      ci_triggers.owner_id
+      ci_triggers.project_id
+      ci_unit_tests.project_id
+      ci_variables.project_id
+      dast_profiles_pipelines.ci_pipeline_id
+      dast_scanner_profiles_builds.ci_build_id
+      dast_site_profiles_builds.ci_build_id
+      dast_site_profiles_pipelines.ci_pipeline_id
+      external_pull_requests.project_id
+      merge_requests.head_pipeline_id
+      merge_trains.pipeline_id
+      requirements_management_test_reports.build_id
+      security_scans.build_id
+      vulnerability_feedback.pipeline_id
+      vulnerability_occurrence_pipelines.pipeline_id
+      vulnerability_statistics.latest_pipeline_id
+    ).freeze
+  end
+
+  def foreign_keys_for(table_name)
+    ApplicationRecord.connection.foreign_keys(table_name)
+  end
+
+  def is_cross_db?(fk_record)
+    Gitlab::Database::GitlabSchema.table_schemas([fk_record.from_table, fk_record.to_table]).many?
+  end
+
+  it 'onlies have allowed list of cross-database foreign keys', :aggregate_failures do
+    all_tables = ApplicationRecord.connection.data_sources
+
+    all_tables.each do |table|
+      foreign_keys_for(table).each do |fk|
+        if is_cross_db?(fk)
+          column = "#{fk.from_table}.#{fk.column}"
+          expect(allowed_cross_database_foreign_keys).to include(column), "Found extra cross-database foreign key #{column} referencing #{fk.to_table} with constraint name #{fk.name}. When a foreign key references another database you must use a Loose Foreign Key instead https://docs.gitlab.com/ee/development/database/loose_foreign_keys.html ."
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index 5e107109fc9..64dcdb9628a 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -18,7 +18,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: [], after_adding_partitions: nil) }
     let(:connection) { ActiveRecord::Base.connection }
-    let(:table) { "some_table" }
+    let(:table) { "issues" }
 
     before do
       allow(connection).to receive(:table_exists?).and_call_original
@@ -36,6 +36,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
 
     it 'creates the partition' do
+      expect(connection).to receive(:execute).with("LOCK TABLE \"#{table}\" IN ACCESS EXCLUSIVE MODE")
       expect(connection).to receive(:execute).with(partitions.first.to_sql)
       expect(connection).to receive(:execute).with(partitions.second.to_sql)
 
diff --git a/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
index 636a09e5710..1cec0463055 100644
--- a/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Partitioning::SlidingListStrategy do
   let(:connection) { ActiveRecord::Base.connection }
   let(:table_name) { :_test_partitioned_test }
-  let(:model) { double('model', table_name: table_name, ignored_columns: %w[partition]) }
+  let(:model) { double('model', table_name: table_name, ignored_columns: %w[partition], connection: connection) }
   let(:next_partition_if) { double('next_partition_if') }
   let(:detach_partition_if) { double('detach_partition_if') }
 
@@ -94,7 +94,8 @@ RSpec.describe Gitlab::Database::Partitioning::SlidingListStrategy do
       let(:detach_partition_if) { ->(p) { p != 5 } }
 
       it 'is the leading set of partitions before that value' do
-        expect(strategy.extra_partitions.map(&:value)).to contain_exactly(1, 2, 3, 4)
+        # should not contain partition 2 since it's the default value for the partition column
+        expect(strategy.extra_partitions.map(&:value)).to contain_exactly(1, 3, 4)
       end
     end
 
@@ -102,7 +103,7 @@ RSpec.describe Gitlab::Database::Partitioning::SlidingListStrategy do
       let(:detach_partition_if) { proc { true } }
 
       it 'is all but the most recent partition', :aggregate_failures do
-        expect(strategy.extra_partitions.map(&:value)).to contain_exactly(1, 2, 3, 4, 5, 6, 7, 8, 9)
+        expect(strategy.extra_partitions.map(&:value)).to contain_exactly(1, 3, 4, 5, 6, 7, 8, 9)
 
         expect(strategy.current_partitions.map(&:value).max).to eq(10)
       end
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table_spec.rb
index c43b51e10a0..3072c413246 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table_spec.rb
@@ -3,14 +3,15 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable, '#perform' do
-  subject { described_class.new }
+  subject(:backfill_job) { described_class.new(connection: connection) }
 
+  let(:connection) { ActiveRecord::Base.connection }
   let(:source_table) { '_test_partitioning_backfills' }
   let(:destination_table) { "#{source_table}_part" }
   let(:unique_key) { 'id' }
 
   before do
-    allow(subject).to receive(:transaction_open?).and_return(false)
+    allow(backfill_job).to receive(:transaction_open?).and_return(false)
   end
 
   context 'when the destination table exists' do
@@ -50,10 +51,9 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
       stub_const("#{described_class}::SUB_BATCH_SIZE", 2)
       stub_const("#{described_class}::PAUSE_SECONDS", pause_seconds)
 
-      allow(subject).to receive(:sleep)
+      allow(backfill_job).to receive(:sleep)
     end
 
-    let(:connection) { ActiveRecord::Base.connection }
     let(:source_model) { Class.new(ActiveRecord::Base) }
     let(:destination_model) { Class.new(ActiveRecord::Base) }
     let(:timestamp) { Time.utc(2020, 1, 2).round }
@@ -66,7 +66,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
     it 'copies data into the destination table idempotently' do
       expect(destination_model.count).to eq(0)
 
-      subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+      backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
 
       expect(destination_model.count).to eq(3)
 
@@ -76,7 +76,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
         expect(destination_record.attributes).to eq(source_record.attributes)
       end
 
-      subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+      backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
 
       expect(destination_model.count).to eq(3)
     end
@@ -87,13 +87,13 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
         expect(bulk_copy).to receive(:copy_between).with(source3.id, source3.id)
       end
 
-      subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+      backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
     end
 
     it 'pauses after copying each sub-batch' do
-      expect(subject).to receive(:sleep).with(pause_seconds).twice
+      expect(backfill_job).to receive(:sleep).with(pause_seconds).twice
 
-      subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+      backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
     end
 
     it 'marks each job record as succeeded after processing' do
@@ -103,7 +103,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
       expect(::Gitlab::Database::BackgroundMigrationJob).to receive(:mark_all_as_succeeded).and_call_original
 
       expect do
-        subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+        backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
       end.to change { ::Gitlab::Database::BackgroundMigrationJob.succeeded.count }.from(0).to(1)
     end
 
@@ -111,24 +111,24 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
       create(:background_migration_job, class_name: "::#{described_class.name}",
         arguments: [source1.id, source3.id, source_table, destination_table, unique_key])
 
-      jobs_updated = subject.perform(source1.id, source3.id, source_table, destination_table, unique_key)
+      jobs_updated = backfill_job.perform(source1.id, source3.id, source_table, destination_table, unique_key)
 
       expect(jobs_updated).to eq(1)
     end
 
     context 'when the job is run within an explicit transaction block' do
-      let(:mock_connection) { double('connection') }
+      subject(:backfill_job) { described_class.new(connection: mock_connection) }
 
-      before do
-        allow(subject).to receive(:connection).and_return(mock_connection)
-        allow(subject).to receive(:transaction_open?).and_return(true)
-      end
+      let(:mock_connection) { double('connection') }
 
       it 'raises an error before copying data' do
+        expect(backfill_job).to receive(:transaction_open?).and_call_original
+
+        expect(mock_connection).to receive(:transaction_open?).and_return(true)
         expect(mock_connection).not_to receive(:execute)
 
         expect do
-          subject.perform(1, 100, source_table, destination_table, unique_key)
+          backfill_job.perform(1, 100, source_table, destination_table, unique_key)
         end.to raise_error(/Aborting job to backfill partitioned #{source_table}/)
 
         expect(destination_model.count).to eq(0)
@@ -137,24 +137,25 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::BackfillPartition
   end
 
   context 'when the destination table does not exist' do
+    subject(:backfill_job) { described_class.new(connection: mock_connection) }
+
     let(:mock_connection) { double('connection') }
     let(:mock_logger) { double('logger') }
 
     before do
-      allow(subject).to receive(:connection).and_return(mock_connection)
-      allow(subject).to receive(:logger).and_return(mock_logger)
-
-      expect(mock_connection).to receive(:table_exists?).with(destination_table).and_return(false)
+      allow(backfill_job).to receive(:logger).and_return(mock_logger)
       allow(mock_logger).to receive(:warn)
     end
 
     it 'exits without attempting to copy data' do
+      expect(mock_connection).to receive(:table_exists?).with(destination_table).and_return(false)
       expect(mock_connection).not_to receive(:execute)
 
       subject.perform(1, 100, source_table, destination_table, unique_key)
     end
 
     it 'logs a warning message that the job was skipped' do
+      expect(mock_connection).to receive(:table_exists?).with(destination_table).and_return(false)
       expect(mock_logger).to receive(:warn).with(/#{destination_table} does not exist/)
 
       subject.perform(1, 100, source_table, destination_table, unique_key)
diff --git a/spec/lib/gitlab/database/reflection_spec.rb b/spec/lib/gitlab/database/reflection_spec.rb
index 7c3d797817d..efc5bd1c1e1 100644
--- a/spec/lib/gitlab/database/reflection_spec.rb
+++ b/spec/lib/gitlab/database/reflection_spec.rb
@@ -259,6 +259,66 @@ RSpec.describe Gitlab::Database::Reflection do
     end
   end
 
+  describe '#flavor', :delete do
+    let(:result) { [double] }
+    let(:connection) { database.model.connection }
+
+    def stub_statements(statements)
+      statements = Array.wrap(statements)
+      execute = connection.method(:execute)
+
+      allow(connection).to receive(:execute) do |arg|
+        if statements.include?(arg)
+          result
+        else
+          execute.call(arg)
+        end
+      end
+    end
+
+    it 're-raises exceptions not matching expected messages' do
+      expect(database.model.connection)
+        .to receive(:execute)
+        .and_raise(ActiveRecord::StatementInvalid, 'Something else')
+
+      expect { database.flavor }.to raise_error ActiveRecord::StatementInvalid, /Something else/
+    end
+
+    it 'recognizes Amazon Aurora PostgreSQL' do
+      stub_statements(['SHOW rds.extensions', 'SELECT AURORA_VERSION()'])
+
+      expect(database.flavor).to eq('Amazon Aurora PostgreSQL')
+    end
+
+    it 'recognizes PostgreSQL on Amazon RDS' do
+      stub_statements('SHOW rds.extensions')
+
+      expect(database.flavor).to eq('PostgreSQL on Amazon RDS')
+    end
+
+    it 'recognizes CloudSQL for PostgreSQL' do
+      stub_statements('SHOW cloudsql.iam_authentication')
+
+      expect(database.flavor).to eq('Cloud SQL for PostgreSQL')
+    end
+
+    it 'recognizes Azure Database for PostgreSQL - Flexible Server' do
+      stub_statements(["SELECT datname FROM pg_database WHERE datname = 'azure_maintenance'", 'SHOW azure.extensions'])
+
+      expect(database.flavor).to eq('Azure Database for PostgreSQL - Flexible Server')
+    end
+
+    it 'recognizes Azure Database for PostgreSQL - Single Server' do
+      stub_statements("SELECT datname FROM pg_database WHERE datname = 'azure_maintenance'")
+
+      expect(database.flavor).to eq('Azure Database for PostgreSQL - Single Server')
+    end
+
+    it 'returns nil if can not recognize the flavor' do
+      expect(database.flavor).to be_nil
+    end
+  end
+
   describe '#config' do
     it 'returns a HashWithIndifferentAccess' do
       expect(database.config)
diff --git a/spec/lib/gitlab/database/reindexing/coordinator_spec.rb b/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
index 0afbe46b7f1..bb91617714a 100644
--- a/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
+++ b/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
@@ -6,30 +6,34 @@ RSpec.describe Gitlab::Database::Reindexing::Coordinator do
   include Database::DatabaseHelpers
   include ExclusiveLeaseHelpers
 
-  describe '.perform' do
-    subject { described_class.new(index, notifier).perform }
-
-    let(:index) { create(:postgres_index) }
-    let(:notifier) { instance_double(Gitlab::Database::Reindexing::GrafanaNotifier, notify_start: nil, notify_end: nil) }
-    let(:reindexer) { instance_double(Gitlab::Database::Reindexing::ReindexConcurrently, perform: nil) }
-    let(:action) { create(:reindex_action, index: index) }
+  let(:notifier) { instance_double(Gitlab::Database::Reindexing::GrafanaNotifier, notify_start: nil, notify_end: nil) }
+  let(:index) { create(:postgres_index) }
+  let(:connection) { index.connection }
 
-    let!(:lease) { stub_exclusive_lease(lease_key, uuid, timeout: lease_timeout) }
-    let(:lease_key) { "gitlab/database/reindexing/coordinator/#{Gitlab::Database::PRIMARY_DATABASE_NAME}" }
-    let(:lease_timeout) { 1.day }
-    let(:uuid) { 'uuid' }
+  let!(:lease) { stub_exclusive_lease(lease_key, uuid, timeout: lease_timeout) }
+  let(:lease_key) { "gitlab/database/reindexing/coordinator/#{Gitlab::Database::PRIMARY_DATABASE_NAME}" }
+  let(:lease_timeout) { 1.day }
+  let(:uuid) { 'uuid' }
 
-    around do |example|
-      model = Gitlab::Database.database_base_models[Gitlab::Database::PRIMARY_DATABASE_NAME]
+  around do |example|
+    model = Gitlab::Database.database_base_models[Gitlab::Database::PRIMARY_DATABASE_NAME]
 
-      Gitlab::Database::SharedModel.using_connection(model.connection) do
-        example.run
-      end
+    Gitlab::Database::SharedModel.using_connection(model.connection) do
+      example.run
     end
+  end
 
-    before do
-      swapout_view_for_table(:postgres_indexes)
+  before do
+    swapout_view_for_table(:postgres_indexes)
+  end
 
+  describe '#perform' do
+    subject { described_class.new(index, notifier).perform }
+
+    let(:reindexer) { instance_double(Gitlab::Database::Reindexing::ReindexConcurrently, perform: nil) }
+    let(:action) { create(:reindex_action, index: index) }
+
+    before do
       allow(Gitlab::Database::Reindexing::ReindexConcurrently).to receive(:new).with(index).and_return(reindexer)
       allow(Gitlab::Database::Reindexing::ReindexAction).to receive(:create_for).with(index).and_return(action)
     end
@@ -87,4 +91,40 @@ RSpec.describe Gitlab::Database::Reindexing::Coordinator do
       end
     end
   end
+
+  describe '#drop' do
+    let(:connection) { index.connection }
+
+    subject(:drop) { described_class.new(index, notifier).drop }
+
+    context 'when exclusive lease is granted' do
+      it 'drops the index with lock retries' do
+        expect(lease).to receive(:try_obtain).ordered.and_return(uuid)
+
+        expect_query("SET lock_timeout TO '60000ms'")
+        expect_query("DROP INDEX CONCURRENTLY IF EXISTS \"public\".\"#{index.name}\"")
+        expect_query("RESET idle_in_transaction_session_timeout; RESET lock_timeout")
+
+        expect(Gitlab::ExclusiveLease).to receive(:cancel).ordered.with(lease_key, uuid)
+
+        drop
+      end
+
+      def expect_query(sql)
+        expect(connection).to receive(:execute).ordered.with(sql).and_wrap_original do |method, sql|
+          method.call(sql.sub(/CONCURRENTLY/, ''))
+        end
+      end
+    end
+
+    context 'when exclusive lease is not granted' do
+      it 'does not drop the index' do
+        expect(lease).to receive(:try_obtain).ordered.and_return(false)
+        expect(Gitlab::Database::WithLockRetriesOutsideTransaction).not_to receive(:new)
+        expect(connection).not_to receive(:execute)
+
+        drop
+      end
+    end
+  end
 end