Add latest changes from gitlab-org/gitlab@14-6-stable-eev14.6.0-rc42

34 files changed, 897 insertions, 375 deletions

diff --git a/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb b/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb
index b4010d0fe8d..7ad3eb395a9 100644
--- a/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb
+++ b/spec/lib/gitlab/database/async_indexes/index_creator_spec.rb
@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::AsyncIndexes::IndexCreator do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     subject { described_class.new(async_index) }
 
@@ -10,7 +12,18 @@ RSpec.describe Gitlab::Database::AsyncIndexes::IndexCreator do
 
     let(:index_model) { Gitlab::Database::AsyncIndexes::PostgresAsyncIndex }
 
-    let(:connection) { ApplicationRecord.connection }
+    let(:model) { Gitlab::Database.database_base_models[Gitlab::Database::PRIMARY_DATABASE_NAME] }
+    let(:connection) { model.connection }
+
+    let!(:lease) { stub_exclusive_lease(lease_key, :uuid, timeout: lease_timeout) }
+    let(:lease_key) { "gitlab/database/async_indexes/index_creator/#{Gitlab::Database::PRIMARY_DATABASE_NAME}" }
+    let(:lease_timeout) { described_class::TIMEOUT_PER_ACTION }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        example.run
+      end
+    end
 
     context 'when the index already exists' do
       before do
@@ -40,7 +53,7 @@ RSpec.describe Gitlab::Database::AsyncIndexes::IndexCreator do
     end
 
     it 'skips logic if not able to acquire exclusive lease' do
-      expect(subject).to receive(:try_obtain_lease).and_return(false)
+      expect(lease).to receive(:try_obtain).ordered.and_return(false)
       expect(connection).not_to receive(:execute).with(/CREATE INDEX/)
       expect(async_index).not_to receive(:destroy)
 
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 0182e0f7651..c4364826ee2 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
@@ -17,15 +17,19 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     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) }
 
-    before_all do
-      create(:batched_background_migration_job, status: :failed, attempts: described_class::MAX_ATTEMPTS)
-      create(:batched_background_migration_job, status: :succeeded)
-    end
+    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) }
 
     before do
       travel_to fixed_time
     end
 
+    describe '.except_succeeded' do
+      it 'returns not succeeded jobs' do
+        expect(described_class.except_succeeded).to contain_exactly(pending_job, running_job, stuck_job, failed_job, max_attempts_failed_job)
+      end
+    end
+
     describe '.active' do
       it 'returns active jobs' do
         expect(described_class.active).to contain_exactly(pending_job, running_job, stuck_job)
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 a1c2634f59c..49714cfc4dd 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -23,6 +23,28 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     subject { build(:batched_background_migration) }
 
     it { is_expected.to validate_uniqueness_of(:job_arguments).scoped_to(:job_class_name, :table_name, :column_name) }
+
+    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) }
+
+      it 'raises an exception' do
+        expect { batched_migration.finished! }.to raise_error(ActiveRecord::RecordInvalid)
+
+        expect(batched_migration.reload.status).to eql 'active'
+      end
+    end
+
+    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) }
+
+      it 'finishes the migration' do
+        batched_migration.finished!
+
+        expect(batched_migration.status).to eql 'finished'
+      end
+    end
   end
 
   describe '.queue_order' do
@@ -214,14 +236,20 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     end
   end
 
-  shared_examples_for 'an attr_writer that demodulizes assigned class names' do |attribute_name|
+  shared_examples_for 'an attr_writer that assigns class names' do |attribute_name|
     let(:batched_migration) { build(:batched_background_migration) }
 
     context 'when a module name exists' do
-      it 'removes the module name' do
+      it 'keeps the class with module name' do
+        batched_migration.public_send(:"#{attribute_name}=", 'Foo::Bar')
+
+        expect(batched_migration[attribute_name]).to eq('Foo::Bar')
+      end
+
+      it 'removes leading namespace resolution operator' do
         batched_migration.public_send(:"#{attribute_name}=", '::Foo::Bar')
 
-        expect(batched_migration[attribute_name]).to eq('Bar')
+        expect(batched_migration[attribute_name]).to eq('Foo::Bar')
       end
     end
 
@@ -271,11 +299,11 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
   end
 
   describe '#job_class_name=' do
-    it_behaves_like 'an attr_writer that demodulizes assigned class names', :job_class_name
+    it_behaves_like 'an attr_writer that assigns class names', :job_class_name
   end
 
   describe '#batch_class_name=' do
-    it_behaves_like 'an attr_writer that demodulizes assigned class names', :batch_class_name
+    it_behaves_like 'an attr_writer that assigns class names', :batch_class_name
   end
 
   describe '#migrated_tuple_count' do
diff --git a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
index 9d49db1f018..e7b9c5fcd02 100644
--- a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
+++ b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
@@ -5,24 +5,24 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
   before do
     create_list(:project, 3)
-    create(:identity)
+    create_list(:ci_instance_variable, 2)
   end
 
   subject { described_class.new(models).count }
 
   describe '#count' do
-    let(:models) { [Project, Identity] }
+    let(:models) { [Project, Ci::InstanceVariable] }
 
     context 'when reltuples is up to date' do
       before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
-        ActiveRecord::Base.connection.execute('ANALYZE identities')
+        Project.connection.execute('ANALYZE projects')
+        Ci::InstanceVariable.connection.execute('ANALYZE ci_instance_variables')
       end
 
       it 'uses statistics to do the count' do
         models.each { |model| expect(model).not_to receive(:count) }
 
-        expect(subject).to eq({ Project => 3, Identity => 1 })
+        expect(subject).to eq({ Project => 3, Ci::InstanceVariable => 2 })
       end
     end
 
@@ -31,7 +31,7 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
 
       before do
         models.each do |model|
-          ActiveRecord::Base.connection.execute("ANALYZE #{model.table_name}")
+          model.connection.execute("ANALYZE #{model.table_name}")
         end
       end
 
@@ -45,7 +45,9 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
 
     context 'insufficient permissions' do
       it 'returns an empty hash' do
-        allow(ActiveRecord::Base).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
+        Gitlab::Database.database_base_models.each_value do |base_model|
+          allow(base_model).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
+        end
 
         expect(subject).to eq({})
       end
diff --git a/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb b/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
index 2f261aebf02..37d3e13a7ab 100644
--- a/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
+++ b/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
@@ -5,11 +5,12 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
   before do
     create_list(:project, 3)
+    create_list(:ci_instance_variable, 2)
     create(:identity)
     create(:group)
   end
 
-  let(:models) { [Project, Identity, Group, Namespace] }
+  let(:models) { [Project, Ci::InstanceVariable, Identity, Group, Namespace] }
   let(:strategy) { described_class.new(models) }
 
   subject { strategy.count }
