Add latest changes from gitlab-org/gitlab@14-8-stable-eev14.8.0-rc42

25 files changed, 584 insertions, 209 deletions

diff --git a/spec/lib/gitlab/database/background_migration/batch_metrics_spec.rb b/spec/lib/gitlab/database/background_migration/batch_metrics_spec.rb
index e96862fbc2d..66983733411 100644
--- a/spec/lib/gitlab/database/background_migration/batch_metrics_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batch_metrics_spec.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'fast_spec_helper'
+require 'spec_helper'
 
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchMetrics do
   let(:batch_metrics) { described_class.new }
diff --git a/spec/lib/gitlab/database/background_migration/batch_optimizer_spec.rb b/spec/lib/gitlab/database/background_migration/batch_optimizer_spec.rb
index 95863ce3765..c367f4a4493 100644
--- a/spec/lib/gitlab/database/background_migration/batch_optimizer_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batch_optimizer_spec.rb
@@ -6,7 +6,11 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchOptimizer do
   describe '#optimize' do
     subject { described_class.new(migration, number_of_jobs: number_of_jobs, ema_alpha: ema_alpha).optimize! }
 
-    let(:migration) { create(:batched_background_migration, batch_size: batch_size, sub_batch_size: 100, interval: 120) }
+    let(:migration_params) { {} }
+    let(:migration) do
+      params = { batch_size: batch_size, sub_batch_size: 100, interval: 120 }.merge(migration_params)
+      create(:batched_background_migration, params)
+    end
 
     let(:batch_size) { 10_000 }
 
@@ -87,6 +91,17 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchOptimizer do
 
         expect { subject }.to change { migration.reload.batch_size }.to(2_000_000)
       end
+
+      context 'when max_batch_size is set' do
+        let(:max_batch_size) { 10000 }
+        let(:migration_params) { { max_batch_size: max_batch_size } }
+
+        it 'caps the batch size at max_batch_size' do
+          mock_efficiency(0.7)
+
+          expect { subject }.to change { migration.reload.batch_size }.to(max_batch_size)
+        end
+      end
     end
 
     context 'reaching the lower limit for the batch size' do
diff --git a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
index c4364826ee2..7338ea657b9 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
@@ -7,18 +7,85 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
 
   describe 'associations' do
     it { is_expected.to belong_to(:batched_migration).with_foreign_key(:batched_background_migration_id) }
+    it { is_expected.to have_many(:batched_job_transition_logs).with_foreign_key(:batched_background_migration_job_id) }
+  end
+
+  describe 'state machine' do
+    let_it_be(:job) { create(:batched_background_migration_job, :failed) }
+
+    context 'when a job is running' do
+      it 'logs the transition' do
+        expect(Gitlab::AppLogger).to receive(:info).with( { batched_job_id: job.id, message: 'BatchedJob transition', new_state: :running, previous_state: :failed } )
+
+        expect { job.run! }.to change(job, :started_at)
+      end
+    end
+
+    context 'when a job succeed' do
+      let(:job) { create(:batched_background_migration_job, :running) }
+
+      it 'logs the transition' do
+        expect(Gitlab::AppLogger).to receive(:info).with( { batched_job_id: job.id, message: 'BatchedJob transition', new_state: :succeeded, previous_state: :running } )
+
+        job.succeed!
+      end
+
+      it 'updates the finished_at' do
+        expect { job.succeed! }.to change(job, :finished_at).from(nil).to(Time)
+      end
+
+      it 'creates a new transition log' do
+        job.succeed!
+
+        transition_log = job.batched_job_transition_logs.first
+
+        expect(transition_log.next_status).to eq('succeeded')
+        expect(transition_log.exception_class).to be_nil
+        expect(transition_log.exception_message).to be_nil
+      end
+    end
+
+    context 'when a job fails' do
+      let(:job) { create(:batched_background_migration_job, :running) }
+
+      it 'logs the transition' do
+        expect(Gitlab::AppLogger).to receive(:info).with( { batched_job_id: job.id, message: 'BatchedJob transition', new_state: :failed, previous_state: :running } )
+
+        job.failure!
+      end
+
+      it 'tracks the exception' do
+        expect(Gitlab::ErrorTracking).to receive(:track_exception).with(RuntimeError, { batched_job_id: job.id } )
+
+        job.failure!(error: RuntimeError.new)
+      end
+
+      it 'updates the finished_at' do
+        expect { job.failure! }.to change(job, :finished_at).from(nil).to(Time)
+      end
+
+      it 'creates a new transition log' do
+        job.failure!(error: RuntimeError.new)
+
+        transition_log = job.batched_job_transition_logs.first
+
+        expect(transition_log.next_status).to eq('failed')
+        expect(transition_log.exception_class).to eq('RuntimeError')
+        expect(transition_log.exception_message).to eq('RuntimeError')
+      end
+    end
   end
 
   describe 'scopes' do
     let_it_be(:fixed_time) { Time.new(2021, 04, 27, 10, 00, 00, 00) }
 
-    let_it_be(:pending_job) { create(:batched_background_migration_job, status: :pending, updated_at: fixed_time) }
-    let_it_be(:running_job) { create(:batched_background_migration_job, status: :running, updated_at: fixed_time) }
-    let_it_be(:stuck_job) { create(:batched_background_migration_job, status: :pending, updated_at: fixed_time - described_class::STUCK_JOBS_TIMEOUT) }
-    let_it_be(:failed_job) { create(:batched_background_migration_job, status: :failed, attempts: 1) }
+    let_it_be(:pending_job) { create(:batched_background_migration_job, :pending, updated_at: fixed_time) }
+    let_it_be(:running_job) { create(:batched_background_migration_job, :running, updated_at: fixed_time) }
+    let_it_be(:stuck_job) { create(:batched_background_migration_job, :pending, updated_at: fixed_time - described_class::STUCK_JOBS_TIMEOUT) }
+    let_it_be(:failed_job) { create(:batched_background_migration_job, :failed, attempts: 1) }
 
-    let!(:max_attempts_failed_job) { create(:batched_background_migration_job, status: :failed, attempts: described_class::MAX_ATTEMPTS) }
-    let!(:succeeded_job) { create(:batched_background_migration_job, status: :succeeded) }
+    let!(:max_attempts_failed_job) { create(:batched_background_migration_job, :failed, attempts: described_class::MAX_ATTEMPTS) }
+    let!(:succeeded_job) { create(:batched_background_migration_job, :succeeded) }
 
     before do
       travel_to fixed_time
@@ -82,10 +149,10 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     subject { job.time_efficiency }
 
     let(:migration) { build(:batched_background_migration, interval: 120.seconds) }