@@ -20,7 +21,8 @@ RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
         Project => threshold + 1,
         Identity => threshold - 1,
         Group => threshold + 1,
-        Namespace => threshold + 1
+        Namespace => threshold + 1,
+        Ci::InstanceVariable => threshold + 1
       }
     end
 
@@ -43,12 +45,14 @@ RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
         expect(Project).not_to receive(:count)
         expect(Group).not_to receive(:count)
         expect(Namespace).not_to receive(:count)
+        expect(Ci::InstanceVariable).not_to receive(:count)
 
         result = subject
         expect(result[Project]).to eq(3)
         expect(result[Group]).to eq(1)
         # 1-Group, 3 namespaces for each project and 3 project namespaces for each project
         expect(result[Namespace]).to eq(7)
+        expect(result[Ci::InstanceVariable]).to eq(2)
       end
     end
 
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index eef248afdf2..796c14c1038 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -140,6 +140,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
   end
 
   describe '#service_discovery_enabled?' do
+    it 'returns false when running inside a Rake task' do
+      allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
+
+      config = described_class.new(ActiveRecord::Base)
+      config.service_discovery[:record] = 'foo'
+
+      expect(config.service_discovery_enabled?).to eq(false)
+    end
+
     it 'returns true when a record is configured' do
       config = described_class.new(ActiveRecord::Base)
       config.service_discovery[:record] = 'foo'
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index 37b83729125..3c7819c04b6 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -487,25 +487,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
   end
 
-  describe 'primary connection re-use', :reestablished_active_record_base do
+  describe 'primary connection re-use', :reestablished_active_record_base, :add_ci_connection do
     let(:model) { Ci::ApplicationRecord }
 
-    around do |example|
-      if Gitlab::Database.has_config?(:ci)
-        example.run
-      else
-        # fake additional Database
-        model.establish_connection(
-          ActiveRecord::DatabaseConfigurations::HashConfig.new(Rails.env, 'ci', ActiveRecord::Base.connection_db_config.configuration_hash)
-        )
-
-        example.run
-
-        # Cleanup connection_specification_name for Ci::ApplicationRecord
-        model.remove_connection
-      end
-    end
-
     describe '#read' do
       it 'returns ci replica connection' do
         expect { |b| lb.read(&b) }.to yield_with_args do |args|
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
index e9bc465b1c7..f05910e5123 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -4,9 +4,10 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   let(:load_balancer) do
-    Gitlab::Database::LoadBalancing::LoadBalancer.new(
-      Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
-    )
+    configuration = Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+    configuration.service_discovery[:record] = 'localhost'
+
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(configuration)
   end
 
   let(:service) do
@@ -86,6 +87,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         service.perform_service_discovery
       end
     end
+
     context 'with failures' do
       before do
         allow(Gitlab::ErrorTracking).to receive(:track_exception)
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
index de2ad662d16..31be3963565 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -5,7 +5,9 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_gitlab_redis_queues do
   let(:middleware) { described_class.new }
   let(:worker) { worker_class.new }
-  let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } }
+  let(:location) {'0/D525E3A8' }
+  let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
+  let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
 
   before do
     skip_feature_flags_yaml_validation
@@ -60,9 +62,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
     end
 
     shared_examples_for 'replica is up to date' do |expected_strategy|
-      let(:location) {'0/D525E3A8' }
-      let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
-
       it 'does not stick to the primary', :aggregate_failures do
         expect(ActiveRecord::Base.load_balancer)
           .to receive(:select_up_to_date_host)
@@ -77,9 +76,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
       include_examples 'load balancing strategy', expected_strategy
     end
 
-    shared_examples_for 'sticks based on data consistency' do |data_consistency|
-      include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker
-
+    shared_examples_for 'sticks based on data consistency' do
       context 'when load_balancing_for_test_data_consistency_worker is disabled' do
         before do
           stub_feature_flags(load_balancing_for_test_data_consistency_worker: false)
@@ -116,23 +113,78 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         it_behaves_like 'replica is up to date', 'replica'
       end
 
-      context 'when legacy wal location is set' do
-        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
+      context 'when database location is not set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
 
-        before do
-          allow(ActiveRecord::Base.load_balancer)
-            .to receive(:select_up_to_date_host)
-            .with('0/D525E3A8')
-            .and_return(true)
-        end
+        include_examples 'stick to the primary', 'primary_no_wal'
+      end
+    end
 
-        it_behaves_like 'replica is up to date', 'replica'
+    shared_examples_for 'sleeps when necessary' do
+      context 'when WAL locations are blank', :freeze_time do
+        let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", "wal_locations" => {}, "created_at" => Time.current.to_f - (described_class::MINIMUM_DELAY_INTERVAL_SECONDS - 0.3) } }
+
+        it 'does not sleep' do
+          expect(middleware).not_to receive(:sleep)
+
+          run_middleware
+        end
       end
 
-      context 'when database location is not set' do
-        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
+      context 'when WAL locations are present', :freeze_time do
+        let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations, "created_at" => Time.current.to_f - elapsed_time } }
 
-        include_examples 'stick to the primary', 'primary_no_wal'
+        context 'when delay interval has not elapsed' do
+          let(:elapsed_time) { described_class::MINIMUM_DELAY_INTERVAL_SECONDS - 0.3 }
+
+          context 'when replica is up to date' do
+            before do
+              Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+                allow(lb).to receive(:select_up_to_date_host).and_return(true)
+              end
+            end
+
+            it 'does not sleep' do
+              expect(middleware).not_to receive(:sleep)
+
+              run_middleware
+            end
+          end
+
+          context 'when replica is not up to date' do
+            before do
+              Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+                allow(lb).to receive(:select_up_to_date_host).and_return(false, true)
+              end
+            end
+
+            it 'sleeps until the minimum delay is reached' do
+              expect(middleware).to receive(:sleep).with(be_within(0.01).of(described_class::MINIMUM_DELAY_INTERVAL_SECONDS - elapsed_time))
+
+              run_middleware
+            end
+          end
+        end
+
+        context 'when delay interval has elapsed' do
+          let(:elapsed_time) { described_class::MINIMUM_DELAY_INTERVAL_SECONDS + 0.3 }
+
+          it 'does not sleep' do
+            expect(middleware).not_to receive(:sleep)
+
+            run_middleware
+          end
+        end
+
+        context 'when created_at is in the future' do
+          let(:elapsed_time) { -5 }
+
+          it 'does not sleep' do
+            expect(middleware).not_to receive(:sleep)
+
+            run_middleware
+          end
+        end
       end
     end
 
@@ -146,10 +198,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
 
       include_examples 'stick to the primary', 'primary'
+
+      context 'when delay interval has not elapsed', :freeze_time do
+        let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations, "created_at" => Time.current.to_f - elapsed_time } }
+        let(:elapsed_time) { described_class::MINIMUM_DELAY_INTERVAL_SECONDS - 0.3 }
+
+        it 'does not sleep' do
+          expect(middleware).not_to receive(:sleep)
+
+          run_middleware
+        end
+      end
     end
 
     context 'when worker data consistency is :delayed' do
-      include_examples 'sticks based on data consistency', :delayed
+      include_context 'data consistency worker class', :delayed, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'sticks based on data consistency'
+      include_examples 'sleeps when necessary'
 
       context 'when replica is not up to date' do
         before do
@@ -177,7 +243,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         end
 
         context 'when job is retried' do
-          let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8', 'retry_count' => 0 } }
+          let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations, 'retry_count' => 0 } }
 
           context 'and replica still lagging behind' do
             include_examples 'stick to the primary', 'primary'
@@ -195,7 +261,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
     end
 
     context 'when worker data consistency is :sticky' do
-      include_examples 'sticks based on data consistency', :sticky
+      include_context 'data consistency worker class', :sticky, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'sticks based on data consistency'
+      include_examples 'sleeps when necessary'
 
       context 'when replica is not up to date' do
         before do
@@ -255,7 +324,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
   end
 
   def run_middleware
-    middleware.call(worker, job, double(:queue)) { yield }
+    middleware.call(worker, job, double(:queue)) { yield if block_given? }
   rescue described_class::JobReplicaNotUpToDate
     # we silence errors here that cause the job to retry
   end
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index d88554614cf..f3139bb1b4f 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -256,15 +256,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
       expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
-
-    it 'removes the old key' do
-      Gitlab::Redis::SharedState.with do |redis|
-        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
-      end
-
-      sticking.unstick(:user, 4)
-      expect(sticking.last_write_location_for(:user, 4)).to be_nil
-    end
   end
 
   describe '#last_write_location_for' do
@@ -273,14 +264,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
       expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     end
-
-    it 'falls back to reading the old key' do
-      Gitlab::Redis::SharedState.with do |redis|
-        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
-      end
-
-      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
-    end
   end
 
   describe '#redis_key_for' do
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index 65ffe539910..45878b2e266 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -38,6 +38,24 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
+  describe '.primary_only?' do
+    it 'returns true if all load balancers have no replicas' do
+      described_class.each_load_balancer do |lb|
+        allow(lb).to receive(:primary_only?).and_return(true)
+      end
+
+      expect(described_class.primary_only?).to eq(true)
+    end
+
+    it 'returns false if at least one has replicas' do
+      described_class.each_load_balancer.with_index do |lb, index|
+        allow(lb).to receive(:primary_only?).and_return(index != 0)
+      end
+
+      expect(described_class.primary_only?).to eq(false)
+    end
+  end
+
   describe '.release_hosts' do
     it 'releases the host of every load balancer' do
       described_class.each_load_balancer do |lb|
diff --git a/spec/lib/gitlab/database/loose_foreign_keys_spec.rb b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
new file mode 100644
index 00000000000..13f2d31bc32
--- /dev/null
+++ b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LooseForeignKeys do
+  describe 'verify all definitions' do
+    subject(:definitions) { described_class.definitions }
+
+    it 'all definitions have assigned a known gitlab_schema and on_delete' do
+      is_expected.to all(have_attributes(
+                           options: a_hash_including(
+                             column: be_a(String),
+                             gitlab_schema: be_in(Gitlab::Database.schemas_to_base_models.symbolize_keys.keys),
+                             on_delete: be_in([:async_delete, :async_nullify])
+                           ),
+                           from_table: be_a(String),
+                           to_table: be_a(String)
+                         ))
+    end
+
+    describe 'ensuring database integrity' do
+      def base_models_for(table)
+        parent_table_schema = Gitlab::Database::GitlabSchema.table_schema(table)
+        Gitlab::Database.schemas_to_base_models.fetch(parent_table_schema)
+      end
+
+      it 'all `to_table` tables are present' do
+        definitions.each do |definition|
+          base_models_for(definition.to_table).each do |model|
+            expect(model.connection).to be_table_exist(definition.to_table)
+          end
+        end
+      end
+
+      it 'all `from_table` tables are present' do
+        definitions.each do |definition|
+          base_models_for(definition.from_table).each do |model|
+            expect(model.connection).to be_table_exist(definition.from_table)
+            expect(model.connection).to be_column_exist(definition.from_table, definition.column)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migration_helpers/loose_foreign_key_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers/loose_foreign_key_helpers_spec.rb
index f1dbfbbff18..25fc676d09e 100644
--- a/spec/lib/gitlab/database/migration_helpers/loose_foreign_key_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers/loose_foreign_key_helpers_spec.rb
@@ -47,11 +47,16 @@ RSpec.describe Gitlab::Database::MigrationHelpers::LooseForeignKeyHelpers do
       record_to_be_deleted.delete
 
       expect(LooseForeignKeys::DeletedRecord.count).to eq(1)
-      deleted_record = LooseForeignKeys::DeletedRecord.all.first
+
+      arel_table = LooseForeignKeys::DeletedRecord.arel_table
+      deleted_record = LooseForeignKeys::DeletedRecord
+        .select(arel_table[Arel.star], arel_table[:partition].as('partition_number')) # aliasing the ignored partition column to partition_number
+        .all
+        .first
 
       expect(deleted_record.primary_key_value).to eq(record_to_be_deleted.id)
       expect(deleted_record.fully_qualified_table_name).to eq('public._test_loose_fk_test_table')
-      expect(deleted_record.partition).to eq(1)
+      expect(deleted_record.partition_number).to eq(1)
     end
 
     it 'stores multiple record deletions' do
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index ea755f5a368..7f80bed04a4 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -2431,7 +2431,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     let(:issues)         { table(:issues) }
 
     def setup
-      namespace = namespaces.create!(name: 'foo', path: 'foo')
+      namespace = namespaces.create!(name: 'foo', path: 'foo', type: Namespaces::UserNamespace.sti_name)
       projects.create!(namespace_id: namespace.id)
     end
 
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 e42a6c970ea..99c7d70724c 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -7,78 +7,6 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     ActiveRecord::Migration.new.extend(described_class)
   end
 