-    let(:job) { build(:batched_background_migration_job, status: :succeeded, batched_migration: migration) }
+    let(:job) { build(:batched_background_migration_job, :succeeded, batched_migration: migration) }
 
     context 'when job has not yet succeeded' do
-      let(:job) { build(:batched_background_migration_job, status: :running) }
+      let(:job) { build(:batched_background_migration_job, :running) }
 
       it 'returns nil' do
         expect(subject).to be_nil
@@ -130,7 +197,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
   end
 
   describe '#split_and_retry!' do
-    let!(:job) { create(:batched_background_migration_job, batch_size: 10, min_value: 6, max_value: 15, status: :failed, attempts: 3) }
+    let!(:job) { create(:batched_background_migration_job, :failed, batch_size: 10, min_value: 6, max_value: 15, attempts: 3) }
 
     context 'when job can be split' do
       before do
@@ -146,7 +213,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
           min_value: 6,
           max_value: 10,
           batch_size: 5,
-          status: 'failed',
+          status_name: :failed,
           attempts: 0,
           started_at: nil,
           finished_at: nil,
@@ -160,7 +227,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
           min_value: 11,
           max_value: 15,
           batch_size: 5,
-          status: 'failed',
+          status_name: :failed,
           attempts: 0,
           started_at: nil,
           finished_at: nil,
@@ -177,7 +244,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     end
 
     context 'when job is not failed' do
-      let!(:job) { create(:batched_background_migration_job, status: :succeeded) }
+      let!(:job) { create(:batched_background_migration_job, :succeeded) }
 
       it 'raises an exception' do
         expect { job.split_and_retry! }.to raise_error 'Only failed jobs can be split'
@@ -185,7 +252,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     end
 
     context 'when batch size is already 1' do
-      let!(:job) { create(:batched_background_migration_job, batch_size: 1, status: :failed) }
+      let!(:job) { create(:batched_background_migration_job, :failed, batch_size: 1) }
 
       it 'raises an exception' do
         expect { job.split_and_retry! }.to raise_error 'Job cannot be split further'
@@ -204,7 +271,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
 
         expect(job.batch_size).to eq(5)
         expect(job.attempts).to eq(0)
-        expect(job.status).to eq('failed')
+        expect(job.status_name).to eq(:failed)
       end
     end
   end
diff --git a/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb b/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb
new file mode 100644
index 00000000000..c42a0fc5e05
--- /dev/null
+++ b/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJobTransitionLog, type: :model do
+  describe 'associations' do
+    it { is_expected.to belong_to(:batched_job).with_foreign_key(:batched_background_migration_job_id) }
+  end
+
+  describe 'validations' do
+    it { is_expected.to validate_presence_of(:previous_status) }
+    it { is_expected.to validate_presence_of(:next_status) }
+    it { is_expected.to validate_presence_of(:batched_job) }
+    it { is_expected.to validate_length_of(:exception_class).is_at_most(100) }
+    it { is_expected.to validate_length_of(:exception_message).is_at_most(1000) }
+    it { is_expected.to define_enum_for(:previous_status).with_values(%i(pending running failed succeeded)).with_prefix }
+    it { is_expected.to define_enum_for(:next_status).with_values(%i(pending running failed succeeded)).with_prefix }
+  end
+end
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
index 04c18a98ee6..bb2c6b9a3ae 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
@@ -96,13 +96,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       end
 
       let!(:previous_job) do
-        create(:batched_background_migration_job,
+        create(:batched_background_migration_job, :succeeded,
           batched_migration: migration,
           min_value: event1.id,
           max_value: event2.id,
           batch_size: 2,
-          sub_batch_size: 1,
-          status: :succeeded
+          sub_batch_size: 1
         )
       end
 
@@ -144,7 +143,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       context 'when migration has failed jobs' do
         before do
-          previous_job.update!(status: :failed)
+          previous_job.failure!
         end
 
         it 'retries the failed job' do
@@ -172,7 +171,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       context 'when migration has stuck jobs' do
         before do
-          previous_job.update!(status: :running, updated_at: 1.hour.ago - Gitlab::Database::BackgroundMigration::BatchedJob::STUCK_JOBS_TIMEOUT)
+          previous_job.update!(status_event: 'run', updated_at: 1.hour.ago - Gitlab::Database::BackgroundMigration::BatchedJob::STUCK_JOBS_TIMEOUT)
         end
 
         it 'retries the stuck job' do
@@ -186,7 +185,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       context 'when migration has possible stuck jobs' do
         before do
-          previous_job.update!(status: :running, updated_at: 1.hour.from_now - Gitlab::Database::BackgroundMigration::BatchedJob::STUCK_JOBS_TIMEOUT)
+          previous_job.update!(status_event: 'run', updated_at: 1.hour.from_now - Gitlab::Database::BackgroundMigration::BatchedJob::STUCK_JOBS_TIMEOUT)
         end
 
         it 'keeps the migration active' do
@@ -201,13 +200,13 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       context 'when the migration has batches to process and failed jobs' do
         before do
           migration.update!(max_value: event3.id)
-          previous_job.update!(status: :failed)
+          previous_job.failure!
         end
 
         it 'runs next batch then retries the failed job' do
           expect(migration_wrapper).to receive(:perform) do |job_record|
             expect(job_record).to eq(job_relation.last)
-            job_record.update!(status: :succeeded)
+            job_record.succeed!
           end
 
           expect { runner.run_migration_job(migration) }.to change { job_relation.count }.by(1)
@@ -264,12 +263,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       it 'runs all jobs inline until finishing the migration' do
         expect(migration_wrapper).to receive(:perform) do |job_record|
           expect(job_record).to eq(job_relation.first)
-          job_record.update!(status: :succeeded)
+          job_record.succeed!
         end
 
         expect(migration_wrapper).to receive(:perform) do |job_record|
           expect(job_record).to eq(job_relation.last)
-          job_record.update!(status: :succeeded)
+          job_record.succeed!
         end
 
         expect { runner.run_entire_migration(migration) }.to change { job_relation.count }.by(2)
@@ -330,9 +329,9 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
           pause_ms: 0
         }
 
-        create(:batched_background_migration_job, common_attributes.merge(status: :succeeded, min_value: 1, max_value: 2))
-        create(:batched_background_migration_job, common_attributes.merge(status: :pending, min_value: 3, max_value: 4))
-        create(:batched_background_migration_job, common_attributes.merge(status: :failed, min_value: 5, max_value: 6, attempts: 1))
+        create(:batched_background_migration_job, :succeeded, common_attributes.merge(min_value: 1, max_value: 2))
+        create(:batched_background_migration_job, :pending, common_attributes.merge(min_value: 3, max_value: 4))
+        create(:batched_background_migration_job, :failed, common_attributes.merge(min_value: 5, max_value: 6, attempts: 1))
       end
 
       it 'completes the migration' do