-  describe '#bulk_queue_background_migration_jobs_by_range' 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 }
-
-      before do
-        User.class_eval do
-          include EachBatch
-        end
-      end
-
-      context 'with enough rows to bulk queue jobs more than once' do
-        before do
-          stub_const('Gitlab::Database::Migrations::BackgroundMigrationHelpers::JOB_BUFFER_SIZE', 1)
-        end
-
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          end
-        end
-
-        it 'queues jobs in groups of buffer size 1' do
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id1, id2]]])
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id3, id3]]])
-
-          model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-        end
-      end
-
-      context 'with not enough rows to bulk queue jobs more than once' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          end
-        end
-
-        it 'queues jobs in bulk all at once (big buffer size)' do
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id1, id2]],
-                                                                                  ['FooJob', [id3, id3]]])
-
-          model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-        end
-      end
-
-      context 'without specifying batch_size' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob')
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
-          end
-        end
-      end
-    end
-
-    context "when the model doesn't have an ID column" do
-      it 'raises error (for now)' do
-        expect do
-          model.bulk_queue_background_migration_jobs_by_range(ProjectAuthorization, 'FooJob')
-        end.to raise_error(StandardError, /does not have an ID/)
-      end
-    end
-  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 }
@@ -354,161 +282,6 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     end
   end
 
-  describe '#queue_batched_background_migration' do
-    let(:pgclass_info) { instance_double('Gitlab::Database::PgClass', cardinality_estimate: 42) }
-
-    before do
-      allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
-    end
-
-    context 'when such migration already exists' do
-      it 'does not create duplicate migration' do
-        create(
-          :batched_background_migration,
-          job_class_name: 'MyJobClass',
-          table_name: :projects,
-          column_name: :id,
-          interval: 10.minutes,
-          min_value: 5,
-          max_value: 1005,
-          batch_class_name: 'MyBatchClass',
-          batch_size: 200,
-          sub_batch_size: 20,
-          job_arguments: [[:id], [:id_convert_to_bigint]]
-        )
-
-        expect do
-          model.queue_batched_background_migration(
-            'MyJobClass',
-            :projects,
-            :id,
-            [:id], [:id_convert_to_bigint],
-            job_interval: 5.minutes,
-            batch_min_value: 5,
-            batch_max_value: 1000,
-            batch_class_name: 'MyBatchClass',
-            batch_size: 100,
-            sub_batch_size: 10)
-        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
-      end
-    end
-
-    it 'creates the database record for the migration' do
-      expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info)
-
-      expect do
-        model.queue_batched_background_migration(
-          'MyJobClass',
-          :projects,
-          :id,
-          job_interval: 5.minutes,
-          batch_min_value: 5,
-          batch_max_value: 1000,
-          batch_class_name: 'MyBatchClass',
-          batch_size: 100,
-          sub_batch_size: 10)
-      end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-      expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
-        job_class_name: 'MyJobClass',
-        table_name: 'projects',
-        column_name: 'id',
-        interval: 300,
-        min_value: 5,
-        max_value: 1000,
-        batch_class_name: 'MyBatchClass',
-        batch_size: 100,
-        sub_batch_size: 10,
-        job_arguments: %w[],
-        status: 'active',
-        total_tuple_count: pgclass_info.cardinality_estimate)
-    end
-
-    context 'when the job interval is lower than the minimum' do
-      let(:minimum_delay) { described_class::BATCH_MIN_DELAY }
-
-      it 'sets the job interval to the minimum value' do
-        expect do
-          model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: minimum_delay - 1.minute)
-        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-        created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-        expect(created_migration.interval).to eq(minimum_delay)
-      end
-    end
-
-    context 'when additional arguments are passed to the method' do
-      it 'saves the arguments on the database record' do
-        expect do
-          model.queue_batched_background_migration(
-            'MyJobClass',
-            :projects,
-            :id,
-            'my',
-            'arguments',
-            job_interval: 5.minutes,
-            batch_max_value: 1000)
-        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-        expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
-          job_class_name: 'MyJobClass',
-          table_name: 'projects',
-          column_name: 'id',
-          interval: 300,
-          min_value: 1,
-          max_value: 1000,
-          job_arguments: %w[my arguments])
-      end
-    end
-
-    context 'when the max_value is not given' do
-      context 'when records exist in the database' do
-        let!(:event1) { create(:event) }
-        let!(:event2) { create(:event) }
-        let!(:event3) { create(:event) }
-
-        it 'creates the record with the current max value' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-          expect(created_migration.max_value).to eq(event3.id)
-        end
-
-        it 'creates the record with an active status' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_active
-        end
-      end
-
-      context 'when the database is empty' do
-        it 'sets the max value to the min value' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-          expect(created_migration.max_value).to eq(created_migration.min_value)
-        end
-
-        it 'creates the record with a finished status' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :projects, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_finished
-        end
-      end
-    end
-  end
-
   describe '#migrate_async' do
     it 'calls BackgroundMigrationWorker.perform_async' do
       expect(BackgroundMigrationWorker).to receive(:perform_async).with("Class", "hello", "world")
@@ -583,7 +356,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
   end
 
   describe '#finalized_background_migration' do
-    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(:main, BackgroundMigrationWorker) }
+    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(BackgroundMigrationWorker) }
 
     let!(:job_class_name) { 'TestJob' }
     let!(:job_class) { Class.new }
@@ -605,7 +378,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       job_class.define_method(:perform, job_perform_method)
 
       allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-        .with(:main).and_return(job_coordinator)
+        .with('main').and_return(job_coordinator)
 
       expect(job_coordinator).to receive(:migration_class_for)
         .with(job_class_name).at_least(:once) { job_class }
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
new file mode 100644
index 00000000000..c45149d67bf
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers do
+  let(:migration) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  describe '#queue_batched_background_migration' do
+    let(:pgclass_info) { instance_double('Gitlab::Database::PgClass', cardinality_estimate: 42) }
+
+    before do
+      allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
+    end
+
+    context 'when such migration already exists' do
+      it 'does not create duplicate migration' do
+        create(
+          :batched_background_migration,
+          job_class_name: 'MyJobClass',
+          table_name: :projects,
+          column_name: :id,
+          interval: 10.minutes,
+          min_value: 5,
+          max_value: 1005,
+          batch_class_name: 'MyBatchClass',
+          batch_size: 200,
+          sub_batch_size: 20,
+          job_arguments: [[:id], [:id_convert_to_bigint]]
+        )
+
+        expect do
+          migration.queue_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            [:id], [:id_convert_to_bigint],
+            job_interval: 5.minutes,
+            batch_min_value: 5,
+            batch_max_value: 1000,
+            batch_class_name: 'MyBatchClass',
+            batch_size: 100,
+            sub_batch_size: 10)
+        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
+      end
+    end
+
+    it 'creates the database record for the migration' do
+      expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info)
+
+      expect do
+        migration.queue_batched_background_migration(
+          'MyJobClass',
+          :projects,
+          :id,
+          job_interval: 5.minutes,
+          batch_min_value: 5,
+          batch_max_value: 1000,
+          batch_class_name: 'MyBatchClass',
+          batch_size: 100,
+          sub_batch_size: 10)
+      end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+      expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
+        job_class_name: 'MyJobClass',
+        table_name: 'projects',
+        column_name: 'id',
+        interval: 300,
+        min_value: 5,
+        max_value: 1000,
+        batch_class_name: 'MyBatchClass',
+        batch_size: 100,
+        sub_batch_size: 10,
+        job_arguments: %w[],
+        status: 'active',
+        total_tuple_count: pgclass_info.cardinality_estimate)
+    end
+
+    context 'when the job interval is lower than the minimum' do
+      let(:minimum_delay) { described_class::BATCH_MIN_DELAY }
+
+      it 'sets the job interval to the minimum value' do
+        expect do
+          migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: minimum_delay - 1.minute)
+        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+        created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+        expect(created_migration.interval).to eq(minimum_delay)
+      end
+    end
+
+    context 'when additional arguments are passed to the method' do
+      it 'saves the arguments on the database record' do
+        expect do
+          migration.queue_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            'my',
+            'arguments',
+            job_interval: 5.minutes,
+            batch_max_value: 1000)
+        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
+          job_class_name: 'MyJobClass',
+          table_name: 'projects',
+          column_name: 'id',
+          interval: 300,
+          min_value: 1,
+          max_value: 1000,
+          job_arguments: %w[my arguments])
+      end
+    end
+
+    context 'when the max_value is not given' do
+      context 'when records exist in the database' do
+        let!(:event1) { create(:event) }
+        let!(:event2) { create(:event) }
+        let!(:event3) { create(:event) }
+
+        it 'creates the record with the current max value' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+          expect(created_migration.max_value).to eq(event3.id)
+        end
+
+        it 'creates the record with an active status' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_active
+        end
+      end
+
+      context 'when the database is empty' do
+        it 'sets the max value to the min value' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+          expect(created_migration.max_value).to eq(created_migration.min_value)
+        end
+
+        it 'creates the record with a finished status' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :projects, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_finished
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
index 841d2a98a16..902d8e13a63 100644
--- a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
+++ b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
@@ -3,6 +3,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Instrumentation do
   let(:result_dir) { Dir.mktmpdir }
+  let(:connection) { ActiveRecord::Migration.connection }
 
   after do
     FileUtils.rm_rf(result_dir)
@@ -14,11 +15,11 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     let(:migration_version) { '12345' }
 
     it 'executes the given block' do
-      expect { |b| subject.observe(version: migration_version, name: migration_name, &b) }.to yield_control
+      expect { |b| subject.observe(version: migration_version, name: migration_name, connection: connection, &b) }.to yield_control
     end
 
     context 'behavior with observers' do
-      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
       let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
 
@@ -29,7 +30,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       it 'instantiates observer with observation' do
         expect(Gitlab::Database::Migrations::Observers::MigrationObserver)
           .to receive(:new)
-          .with(instance_of(Gitlab::Database::Migrations::Observation), anything) { |observation| expect(observation.version).to eq(migration_version) }
+          .with(instance_of(Gitlab::Database::Migrations::Observation), anything, connection) { |observation| expect(observation.version).to eq(migration_version) }
           .and_return(observer)
 
         subject
@@ -63,7 +64,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'on successful execution' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
       it 'records walltime' do
         expect(subject.walltime).not_to be_nil
@@ -83,7 +84,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'upon failure' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) { raise 'something went wrong' } }
+      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/)
@@ -93,7 +94,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
         subject { instance.observations.first }
 
         before do
-          instance.observe(version: migration_version, name: migration_name) { raise 'something went wrong' }
+          instance.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' }
         rescue StandardError
           # ignore
         end
@@ -125,8 +126,8 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       let(:migration2) { double('migration2', call: nil) }
 
       it 'records observations for all migrations' do
-        subject.observe(version: migration_version, name: migration_name) {}
-        subject.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } rescue nil
+        subject.observe(version: migration_version, name: migration_name, connection: connection) {}
+        subject.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' } rescue nil
 
         expect(subject.observations.size).to eq(2)
       end
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 191ac29e3b3..5a19ae6581d 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
-  subject { described_class.new(observation, directory_path) }
+  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(:connection) { ActiveRecord::Base.connection }
   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 2e70a85fd5b..7b01e39f5f1 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
-  subject { described_class.new(observation, directory_path) }
+  subject { described_class.new(observation, directory_path, connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select 1' }
   let(:directory_path) { Dir.mktmpdir }
   let(:migration_version) { 20210422152437 }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
index 9727a215d71..2515f0d4a06 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
-  subject { described_class.new(observation, double("unused path")) }
+  subject { described_class.new(observation, double("unused path"), connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
 
   def mock_pgss(enabled: true)
     if enabled
diff --git a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
index e689759c574..4b08838d6bb 100644
--- a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
-  subject { described_class.new(observation, double('unused path')) }
+  subject { described_class.new(observation, double('unused path'), connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select pg_database_size(current_database())' }
 
   it 'records the size change' do
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 e65f89747c4..b26bb8fbe41 100644
--- a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
@@ -2,8 +2,9 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
-  subject(:transaction_duration_observer) { described_class.new(observation, directory_path) }
+  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(:directory_path) { Dir.mktmpdir }
   let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
@@ -78,17 +79,17 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
   end
 
   def run_real_transactions
-    ActiveRecord::Base.transaction do
+    ApplicationRecord.transaction do
     end
   end
 
   def run_sub_transactions
-    ActiveRecord::Base.transaction(requires_new: true) do
+    ApplicationRecord.transaction(requires_new: true) do
     end
   end
 
   def run_transaction
-    ActiveRecord::Base.connection_pool.with_connection do |connection|
+    ApplicationRecord.connection_pool.with_connection do |connection|
       Gitlab::Database::SharedModel.using_connection(connection) do
         Gitlab::Database::SharedModel.transaction do
           Gitlab::Database::SharedModel.transaction(requires_new: true) do
diff --git a/spec/lib/gitlab/database/migrations/runner_spec.rb b/spec/lib/gitlab/database/migrations/runner_spec.rb
index 52fb5ec2ba8..4616bd6941e 100644
--- a/spec/lib/gitlab/database/migrations/runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/runner_spec.rb
@@ -76,7 +76,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       it 'runs the unapplied migrations in version order', :aggregate_failures do
         up.run
 
-        expect(migration_runs.map(&:dir)).to eq([:up, :up])
+        expect(migration_runs.map(&:dir)).to match_array([:up, :up])
         expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
       end
     end
@@ -101,7 +101,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       it 'runs the applied migrations for the current branch in reverse order', :aggregate_failures do
         down.run
 
-        expect(migration_runs.map(&:dir)).to eq([:down, :down])
+        expect(migration_runs.map(&:dir)).to match_array([:down, :down])
         expect(migration_runs.map(&:version_to_migrate)).to eq(applied_migrations_this_branch.reverse.map(&:version))
       end
     end
diff --git a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
index b2c4e4b54a4..2ef873e8adb 100644
--- a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
@@ -90,18 +90,6 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
         expect(table_oid('test_partition')).to be_nil
       end
 
-      context 'when the drop_detached_partitions feature flag is disabled' do
-        before do
-          stub_feature_flags(drop_detached_partitions: false)
-        end
-
-        it 'does not drop the partition' do
-          dropper.perform
-
-          expect(table_oid('test_partition')).not_to be_nil
-        end
-      end
-
       context 'removing foreign keys' do
         it 'removes foreign keys from the table before dropping it' do
           expect(dropper).to receive(:drop_detached_partition).and_wrap_original do |drop_method, partition_name|
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index 1c6f5c5c694..5e107109fc9 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -16,7 +16,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     subject(:sync_partitions) { described_class.new(model).sync_partitions }
 
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
-    let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: []) }
+    let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: [], after_adding_partitions: nil) }
     let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "some_table" }
 
@@ -83,7 +83,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
     let(:manager) { described_class.new(model) }
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
-    let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: []) }
+    let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: [], after_adding_partitions: nil) }
     let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "foo" }
 