@@ -359,7 +358,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       context 'when migration fails to complete' do
         it 'raises an error' do
-          batched_migration.batched_jobs.failed.update_all(attempts: Gitlab::Database::BackgroundMigration::BatchedJob::MAX_ATTEMPTS)
+          batched_migration.batched_jobs.with_status(:failed).update_all(attempts: Gitlab::Database::BackgroundMigration::BatchedJob::MAX_ATTEMPTS)
 
           expect do
             runner.finalize(
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 01d61a525e6..ea4ba4dd137 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -26,7 +26,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     context 'when there are failed jobs' do
       let(:batched_migration) { create(:batched_background_migration, status: :active, total_tuple_count: 100) }
-      let!(:batched_job) { create(:batched_background_migration_job, batched_migration: batched_migration, status: :failed) }
+      let!(:batched_job) { create(:batched_background_migration_job, :failed, batched_migration: batched_migration) }
 
       it 'raises an exception' do
         expect { batched_migration.finished! }.to raise_error(ActiveRecord::RecordInvalid)
@@ -37,7 +37,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     context 'when the jobs are completed' do
       let(:batched_migration) { create(:batched_background_migration, status: :active, total_tuple_count: 100) }
-      let!(:batched_job) { create(:batched_background_migration_job, batched_migration: batched_migration, status: :succeeded) }
+      let!(:batched_job) { create(:batched_background_migration_job, :succeeded, batched_migration: batched_migration) }
 
       it 'finishes the migration' do
         batched_migration.finished!
@@ -64,7 +64,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     it 'returns the first active migration according to queue order' do
       expect(described_class.active_migration).to eq(migration2)
-      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
+      create(:batched_background_migration_job, :succeeded, batched_migration: migration1, batch_size: 1000)
     end
   end
 
@@ -84,10 +84,10 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     let!(:migration_without_jobs) { create(:batched_background_migration) }
 
     before do
-      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
-      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 200, status: :failed)
-      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 500, status: :succeeded)
-      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 200, status: :running)
+      create(:batched_background_migration_job, :succeeded, batched_migration: migration1, batch_size: 1000)
+      create(:batched_background_migration_job, :failed, batched_migration: migration1, batch_size: 200)
+      create(:batched_background_migration_job, :succeeded, batched_migration: migration2, batch_size: 500)
+      create(:batched_background_migration_job, :running, batched_migration: migration2, batch_size: 200)
     end
 
     it 'returns totals from successful jobs' do
@@ -268,7 +268,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     subject(:retry_failed_jobs) { batched_migration.retry_failed_jobs! }
 
     context 'when there are failed migration jobs' do
-      let!(:batched_background_migration_job) { create(:batched_background_migration_job, batched_migration: batched_migration, batch_size: 10, min_value: 6, max_value: 15, status: :failed, attempts: 3) }
+      let!(:batched_background_migration_job) { create(:batched_background_migration_job, :failed, batched_migration: batched_migration, batch_size: 10, min_value: 6, max_value: 15, attempts: 3) }
 
       before do
         allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
@@ -312,9 +312,9 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     let(:batched_migration) { create(:batched_background_migration) }
 
     before do
-      create_list(:batched_background_migration_job, 5, status: :succeeded, batch_size: 1_000, batched_migration: batched_migration)
-      create_list(:batched_background_migration_job, 1, status: :running, batch_size: 1_000, batched_migration: batched_migration)
-      create_list(:batched_background_migration_job, 1, status: :failed, batch_size: 1_000, batched_migration: batched_migration)
+      create_list(:batched_background_migration_job, 5, :succeeded, batch_size: 1_000, batched_migration: batched_migration)
+      create_list(:batched_background_migration_job, 1, :running, batch_size: 1_000, batched_migration: batched_migration)
+      create_list(:batched_background_migration_job, 1, :failed, batch_size: 1_000, batched_migration: batched_migration)
     end
 
     it 'sums the batch_size of succeeded jobs' do
@@ -347,7 +347,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     let_it_be(:common_attrs) do
       {
-        status: :succeeded,
         batched_migration: migration,
         finished_at: end_time
       }
@@ -357,7 +356,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
       subject { migration.smoothed_time_efficiency(number_of_jobs: 10) }
 
       it 'returns nil' do
-        create_list(:batched_background_migration_job, 9, **common_attrs)
+        create_list(:batched_background_migration_job, 9, :succeeded, **common_attrs)
 
         expect(subject).to be_nil
       end
@@ -369,6 +368,8 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
       subject { migration.smoothed_time_efficiency(number_of_jobs: number_of_jobs) }
 
+      let!(:jobs) { create_list(:batched_background_migration_job, number_of_jobs, :succeeded, **common_attrs.merge(batched_migration: migration)) }
+
       before do
         expect(migration).to receive_message_chain(:batched_jobs, :successful_in_execution_order, :reverse_order, :limit, :with_preloads)
                                .and_return(jobs)
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
index c1183a15e37..4f5536d8771 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
@@ -35,8 +35,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
     expect(job_instance).to receive(:perform)
     expect(job_instance).to receive(:batch_metrics).and_return(test_metrics)
 
-    expect(job_record).to receive(:update!).with(hash_including(attempts: 1, status: :running)).and_call_original
-
     freeze_time do
       subject
 
@@ -51,11 +49,10 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
 
   context 'when running a job that failed previously' do
     let!(:job_record) do
-      create(:batched_background_migration_job,
+      create(:batched_background_migration_job, :failed,
         batched_migration: active_migration,
         pause_ms: pause_ms,
         attempts: 1,
-        status: :failed,
         finished_at: 1.hour.ago,
         metrics: { 'my_metrics' => 'some_value' }
       )
@@ -67,10 +64,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
       expect(job_instance).to receive(:perform)
       expect(job_instance).to receive(:batch_metrics).and_return(updated_metrics)
 
-      expect(job_record).to receive(:update!).with(
-        hash_including(attempts: 2, status: :running, finished_at: nil, metrics: {})
-      ).and_call_original
-
       freeze_time do
         subject
 
@@ -201,4 +194,44 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
     it_behaves_like 'an error is raised', RuntimeError.new('Something broke!')
     it_behaves_like 'an error is raised', SignalException.new('SIGTERM')
   end
+
+  context 'when the batched background migration does not inherit from BaseJob' do
+    let(:migration_class) { Class.new }
+
+    before do
+      stub_const('Gitlab::BackgroundMigration::Foo', migration_class)
+    end
+
+    let(:connection) { double(:connection) }
+    let(:active_migration) { create(:batched_background_migration, :active, job_class_name: 'Foo') }
+    let!(:job_record) { create(:batched_background_migration_job, batched_migration: active_migration) }
+
+    it 'does not pass any argument' do
+      expect(Gitlab::BackgroundMigration::Foo).to receive(:new).with(no_args).and_return(job_instance)
+
+      expect(job_instance).to receive(:perform)
+
+      described_class.new(connection: connection).perform(job_record)
+    end
+  end
+
+  context 'when the batched background migration inherits from BaseJob' do
+    let(:connection) { double(:connection) }
+    let(:active_migration) { create(:batched_background_migration, :active, job_class_name: 'Foo') }
+    let!(:job_record) { create(:batched_background_migration_job, batched_migration: active_migration) }
+
+    let(:migration_class) { Class.new(::Gitlab::BackgroundMigration::BaseJob) }
+
+    before do
+      stub_const('Gitlab::BackgroundMigration::Foo', migration_class)
+    end
+
+    it 'passes the correct connection' do
+      expect(Gitlab::BackgroundMigration::Foo).to receive(:new).with(connection: connection).and_return(job_instance)
+
+      expect(job_instance).to receive(:perform)
+
+      described_class.new(connection: connection).perform(job_record)
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
index 0844616ee1c..31486240bfa 100644
--- a/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
+++ b/spec/lib/gitlab/database/dynamic_model_helpers_spec.rb
@@ -4,10 +4,11 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::DynamicModelHelpers do
   let(:including_class) { Class.new.include(described_class) }
-  let(:table_name) { 'projects' }
+  let(:table_name) { Project.table_name }
+  let(:connection) { Project.connection }
 
   describe '#define_batchable_model' do
-    subject { including_class.new.define_batchable_model(table_name) }
+    subject { including_class.new.define_batchable_model(table_name, connection: connection) }
 
     it 'is an ActiveRecord model' do
       expect(subject.ancestors).to include(ActiveRecord::Base)
@@ -40,7 +41,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
 
       it 'iterates table in batches' do
         each_batch_size = ->(&block) do
-          subject.each_batch(table_name, of: 1) do |batch|
+          subject.each_batch(table_name, connection: connection, of: 1) do |batch|
             block.call(batch.size)
           end
         end
@@ -56,7 +57,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'raises an error' do
-        expect { subject.each_batch(table_name, of: 1) { |batch| batch.size } }
+        expect { subject.each_batch(table_name, connection: connection, of: 1) { |batch| batch.size } }
           .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
       end
     end
@@ -74,7 +75,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'iterates table in batch ranges' do
-        expect { |b| subject.each_batch_range(table_name, of: 1, &b) }
+        expect { |b| subject.each_batch_range(table_name, connection: connection, of: 1, &b) }
           .to yield_successive_args(
             [first_project.id, first_project.id],
             [second_project.id, second_project.id]
@@ -82,13 +83,13 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'yields only one batch if bigger than the table size' do
-        expect { |b| subject.each_batch_range(table_name, of: 2, &b) }
+        expect { |b| subject.each_batch_range(table_name, connection: connection, of: 2, &b) }
           .to yield_successive_args([first_project.id, second_project.id])
       end
 
       it 'makes it possible to apply a scope' do
         each_batch_limited = ->(&b) do
-          subject.each_batch_range(table_name, scope: ->(table) { table.limit(1) }, of: 1, &b)
+          subject.each_batch_range(table_name, connection: connection, scope: ->(table) { table.limit(1) }, of: 1, &b)
         end
 
         expect { |b| each_batch_limited.call(&b) }
@@ -102,7 +103,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
       end
 
       it 'raises an error' do
-        expect { subject.each_batch_range(table_name, of: 1) { 1 } }
+        expect { subject.each_batch_range(table_name, connection: connection, of: 1) { 1 } }
           .to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
       end
     end
diff --git a/spec/lib/gitlab/database/each_database_spec.rb b/spec/lib/gitlab/database/each_database_spec.rb
index 9327fc4ff78..d526b3bc1ac 100644
--- a/spec/lib/gitlab/database/each_database_spec.rb
+++ b/spec/lib/gitlab/database/each_database_spec.rb
@@ -4,45 +4,97 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::EachDatabase do
   describe '.each_database_connection' do
-    let(:expected_connections) do
-      Gitlab::Database.database_base_models.map { |name, model| [model.connection, name] }
+    before do
+      allow(Gitlab::Database).to receive(:database_base_models)
+        .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
     end
 
-    it 'yields each connection after connecting SharedModel' do
-      expected_connections.each do |connection, _|
-        expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(connection).and_yield
-      end
+    it 'yields each connection after connecting SharedModel', :add_ci_connection do
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+        .with(ActiveRecord::Base.connection).ordered.and_yield
 
-      yielded_connections = []
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+        .with(Ci::ApplicationRecord.connection).ordered.and_yield
 
-      described_class.each_database_connection do |connection, name|
-        yielded_connections << [connection, name]
-      end
-
-      expect(yielded_connections).to match_array(expected_connections)
+      expect { |b| described_class.each_database_connection(&b) }
+        .to yield_successive_args(
+          [ActiveRecord::Base.connection, 'main'],
+          [Ci::ApplicationRecord.connection, 'ci']
+        )
     end
   end
 
   describe '.each_model_connection' do
-    let(:model1) { double(connection: double, table_name: 'table1') }
-    let(:model2) { double(connection: double, table_name: 'table2') }
+    context 'when the model inherits from SharedModel', :add_ci_connection do
+      let(:model1) { Class.new(Gitlab::Database::SharedModel) }
+      let(:model2) { Class.new(Gitlab::Database::SharedModel) }
 
-    before do
-      allow(model1.connection).to receive_message_chain('pool.db_config.name').and_return('name1')
-      allow(model2.connection).to receive_message_chain('pool.db_config.name').and_return('name2')
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models)
+          .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
+      end
+
+      it 'yields each model with SharedModel connected to each database connection' do
+        expect_yielded_models([model1, model2], [
+          { model: model1, connection: ActiveRecord::Base.connection, name: 'main' },
+          { model: model1, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+          { model: model2, connection: ActiveRecord::Base.connection, name: 'main' },
+          { model: model2, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+        ])
+      end
+
+      context 'when the model limits connection names' do
+        before do
+          model1.limit_connection_names = %i[main]
+          model2.limit_connection_names = %i[ci]
+        end
+
+        it 'only yields the model with SharedModel connected to the limited connections' do
+          expect_yielded_models([model1, model2], [
+            { model: model1, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: model2, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ])
+        end
+      end
+    end
+
+    context 'when the model does not inherit from SharedModel' do
+      let(:main_model) { Class.new(ActiveRecord::Base) }
+      let(:ci_model) { Class.new(Ci::ApplicationRecord) }
+
+      let(:main_connection) { double(:connection) }
+      let(:ci_connection) { double(:connection) }
+
+      before do
+        allow(main_model).to receive(:connection).and_return(main_connection)
+        allow(ci_model).to receive(:connection).and_return(ci_connection)
+
+        allow(main_connection).to receive_message_chain('pool.db_config.name').and_return('main')
+        allow(ci_connection).to receive_message_chain('pool.db_config.name').and_return('ci')
+      end
+
+      it 'yields each model after connecting SharedModel' do
+        expect_yielded_models([main_model, ci_model], [
+          { model: main_model, connection: main_connection, name: 'main' },
+          { model: ci_model, connection: ci_connection, name: 'ci' }
+        ])
+      end
     end
 
-    it 'yields each model after connecting SharedModel' do
-      expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model1.connection).and_yield
-      expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model2.connection).and_yield
+    def expect_yielded_models(models_to_iterate, expected_values)
+      times_yielded = 0
+
+      described_class.each_model_connection(models_to_iterate) do |model, name|
+        expected = expected_values[times_yielded]
 
-      yielded_models = []
+        expect(model).to be(expected[:model])
+        expect(model.connection).to be(expected[:connection])
+        expect(name).to eq(expected[:name])
 
-      described_class.each_model_connection([model1, model2]) do |model, name|
-        yielded_models << [model, name]
+        times_yielded += 1
       end
 
-      expect(yielded_models).to match_array([[model1, 'name1'], [model2, 'name2']])
+      expect(times_yielded).to eq(expected_values.size)
     end
   end
 end
diff --git a/spec/lib/gitlab/database/gitlab_schema_spec.rb b/spec/lib/gitlab/database/gitlab_schema_spec.rb
index 255efc99ff6..a5a67c2c918 100644
--- a/spec/lib/gitlab/database/gitlab_schema_spec.rb
+++ b/spec/lib/gitlab/database/gitlab_schema_spec.rb
@@ -44,6 +44,8 @@ RSpec.describe Gitlab::Database::GitlabSchema do
       'my_schema.ci_builds'             | :gitlab_ci
       'information_schema.columns'      | :gitlab_shared
       'audit_events_part_5fc467ac26'    | :gitlab_main
+      '_test_gitlab_main_table'         | :gitlab_main
+      '_test_gitlab_ci_table'           | :gitlab_ci
       '_test_my_table'                  | :gitlab_shared
       'pg_attribute'                    | :gitlab_shared
       'my_other_table'                  | :undefined_my_other_table
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index 796c14c1038..e87c9c20707 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -2,11 +2,18 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
+RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
   let(:configuration_hash) { {} }
   let(:db_config) { ActiveRecord::DatabaseConfigurations::HashConfig.new('test', 'ci', configuration_hash) }
   let(:model) { double(:model, connection_db_config: db_config) }
 
+  before do
+    # It's confusing to think about these specs with this enabled by default so
+    # we make it disabled by default and just write the specific spec for when
+    # it's enabled
+    stub_feature_flags(force_no_sharing_primary_model: false)
+  end
+
   describe '.for_model' do
     context 'when load balancing is not configured' do
       it 'uses the default settings' do
@@ -233,11 +240,23 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
     end
 
     context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
-      it 'the primary connection uses main connection' do
+      before do
         stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
+      end
 
+      it 'the primary connection uses main connection' do
         expect(config.primary_connection_specification_name).to eq('ActiveRecord::Base')
       end
+
+      context 'when force_no_sharing_primary_model feature flag is enabled' do
+        before do
+          stub_feature_flags(force_no_sharing_primary_model: true)
+        end
+
+        it 'the primary connection uses ci connection' do
+          expect(config.primary_connection_specification_name).to eq('Ci::ApplicationRecord')
+        end
+      end
     end
 
     context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=unknown' do
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
index 953d83d3b48..20519a759b2 100644
--- a/spec/lib/gitlab/database/load_balancing/setup_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
@@ -130,6 +130,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: false,
           ff_use_model_load_balancing: nil,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'ci_replica', write: 'ci' }
@@ -140,6 +141,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
           request_store_active: false,
           ff_use_model_load_balancing: nil,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'ci_replica', write: 'main' }