@@ -101,28 +101,10 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       ]
     end
 
-    context 'with the partition_pruning feature flag enabled' do
-      before do
-        stub_feature_flags(partition_pruning: true)
-      end
-
-      it 'detaches each extra partition' do
-        extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) }
-
-        sync_partitions
-      end
-    end
+    it 'detaches each extra partition' do
+      extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) }
 
-    context 'with the partition_pruning feature flag disabled' do
-      before do
-        stub_feature_flags(partition_pruning: false)
-      end
-
-      it 'returns immediately' do
-        expect(manager).not_to receive(:detach)
-
-        sync_partitions
-      end
+      sync_partitions
     end
   end
 
diff --git a/spec/lib/gitlab/database/partitioning/single_numeric_list_partition_spec.rb b/spec/lib/gitlab/database/partitioning/single_numeric_list_partition_spec.rb
new file mode 100644
index 00000000000..9941241e846
--- /dev/null
+++ b/spec/lib/gitlab/database/partitioning/single_numeric_list_partition_spec.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::SingleNumericListPartition do
+  describe '.from_sql' do
+    subject(:parsed_partition) { described_class.from_sql(table, partition_name, definition) }
+
+    let(:table) { 'partitioned_table' }
+    let(:partition_value) { 0 }
+    let(:partition_name) { "partitioned_table_#{partition_value}" }
+    let(:definition) { "FOR VALUES IN ('#{partition_value}')" }
+
+    it 'uses specified table name' do
+      expect(parsed_partition.table).to eq(table)
+    end
+
+    it 'uses specified partition name' do
+      expect(parsed_partition.partition_name).to eq(partition_name)
+    end
+
+    it 'parses the definition' do
+      expect(parsed_partition.value).to eq(partition_value)
+    end
+  end
+
+  describe '#partition_name' do
+    it 'is the explicit name if provided' do
+      expect(described_class.new('table', 1, partition_name: 'some_other_name').partition_name).to eq('some_other_name')
+    end
+
+    it 'defaults to the table name followed by the partition value' do
+      expect(described_class.new('table', 1).partition_name).to eq('table_1')
+    end
+  end
+
+  context 'sorting' do
+    it 'is incomparable if the tables do not match' do
+      expect(described_class.new('table1', 1) <=> described_class.new('table2', 2)).to be_nil
+    end
+
+    it 'sorts by the value when the tables match' do
+      expect(described_class.new('table1', 1) <=> described_class.new('table1', 2)).to eq(1 <=> 2)
+    end
+
+    it 'sorts by numeric value rather than text value' do
+      expect(described_class.new('table', 10)).to be > described_class.new('table', 9)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
new file mode 100644
index 00000000000..636a09e5710
--- /dev/null
+++ b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
@@ -0,0 +1,214 @@
+# frozen_string_literal: true
+
+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(:next_partition_if) { double('next_partition_if') }
+  let(:detach_partition_if) { double('detach_partition_if') }
+
+  subject(:strategy) do
+    described_class.new(model, :partition,
+                        next_partition_if: next_partition_if,
+                        detach_partition_if: detach_partition_if)
+  end
+
+  before do
+    connection.execute(<<~SQL)
+      create table #{table_name}
+        (
+          id serial not null,
+          partition bigint not null default 2,
+          created_at timestamptz not null,
+          primary key (id, partition)
+        )
+        partition by list(partition);
+
+      create table #{table_name}_1
+      partition of #{table_name} for values in (1);
+
+      create table #{table_name}_2
+      partition of #{table_name} for values in (2);
+    SQL
+  end
+
+  describe '#current_partitions' do
+    it 'detects both partitions' do
+      expect(strategy.current_partitions).to eq([
+        Gitlab::Database::Partitioning::SingleNumericListPartition.new(table_name, 1, partition_name: '_test_partitioned_test_1'),
+        Gitlab::Database::Partitioning::SingleNumericListPartition.new(table_name, 2, partition_name: '_test_partitioned_test_2')
+      ])
+    end
+  end
+
+  describe '#active_partition' do
+    it 'is the partition with the largest value' do
+      expect(strategy.active_partition.value).to eq(2)
+    end
+  end
+
+  describe '#missing_partitions' do
+    context 'when next_partition_if returns true' do
+      let(:next_partition_if) { proc { true } }
+
+      it 'is a partition definition for the next partition in the series' do
+        extra = strategy.missing_partitions
+
+        expect(extra.length).to eq(1)
+        expect(extra.first.value).to eq(3)
+      end
+    end
+
+    context 'when next_partition_if returns false' do
+      let(:next_partition_if) { proc { false } }
+
+      it 'is empty' do
+        expect(strategy.missing_partitions).to be_empty
+      end
+    end
+
+    context 'when there are no partitions for the table' do
+      it 'returns a partition for value 1' do
+        connection.execute("drop table #{table_name}_1; drop table #{table_name}_2;")
+
+        missing_partitions = strategy.missing_partitions
+
+        expect(missing_partitions.size).to eq(1)
+        missing_partition = missing_partitions.first
+
+        expect(missing_partition.value).to eq(1)
+      end
+    end
+  end
+
+  describe '#extra_partitions' do
+    before do
+      (3..10).each do |i|
+        connection.execute("CREATE TABLE #{table_name}_#{i} PARTITION OF #{table_name} FOR VALUES IN (#{i})")
+      end
+    end
+
+    context 'when some partitions are true for detach_partition_if' 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)
+      end
+    end
+
+    context 'when all partitions are true for detach_partition_if' 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.current_partitions.map(&:value).max).to eq(10)
+      end
+    end
+  end
+
+  describe '#initial_partition' do
+    it 'starts with the value 1', :aggregate_failures do
+      initial_partition = strategy.initial_partition
+      expect(initial_partition.value).to eq(1)
+      expect(initial_partition.table).to eq(strategy.table_name)
+      expect(initial_partition.partition_name).to eq("#{strategy.table_name}_1")
+    end
+  end
+
+  describe '#next_partition' do
+    it 'is one after the active partition', :aggregate_failures do
+      expect(strategy).to receive(:active_partition).and_return(double(value: 5))
+      next_partition = strategy.next_partition
+
+      expect(next_partition.value).to eq(6)
+      expect(next_partition.table).to eq(strategy.table_name)
+      expect(next_partition.partition_name).to eq("#{strategy.table_name}_6")
+    end
+  end
+
+  describe '#ensure_partitioning_column_ignored!' do
+    it 'raises when the column is not ignored' do
+      expect do
+        Class.new(ApplicationRecord) do
+          include PartitionedTable
+
+          partitioned_by :partition, strategy: :sliding_list,
+                       next_partition_if: proc { false },
+                       detach_partition_if: proc { false }
+        end
+      end.to raise_error(/ignored_columns/)
+    end
+
+    it 'does not raise when the column is ignored' do
+      expect do
+        Class.new(ApplicationRecord) do
+          include PartitionedTable
+
+          self.ignored_columns = [:partition]
+
+          partitioned_by :partition, strategy: :sliding_list,
+                         next_partition_if: proc { false },
+                         detach_partition_if: proc { false }
+        end
+      end.not_to raise_error
+    end
+  end
+  context 'redirecting inserts as the active partition changes' do
+    let(:model) do
+      Class.new(ApplicationRecord) do
+        include PartitionedTable
+
+        self.table_name = '_test_partitioned_test'
+        self.primary_key = :id
+
+        self.ignored_columns = %w[partition]
+
+        # method().call cannot be detected by rspec, so we add a layer of indirection here
+        def self.next_partition_if_wrapper(...)
+          next_partition?(...)
+        end
+
+        def self.detach_partition_if_wrapper(...)
+          detach_partition?(...)
+        end
+        partitioned_by :partition, strategy: :sliding_list,
+                       next_partition_if: method(:next_partition_if_wrapper),
+                       detach_partition_if: method(:detach_partition_if_wrapper)
+
+        def self.next_partition?(current_partition)
+        end
+
+        def self.detach_partition?(partition)
+        end
+      end
+    end
+
+    it 'redirects to the new partition', :aggregate_failures do
+      partition_2_model = model.create! # Goes in partition 2
+
+      allow(model).to receive(:next_partition?) do
+        model.partitioning_strategy.active_partition.value < 3
+      end
+
+      allow(model).to receive(:detach_partition?).and_return(false)
+
+      Gitlab::Database::Partitioning::PartitionManager.new(model).sync_partitions
+
+      partition_3_model = model.create!
+
+      # Rails doesn't pick up on database default changes, so we need to reload
+      # We also want to grab the partition column to verify what it was set to.
+      # In normal operation we make rails ignore it so that we can use a changing default
+      # So we force select * to load it
+      all_columns = model.select(model.arel_table[Arel.star])
+      partition_2_model = all_columns.find(partition_2_model.id)
+      partition_3_model = all_columns.find(partition_3_model.id)
+
+      expect(partition_2_model.partition).to eq(2)
+      expect(partition_3_model.partition).to eq(3)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/query_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzer_spec.rb
index 82a1c7143d5..34c72893c53 100644
--- a/spec/lib/gitlab/database/query_analyzer_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzer_spec.rb
@@ -128,11 +128,20 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
 
     it 'does not call analyze on suppressed analyzers' do
       expect(analyzer).to receive(:suppressed?).and_return(true)