@@ -150,6 +152,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: false,
           ff_use_model_load_balancing: nil,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'main_replica', write: 'main' }
@@ -160,60 +163,77 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
           request_store_active: false,
           ff_use_model_load_balancing: nil,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'main_replica', write: 'main' }
           }
         },
-        "with FF disabled without RequestStore it uses main" => {
+        "with FF use_model_load_balancing disabled without RequestStore it uses main" => {
           env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: false,
           ff_use_model_load_balancing: false,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'main_replica', write: 'main' }
           }
         },
-        "with FF enabled without RequestStore sticking of FF does not work, so it fallbacks to use main" => {
+        "with FF use_model_load_balancing enabled without RequestStore sticking of FF does not work, so it fallbacks to use main" => {
           env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: false,
           ff_use_model_load_balancing: true,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'main_replica', write: 'main' }
           }
         },
-        "with FF disabled with RequestStore it uses main" => {
+        "with FF use_model_load_balancing disabled with RequestStore it uses main" => {
           env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: true,
           ff_use_model_load_balancing: false,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'main_replica', write: 'main' }
           }
         },
-        "with FF enabled with RequestStore it sticks FF and uses CI connection" => {
+        "with FF use_model_load_balancing enabled with RequestStore it sticks FF and uses CI connection" => {
           env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
           request_store_active: true,
           ff_use_model_load_balancing: true,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'ci_replica', write: 'ci' }
           }
         },
-        "with re-use and FF enabled with RequestStore it sticks FF and uses CI connection for reads" => {
+        "with re-use and ff_use_model_load_balancing enabled and FF force_no_sharing_primary_model disabled with RequestStore it sticks FF and uses CI connection for reads" => {
           env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
           env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
           request_store_active: true,
           ff_use_model_load_balancing: true,
+          ff_force_no_sharing_primary_model: false,
           expectations: {
             main: { read: 'main_replica', write: 'main' },
             ci: { read: 'ci_replica', write: 'main' }
           }
+        },
+        "with re-use and ff_use_model_load_balancing enabled and FF force_no_sharing_primary_model enabled with RequestStore it sticks FF and uses CI connection for reads" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
+          request_store_active: true,
+          ff_use_model_load_balancing: true,
+          ff_force_no_sharing_primary_model: true,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'ci_replica', write: 'ci' }
+          }
         }
       }
     end
@@ -243,6 +263,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
       around do |example|
         if request_store_active
           Gitlab::WithRequestStore.with_request_store do
+            stub_feature_flags(force_no_sharing_primary_model: ff_force_no_sharing_primary_model)
             RequestStore.clear!
 
             example.run
diff --git a/spec/lib/gitlab/database/loose_foreign_keys_spec.rb b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
index 13f2d31bc32..ed11699e494 100644
--- a/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
+++ b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
@@ -18,6 +18,45 @@ RSpec.describe Gitlab::Database::LooseForeignKeys do
                          ))
     end
 
+    context 'ensure keys are sorted' do
+      it 'does not have any keys that are out of order' do
+        parsed = YAML.parse_file(described_class.loose_foreign_keys_yaml_path)
+        mapping = parsed.children.first
+        table_names = mapping.children.select(&:scalar?).map(&:value)
+        expect(table_names).to eq(table_names.sort), "expected sorted table names in the YAML file"
+      end
+    end
+
+    context 'ensure no duplicates are found' do
+      it 'does not have duplicate tables defined' do
+        # since we use hash to detect duplicate hash keys we need to parse YAML document
+        parsed = YAML.parse_file(described_class.loose_foreign_keys_yaml_path)
+        expect(parsed).to be_document
+        expect(parsed.children).to be_one, "YAML has a single document"
+
+        # require hash
+        mapping = parsed.children.first
+        expect(mapping).to be_mapping, "YAML has a top-level hash"
+
+        # find all scalars with names
+        table_names = mapping.children.select(&:scalar?).map(&:value)
+        expect(table_names).not_to be_empty, "YAML has a non-zero tables defined"
+
+        # expect to not have duplicates
+        expect(table_names).to contain_exactly(*table_names.uniq)
+      end
+
+      it 'does not have duplicate column definitions' do
+        # ignore other modifiers
+        all_definitions = definitions.map do |definition|
+          { from_table: definition.from_table, to_table: definition.to_table, column: definition.column }
+        end
+
+        # expect to not have duplicates
+        expect(all_definitions).to contain_exactly(*all_definitions.uniq)
+      end
+    end
+
     describe 'ensuring database integrity' do
       def base_models_for(table)
         parent_table_schema = Gitlab::Database::GitlabSchema.table_schema(table)
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 7e3de32b965..d71a4f81901 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -14,6 +14,54 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     allow(model).to receive(:puts)
   end
 
+  describe 'overridden dynamic model helpers' do
+    let(:test_table) { '__test_batching_table' }
+
+    before do
+      model.connection.execute(<<~SQL)
+        CREATE TABLE #{test_table} (
+          id integer NOT NULL PRIMARY KEY,
+          name text NOT NULL
+        );
+
+        INSERT INTO #{test_table} (id, name)
+        VALUES (1, 'bob'), (2, 'mary'), (3, 'amy');
+      SQL
+    end
+
+    describe '#define_batchable_model' do
+      it 'defines a batchable model with the migration connection' do
+        expect(model.define_batchable_model(test_table).count).to eq(3)
+      end
+    end
+
+    describe '#each_batch' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        each_batch_name = ->(&block) do
+          model.each_batch(test_table, of: 2) do |batch|
+            block.call(batch.pluck(:name))
+          end
+        end
+
+        expect { |b| each_batch_name.call(&b) }.to yield_successive_args(%w[bob mary], %w[amy])
+      end
+    end
+
+    describe '#each_batch_range' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        expect { |b| model.each_batch_range(test_table, of: 2, &b) }.to yield_successive_args([1, 2], [3, 3])
+      end
+    end
+  end
+
   describe '#remove_timestamps' do
     it 'can remove the default timestamps' do
       Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
@@ -442,6 +490,60 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#remove_foreign_key_if_exists' do
+    context 'when the foreign key does not exist' do
+      before do
+        allow(model).to receive(:foreign_key_exists?).and_return(false)
+      end
+
+      it 'does nothing' do
+        expect(model).not_to receive(:remove_foreign_key)
+
+        model.remove_foreign_key_if_exists(:projects, :users, column: :user_id)
+      end
+    end
+
+    context 'when the foreign key exists' do
+      before do
+        allow(model).to receive(:foreign_key_exists?).and_return(true)
+      end
+
+      it 'removes the foreign key' do
+        expect(model).to receive(:remove_foreign_key).with(:projects, :users, { column: :user_id })
+
+        model.remove_foreign_key_if_exists(:projects, :users, column: :user_id)
+      end
+
+      context 'when the target table is not given' do
+        it 'passes the options as the second parameter' do
+          expect(model).to receive(:remove_foreign_key).with(:projects, { column: :user_id })
+
+          model.remove_foreign_key_if_exists(:projects, column: :user_id)
+        end
+      end
+
+      context 'when the reverse_lock_order option is given' do
+        it 'requests for lock before removing the foreign key' do
+          expect(model).to receive(:transaction_open?).and_return(true)
+          expect(model).to receive(:execute).with(/LOCK TABLE users, projects/)
+          expect(model).not_to receive(:remove_foreign_key).with(:projects, :users)
+
+          model.remove_foreign_key_if_exists(:projects, :users, column: :user_id, reverse_lock_order: true)
+        end
+
+        context 'when not inside a transaction' do
+          it 'does not lock' do
+            expect(model).to receive(:transaction_open?).and_return(false)
+            expect(model).not_to receive(:execute).with(/LOCK TABLE users, projects/)
+            expect(model).to receive(:remove_foreign_key).with(:projects, :users, { column: :user_id })
+
+            model.remove_foreign_key_if_exists(:projects, :users, column: :user_id, reverse_lock_order: true)
+          end
+        end
+      end
+    end
+  end
+
   describe '#add_concurrent_foreign_key' do
     before do
       allow(model).to receive(:foreign_key_exists?).and_return(false)
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 0abb76b9f8a..96dc3a0fc28 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -299,7 +299,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
 
       before do
         allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-          .with('main').and_return(coordinator)
+          .with(tracking_database).and_return(coordinator)
 
         expect(coordinator).to receive(:migration_class_for)
           .with(job_class_name).at_least(:once) { job_class }