+      expect(analyzer).to receive(:requires_tracking?).and_return(false)
       expect(analyzer).not_to receive(:analyze)
 
       expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
     end
 
+    it 'does call analyze on suppressed analyzers if some queries require tracking' do
+      expect(analyzer).to receive(:suppressed?).and_return(true)
+      expect(analyzer).to receive(:requires_tracking?).and_return(true)
+      expect(analyzer).to receive(:analyze)
+
+      expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
+    end
+
     def process_sql(sql)
       described_class.instance.within do
         ApplicationRecord.load_balancer.read_write do |connection|
diff --git a/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb b/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
index ab5f05e3ec4..86e74cf5177 100644
--- a/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
@@ -17,7 +17,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
     process_sql(ActiveRecord::Base, "SELECT 1 FROM projects")
   end
 
-  context 'properly observes all queries', :mocked_ci_connection do
+  context 'properly observes all queries', :add_ci_connection do
     using RSpec::Parameterized::TableSyntax
 
     where do
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 eb8ccb0bd89..c41b4eeea10 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
@@ -92,6 +92,23 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
           end
         end
       end
+
+      context 'when comments are added to the front of query strings' do
+        around do |example|
+          prepend_comment_was = Marginalia::Comment.prepend_comment
+          Marginalia::Comment.prepend_comment = true
+
+          example.run
+
+          Marginalia::Comment.prepend_comment = prepend_comment_was
+        end
+
+        it 'raises error' do
+          Project.transaction do
+            expect { run_queries }.to raise_error /Cross-database data modification/
+          end
+        end
+      end
     end
 
     context 'when executing a SELECT FOR UPDATE query' do
@@ -164,4 +181,49 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       end.to raise_error /Cross-database data modification.*The gitlab_schema was undefined/
     end
   end
+
+  context 'when execution is rescued with StandardError' do
+    it 'raises cross-database data modification exception' do
+      expect do
+        Project.transaction do
+          project.touch
+          project.connection.execute('UPDATE foo_bars_undefined_table SET a=1 WHERE id = -1')
+        end
+      rescue StandardError
+        # Ensures that standard rescue does not silence errors
+      end.to raise_error /Cross-database data modification.*The gitlab_schema was undefined/
+    end
+  end
+
+  context 'when uniquiness validation is tested', type: :model do
+    subject { build(:ci_variable) }
+
+    it 'does not raise exceptions' do
+      expect do
+        is_expected.to validate_uniqueness_of(:key).scoped_to(:project_id, :environment_scope).with_message(/\(\w+\) has already been taken/)
+      end.not_to raise_error
+    end
+  end
+
+  context 'when doing rollback in a suppressed block' do
+    it 'does not raise misaligned transactions exception' do
+      expect do
+        # This is non-materialised transaction:
+        # 1. the transaction will be open on a write (project.touch) (in a suppressed block)
+        # 2. the rescue will be handled outside of suppressed block
+        #
+        # This will create misaligned boundaries since BEGIN
+        # of transaction will be executed within a suppressed block
+        Project.transaction do
+          described_class.with_suppressed do
+            project.touch
+
+            raise 'force rollback'
+          end
+
+          # the ensure of `.transaction` executes `ROLLBACK TO SAVEPOINT`
+        end
+      end.to raise_error /force rollback/
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/reindexing/coordinator_spec.rb b/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
index 085fd3061ad..0afbe46b7f1 100644
--- a/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
+++ b/spec/lib/gitlab/database/reindexing/coordinator_spec.rb
@@ -15,10 +15,18 @@ RSpec.describe Gitlab::Database::Reindexing::Coordinator do
     let(:action) { create(:reindex_action, index: index) }
 
     let!(:lease) { stub_exclusive_lease(lease_key, uuid, timeout: lease_timeout) }