@@ -403,7 +403,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       end
 
       context 'when a specific coordinator is given' do
-        let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.for_tracking_database('main') }
+        let(:coordinator) { Gitlab::BackgroundMigration::JobCoordinator.for_tracking_database(tracking_database) }
 
         it 'uses that coordinator' do
           expect(coordinator).to receive(:perform_in).with(10.minutes, 'Class', 'Hello', 'World').and_call_original
@@ -438,6 +438,16 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     it_behaves_like 'helpers that enqueue background migrations', BackgroundMigrationWorker, 'main'
   end
 
+  context 'when the migration is running against the ci database', if: Gitlab::Database.has_config?(:ci) do
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(::Ci::ApplicationRecord.connection) do
+        example.run
+      end
+    end
+
+    it_behaves_like 'helpers that enqueue background migrations', BackgroundMigration::CiDatabaseWorker, 'ci'
+  end
+
   describe '#delete_job_tracking' do
     let!(:job_class_name) { 'TestJob' }
 
diff --git a/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
index c45149d67bf..37efff165c7 100644
--- a/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
@@ -59,6 +59,7 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
           batch_max_value: 1000,
           batch_class_name: 'MyBatchClass',
           batch_size: 100,
+          max_batch_size: 10000,
           sub_batch_size: 10)
       end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
 
@@ -71,6 +72,7 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
         max_value: 1000,
         batch_class_name: 'MyBatchClass',
         batch_size: 100,
+        max_batch_size: 10000,
         sub_batch_size: 10,
         job_arguments: %w[],
         status: 'active',
diff --git a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
index 902d8e13a63..fd8303c379c 100644
--- a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
+++ b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
@@ -66,55 +66,43 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     context 'on successful execution' do
       subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
-      it 'records walltime' do
+      it 'records a valid observation', :aggregate_failures do
         expect(subject.walltime).not_to be_nil
-      end
-
-      it 'records success' do
         expect(subject.success).to be_truthy
-      end
-
-      it 'records the migration version' do
         expect(subject.version).to eq(migration_version)
-      end
-
-      it 'records the migration name' do
         expect(subject.name).to eq(migration_name)
       end
     end
 
     context 'upon failure' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' } }
-
-      it 'raises the exception' do
-        expect { subject }.to raise_error(/something went wrong/)
-      end
-
-      context 'retrieving observations' do
-        subject { instance.observations.first }
-
-        before do
-          instance.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' }
-        rescue StandardError
-          # ignore
-        end
+      where(exception: ['something went wrong', SystemStackError, Interrupt])
 
+      with_them do
         let(:instance) { described_class.new(result_dir: result_dir) }
 
-        it 'records walltime' do
-          expect(subject.walltime).not_to be_nil
-        end
-
-        it 'records failure' do
-          expect(subject.success).to be_falsey
-        end
+        subject(:observe) { instance.observe(version: migration_version, name: migration_name, connection: connection) { raise exception } }
 
-        it 'records the migration version' do
-          expect(subject.version).to eq(migration_version)
+        it 'raises the exception' do
+          expect { observe }.to raise_error(exception)
         end
 
-        it 'records the migration name' do
-          expect(subject.name).to eq(migration_name)
+        context 'retrieving observations' do
+          subject { instance.observations.first }
+
+          before do
+            observe
+            # rubocop:disable Lint/RescueException
+          rescue Exception
+            # rubocop:enable Lint/RescueException
+            # ignore (we expect this exception)
+          end
+
+          it 'records a valid observation', :aggregate_failures do
+            expect(subject.walltime).not_to be_nil
+            expect(subject.success).to be_falsey
+            expect(subject.version).to eq(migration_version)
+            expect(subject.name).to eq(migration_name)
+          end
         end
       end
     end
diff --git a/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb b/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
index 076fb9e8215..50ad77caaf1 100644
--- a/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
+++ b/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
@@ -3,7 +3,8 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries do
-    let(:migration) { double }
+    let(:connection) { ActiveRecord::Base.connection }
+    let(:migration) { double(connection: connection) }
     let(:return_value) { double }
     let(:class_def) do
       Class.new do
@@ -40,6 +41,18 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
         expect(result).to eq(return_value)
       end
     end
+
+    describe '#migration_connection' do
+      subject { class_def.new(migration).migration_connection }
+
+      it 'retrieves actual migration connection from #migration' do
+        expect(migration).to receive(:connection).and_return(return_value)
+
+        result = subject
+
+        expect(result).to eq(return_value)
+      end
+    end
   end
 
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries do
@@ -96,7 +109,8 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
 
     context 'with transactions enabled and lock retries enabled' do
       let(:receiver) { double('receiver', use_transaction?: true)}
-      let(:migration) { double('migration', enable_lock_retries?: true) }
+      let(:migration) { double('migration', migration_connection: connection, enable_lock_retries?: true) }
+      let(:connection) { ActiveRecord::Base.connection }
 
       it 'calls super method' do
         p = proc { }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
index 5a19ae6581d..a757cac0a2a 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
@@ -5,7 +5,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
   subject { described_class.new(observation, directory_path, connection) }
 
   let(:connection) { ActiveRecord::Migration.connection }
-  let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
+  let(:observation) { Gitlab::Database::Migrations::Observation.new(version: migration_version, name: migration_name) }
   let(:query) { "select date_trunc('day', $1::timestamptz) + $2 * (interval '1 hour')" }
   let(:query_binds) { [Time.current, 3] }
   let(:directory_path) { Dir.mktmpdir }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
index 7b01e39f5f1..eb66972e5ab 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
@@ -4,7 +4,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
   subject { described_class.new(observation, directory_path, connection) }
 
-  let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
+  let(:observation) { Gitlab::Database::Migrations::Observation.new(version: migration_version, name: migration_name) }
   let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select 1' }
   let(:directory_path) { Dir.mktmpdir }
diff --git a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
index b26bb8fbe41..f433e25b2ba 100644
--- a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
@@ -5,7 +5,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
   subject(:transaction_duration_observer) { described_class.new(observation, directory_path, connection) }
 
   let(:connection) { ActiveRecord::Migration.connection }
-  let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
+  let(:observation) { Gitlab::Database::Migrations::Observation.new(version: migration_version, name: migration_name) }
   let(:directory_path) { Dir.mktmpdir }
   let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
   let(:transaction_duration) { Gitlab::Json.parse(File.read(log_file)) }
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
index e5a8143fcc3..f94a40c93e1 100644
--- a/spec/lib/gitlab/database/no_cross_db_foreign_keys_spec.rb
+++ b/spec/lib/gitlab/database/no_cross_db_foreign_keys_spec.rb
@@ -3,59 +3,12 @@
 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 .
+  # Since we don't expect to have any cross-database foreign keys
+  # this is empty. If we will have an entry like
+  # `ci_daily_build_group_report_results.project_id`
+  # should be added.
   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
+    %w[].freeze
   end
 
   def foreign_keys_for(table_name)
diff --git a/spec/lib/gitlab/database/query_analyzers/prevent_cross_database_modification_spec.rb b/spec/lib/gitlab/database/query_analyzers/prevent_cross_database_modification_spec.rb
index c41b4eeea10..22a70dc7df0 100644
--- a/spec/lib/gitlab/database/query_analyzers/prevent_cross_database_modification_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzers/prevent_cross_database_modification_spec.rb
@@ -14,23 +14,41 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     Gitlab::Database::QueryAnalyzer.instance.within { example.run }
   end
 
-  shared_examples 'successful examples' do
+  describe 'context and suppress key names' do
+    describe '.context_key' do
+      it 'contains class name' do
+        expect(described_class.context_key)
+          .to eq 'analyzer_prevent_cross_database_modification_context'.to_sym
+      end
+    end
+
+    describe '.suppress_key' do
+      it 'contains class name' do
+        expect(described_class.suppress_key)
+          .to eq 'analyzer_prevent_cross_database_modification_suppressed'.to_sym
+      end
+    end
+  end
+
+  shared_examples 'successful examples' do |model:|
+    let(:model) { model }
+
     context 'outside transaction' do
       it { expect { run_queries }.not_to raise_error }
     end
 
-    context 'within transaction' do
+    context "within #{model} transaction" do
       it do
-        Project.transaction do
+        model.transaction do
           expect { run_queries }.not_to raise_error
         end
       end
     end
 
-    context 'within nested transaction' do
+    context "within nested #{model} transaction" do
       it do
-        Project.transaction(requires_new: true) do
-          Project.transaction(requires_new: true) do
+        model.transaction(requires_new: true) do
+          model.transaction(requires_new: true) do
             expect { run_queries }.not_to raise_error
           end
         end
@@ -38,13 +56,26 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
   end
 
+  shared_examples 'cross-database modification errors' do |model:|
+    let(:model) { model }
+
+    context "within #{model} transaction" do
+      it 'raises error' do
+        model.transaction do
+          expect { run_queries }.to raise_error /Cross-database data modification/
+        end
+      end
+    end
+  end
+
   context 'when CI and other tables are read in a transaction' do
     def run_queries
       pipeline.reload
       project.reload
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Project
+    include_examples 'successful examples', model: Ci::Pipeline
   end
 
   context 'when only CI data is modified' do
@@ -53,7 +84,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.reload
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Ci::Pipeline
+
+    include_examples 'cross-database modification errors', model: Project
   end
 
   context 'when other data is modified' do
@@ -62,7 +95,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.touch
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Project
+
+    include_examples 'cross-database modification errors', model: Ci::Pipeline
   end
 
   context 'when both CI and other data is modified' do
@@ -144,7 +179,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
         project.save!
       end
 
-      include_examples 'successful examples'
+      include_examples 'successful examples', model: Ci::Pipeline
+
+      include_examples 'cross-database modification errors', model: Project
     end
 
     describe '.allow_cross_database_modification_within_transaction' do
diff --git a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
index 0282a7af0df..6c32fb3ca17 100644
--- a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
@@ -5,7 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
   let(:env) { {} }
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
-  let(:subject) { described_class.new(env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(connection: connection, env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:connection) { ActiveRecord::Base.retrieve_connection }
 
   let(:timing_configuration) do
     [
@@ -67,7 +68,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             WHERE t.relkind = 'r' AND l.mode = 'ExclusiveLock' AND t.relname = '#{Project.table_name}'
           """
 
-          expect(ActiveRecord::Base.connection.execute(check_exclusive_lock_query).to_a).to be_present
+          expect(connection.execute(check_exclusive_lock_query).to_a).to be_present
         end
       end
 
@@ -96,8 +97,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
               lock_fiber.resume
             end
 
-            ActiveRecord::Base.transaction do
-              ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+            connection.transaction do
+              connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
               lock_acquired = true
             end
           end
@@ -115,7 +116,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
         context 'setting the idle transaction timeout' do
           context 'when there is no outer transaction: disable_ddl_transaction! is set in the migration' do
             it 'does not disable the idle transaction timeout' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+              allow(connection).to receive(:transaction_open?).and_return(false)
               allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
               allow(subject).to receive(:run_block_with_lock_timeout).once
 
@@ -127,7 +128,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
           context 'when there is outer transaction: disable_ddl_transaction! is not set in the migration' do
             it 'disables the idle transaction timeout so the code can sleep and retry' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+              allow(connection).to receive(:transaction_open?).and_return(true)
 
               n = 0
               allow(subject).to receive(:run_block_with_lock_timeout).twice do
@@ -184,8 +185,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             subject.run(raise_on_exhaustion: true) do
               lock_attempts += 1
 
-              ActiveRecord::Base.transaction do
-                ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+              connection.transaction do
+                connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
                 lock_acquired = true
               end
             end
@@ -199,11 +200,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
       context 'when statement timeout is reached' do
         it 'raises StatementInvalid error' do
           lock_acquired = false
-          ActiveRecord::Base.connection.execute("SET statement_timeout='100ms'")
+          connection.execute("SET statement_timeout='100ms'")
 
           expect do
             subject.run do
-              ActiveRecord::Base.connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
+              connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
               lock_acquired = true
             end
           end.to raise_error(ActiveRecord::StatementInvalid)
@@ -216,11 +217,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
   context 'restore local database variables' do
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW lock_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW lock_timeout").to_a }
     end
 
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
     end
   end
 
@@ -228,8 +229,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET lock_timeout` using the configured timeout value in milliseconds' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("SET lock_timeout TO '15ms'").and_call_original
+      expect(connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
+      expect(connection).to receive(:execute).with("SET lock_timeout TO '15ms'").and_call_original
 
       subject.run { }
     end
diff --git a/spec/lib/gitlab/database/with_lock_retries_spec.rb b/spec/lib/gitlab/database/with_lock_retries_spec.rb
index c2c818aa106..6b35ccafabc 100644
--- a/spec/lib/gitlab/database/with_lock_retries_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::WithLockRetries do
   let(:env) { {} }
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
-  let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(connection: connection, env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
   let(:allow_savepoints) { true }
   let(:connection) { ActiveRecord::Base.retrieve_connection }