-    let(:lease_key) { 'gitlab/database/reindexing/coordinator' }
+    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]
+
+      Gitlab::Database::SharedModel.using_connection(model.connection) do
+        example.run
+      end
+    end
+
     before do
       swapout_view_for_table(:postgres_indexes)
 
diff --git a/spec/lib/gitlab/database/reindexing_spec.rb b/spec/lib/gitlab/database/reindexing_spec.rb
index 13aff343432..0c576505e07 100644
--- a/spec/lib/gitlab/database/reindexing_spec.rb
+++ b/spec/lib/gitlab/database/reindexing_spec.rb
@@ -6,6 +6,63 @@ RSpec.describe Gitlab::Database::Reindexing do
   include ExclusiveLeaseHelpers
   include Database::DatabaseHelpers
 
+  describe '.invoke' do
+    let(:databases) { Gitlab::Database.database_base_models }
+    let(:databases_count) { databases.count }
+
+    it 'cleans up any leftover indexes' do
+      expect(described_class).to receive(:cleanup_leftovers!).exactly(databases_count).times
+
+      described_class.invoke
+    end
+
+    context 'when there is an error raised' do
+      it 'logs and re-raise' do
+        expect(described_class).to receive(:automatic_reindexing).and_raise('Unexpected!')
+        expect(Gitlab::AppLogger).to receive(:error)
+
+        expect { described_class.invoke }.to raise_error('Unexpected!')
+      end
+    end
+
+    context 'when async index creation is enabled' do
+      it 'executes async index creation prior to any reindexing actions' do
+        stub_feature_flags(database_async_index_creation: true)
+
+        expect(Gitlab::Database::AsyncIndexes).to receive(:create_pending_indexes!).ordered.exactly(databases_count).times
+        expect(described_class).to receive(:automatic_reindexing).ordered.exactly(databases_count).times
+
+        described_class.invoke
+      end
+    end
+
+    context 'when async index creation is disabled' do
+      it 'does not execute async index creation' do
+        stub_feature_flags(database_async_index_creation: false)
+
+        expect(Gitlab::Database::AsyncIndexes).not_to receive(:create_pending_indexes!)
+
+        described_class.invoke
+      end
+    end
+
+    context 'calls automatic reindexing' do
+      it 'uses all candidate indexes' do
+        expect(described_class).to receive(:automatic_reindexing).exactly(databases_count).times
+
+        described_class.invoke
+      end
+
+      context 'when explicit database is given' do
+        it 'skips other databases' do
+          expect(described_class).to receive(:automatic_reindexing).once
+
+          described_class.invoke(Gitlab::Database::PRIMARY_DATABASE_NAME)
+        end
+      end
+    end
+  end
+
   describe '.automatic_reindexing' do
     subject { described_class.automatic_reindexing(maximum_records: limit) }
 
@@ -133,10 +190,19 @@ RSpec.describe Gitlab::Database::Reindexing do
   end
 
   describe '.cleanup_leftovers!' do
-    subject { described_class.cleanup_leftovers! }
+    subject(:cleanup_leftovers) { described_class.cleanup_leftovers! }
+
+    let(:model) { Gitlab::Database.database_base_models[Gitlab::Database::PRIMARY_DATABASE_NAME] }
+    let(:connection) { model.connection }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        example.run
+      end
+    end
 
     before do
-      ApplicationRecord.connection.execute(<<~SQL)
+      connection.execute(<<~SQL)
         CREATE INDEX foobar_ccnew ON users (id);
         CREATE INDEX foobar_ccnew1 ON users (id);
       SQL
@@ -150,11 +216,11 @@ RSpec.describe Gitlab::Database::Reindexing do
       expect_query("DROP INDEX CONCURRENTLY IF EXISTS \"public\".\"foobar_ccnew1\"")
       expect_query("RESET idle_in_transaction_session_timeout; RESET lock_timeout")
 
-      subject
+      cleanup_leftovers
     end
 
     def expect_query(sql)
-      expect(ApplicationRecord.connection).to receive(:execute).ordered.with(sql).and_wrap_original do |method, sql|
+      expect(connection).to receive(:execute).ordered.with(sql).and_wrap_original do |method, sql|
         method.call(sql.sub(/CONCURRENTLY/, ''))
       end
     end
diff --git a/spec/lib/gitlab/database/shared_model_spec.rb b/spec/lib/gitlab/database/shared_model_spec.rb
index 94f2b5a3434..54af4a0c4dc 100644
--- a/spec/lib/gitlab/database/shared_model_spec.rb
+++ b/spec/lib/gitlab/database/shared_model_spec.rb
@@ -84,4 +84,16 @@ RSpec.describe Gitlab::Database::SharedModel do
       expect(described_class.connection).to be(original_connection)
     end
   end
+
+  describe '#connection_db_config' do
+    it 'returns the class connection_db_config' do
+      shared_model_class = Class.new(described_class) do
+        self.table_name = 'postgres_async_indexes'
+      end
+
+      shared_model = shared_model_class.new
+
+      expect(shared_model.connection_db_config). to eq(described_class.connection_db_config)
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/type/json_pg_safe_spec.rb b/spec/lib/gitlab/database/type/json_pg_safe_spec.rb
new file mode 100644
index 00000000000..91dc6f39aa7
--- /dev/null
+++ b/spec/lib/gitlab/database/type/json_pg_safe_spec.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Type::JsonPgSafe do
+  let(:type) { described_class.new }
+
+  describe '#serialize' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { type.serialize(value) }
+
+    where(:value, :json) do
+      nil                            | nil
+      1                              | '1'
+      1.0                            | '1.0'
+      "str\0ing\u0000"               | '"string"'
+      ["\0arr", "a\u0000y"]          | '["arr","ay"]'
+      { "key\0" => "value\u0000\0" } | '{"key":"value"}'
+    end
+
+    with_them do
+      it { is_expected.to eq(json) }
+    end
+  end
+end