Add latest changes from gitlab-org/gitlab@15-1-stable-eev15.1.0-rc42

21 files changed, 704 insertions, 346 deletions

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 f147e8204e6..97459d4a7be 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
@@ -311,6 +311,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
     let(:table_name) { :_test_batched_migrations_test_table }
     let(:column_name) { :some_id }
     let(:job_arguments) { [:some_id, :some_id_convert_to_bigint] }
+    let(:gitlab_schemas) { Gitlab::Database.gitlab_schemas_for_connection(connection) }
 
     let(:migration_status) { :active }
 
@@ -358,7 +359,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       it 'completes the migration' do
         expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
-          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .with(gitlab_schemas, 'CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
           .and_return(batched_migration)
 
         expect(batched_migration).to receive(:finalize!).and_call_original
@@ -399,7 +400,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
       it 'is a no-op' do
         expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
-          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .with(gitlab_schemas, 'CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
           .and_return(batched_migration)
 
         configuration = {
@@ -426,7 +427,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
     context 'when the migration does not exist' do
       it 'is a no-op' do
         expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
-          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, [:some, :other, :arguments])
+          .with(gitlab_schemas, 'CopyColumnUsingBackgroundMigrationJob', table_name, column_name, [:some, :other, :arguments])
           .and_return(nil)
 
         configuration = {
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 a1c979bba50..8819171cfd0 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -78,23 +78,41 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
   end
 
   describe '.active_migration' do
+    let(:connection) { Gitlab::Database.database_base_models[:main].connection }
     let!(:migration1) { create(:batched_background_migration, :finished) }
 
-    context 'without migrations on hold' do
+    subject(:active_migration) { described_class.active_migration(connection: connection) }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        example.run
+      end
+    end
+
+    context 'when there are no migrations on hold' do
       let!(:migration2) { create(:batched_background_migration, :active) }
       let!(:migration3) { create(:batched_background_migration, :active) }
 
       it 'returns the first active migration according to queue order' do
-        expect(described_class.active_migration).to eq(migration2)
+        expect(active_migration).to eq(migration2)
       end
     end
 
-    context 'with migrations are on hold' do
+    context 'when there are migrations on hold' do
       let!(:migration2) { create(:batched_background_migration, :active, on_hold_until: 10.minutes.from_now) }
       let!(:migration3) { create(:batched_background_migration, :active, on_hold_until: 2.minutes.ago) }
 
       it 'returns the first active migration that is not on hold according to queue order' do
-        expect(described_class.active_migration).to eq(migration3)
+        expect(active_migration).to eq(migration3)
+      end
+    end
+
+    context 'when there are migrations not available for the current connection' do
+      let!(:migration2) { create(:batched_background_migration, :active, gitlab_schema: :gitlab_not_existing) }
+      let!(:migration3) { create(:batched_background_migration, :active, gitlab_schema: :gitlab_main) }
+
+      it 'returns the first active migration that is available for the current connection' do
+        expect(active_migration).to eq(migration3)
       end
     end
   end
@@ -553,25 +571,43 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
   end
 
   describe '.for_configuration' do
-    let!(:migration) do
-      create(
-        :batched_background_migration,
+    let!(:attributes) do
+      {
         job_class_name: 'MyJobClass',
         table_name: :projects,
         column_name: :id,
-        job_arguments: [[:id], [:id_convert_to_bigint]]
-      )
+        job_arguments: [[:id], [:id_convert_to_bigint]],
+        gitlab_schema: :gitlab_main
+      }
     end
 
+    let!(:migration) { create(:batched_background_migration, attributes) }
+
     before do
-      create(:batched_background_migration, job_class_name: 'OtherClass')
-      create(:batched_background_migration, table_name: 'other_table')
-      create(:batched_background_migration, column_name: 'other_column')
-      create(:batched_background_migration, job_arguments: %w[other arguments])
+      create(:batched_background_migration, attributes.merge(job_class_name: 'OtherClass'))
+      create(:batched_background_migration, attributes.merge(table_name: 'other_table'))
+      create(:batched_background_migration, attributes.merge(column_name: 'other_column'))
+      create(:batched_background_migration, attributes.merge(job_arguments: %w[other arguments]))
     end
 
     it 'finds the migration matching the given configuration parameters' do
-      actual = described_class.for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
+      actual = described_class.for_configuration(:gitlab_main, 'MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
+
+      expect(actual).to contain_exactly(migration)
+    end
+
+    it 'filters by gitlab schemas available for the connection' do
+      actual = described_class.for_configuration(:gitlab_ci, 'MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
+
+      expect(actual).to be_empty
+    end
+
+    it 'doesn not filter by gitlab schemas available for the connection if the column is nor present' do
+      skip_if_multiple_databases_not_setup
+
+      expect(described_class).to receive(:gitlab_schema_column_exists?).and_return(false)
+
+      actual = described_class.for_configuration(:gitlab_main, 'MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
 
       expect(actual).to contain_exactly(migration)
     end
@@ -579,7 +615,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
   describe '.find_for_configuration' do
     it 'returns nill if such migration does not exists' do
-      expect(described_class.find_for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to be_nil
+      expect(described_class.find_for_configuration(:gitlab_main, 'MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to be_nil
     end
 
     it 'returns the migration when it exists' do
@@ -588,10 +624,25 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
         job_class_name: 'MyJobClass',
         table_name: :projects,
         column_name: :id,
-        job_arguments: [[:id], [:id_convert_to_bigint]]
+        job_arguments: [[:id], [:id_convert_to_bigint]],
+        gitlab_schema: :gitlab_main
       )
 
-      expect(described_class.find_for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to eq(migration)
+      expect(described_class.find_for_configuration(:gitlab_main, 'MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to eq(migration)
+    end
+  end
+
+  describe '.for_gitlab_schema' do
+    let!(:migration) { create(:batched_background_migration, gitlab_schema: :gitlab_main) }
+
+    before do
+      create(:batched_background_migration, gitlab_schema: :gitlab_not_existing)
+    end
+
+    it 'finds the migrations matching the given gitlab schema' do
+      actual = described_class.for_gitlab_schema(:gitlab_main)
+
+      expect(actual).to contain_exactly(migration)
     end
   end
 end
diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index 028bdce852e..811d4fad95c 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -384,4 +384,58 @@ RSpec.describe Gitlab::Database::BatchCount do
       subject { described_class.method(:batch_sum) }
     end
   end
+
+  describe '#batch_average' do
+    let(:model) { Issue }
+    let(:column) { :weight }
+
+    before do
+      Issue.update_all(weight: 2)
+    end
+
+    it 'returns the average of values in the given column' do
+      expect(described_class.batch_average(model, column)).to eq(2)
+    end
+
+    it 'works when given an Arel column' do
+      expect(described_class.batch_average(model, model.arel_table[column])).to eq(2)
+    end
+
+    it 'works with a batch size of 50K' do
+      expect(described_class.batch_average(model, column, batch_size: 50_000)).to eq(2)
+    end
+
+    it 'works with start and finish provided' do
+      expect(described_class.batch_average(model, column, start: model.minimum(:id), finish: model.maximum(:id))).to eq(2)
+    end
+
+    it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE}" do
+      min_id = model.minimum(:id)
+      relation = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE)
+
+      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
+
+      described_class.batch_average(model, column)
+    end
+
+    it_behaves_like 'when a transaction is open' do
+      subject { described_class.batch_average(model, column) }
+    end
+
+    it_behaves_like 'disallowed configurations', :batch_average do
+      let(:args) { [model, column] }
+      let(:default_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE }
+      let(:small_batch_size)   { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE - 1 }
+    end
+
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :itself }
+      let(:operation) { :average }
+      let(:operation_args) { [column] }
+
+      subject { described_class.method(:batch_average) }
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/each_database_spec.rb b/spec/lib/gitlab/database/each_database_spec.rb
index 191f7017b4c..8345cdfb8fb 100644
--- a/spec/lib/gitlab/database/each_database_spec.rb
+++ b/spec/lib/gitlab/database/each_database_spec.rb
@@ -61,7 +61,11 @@ RSpec.describe Gitlab::Database::EachDatabase do
 
     context 'when shared connections are not included' do
       it 'only yields the unshared connections' do
-        expect(Gitlab::Database).to receive(:db_config_share_with).twice.and_return(nil, 'main')
+        if Gitlab::Database.has_config?(:ci)
+          expect(Gitlab::Database).to receive(:db_config_share_with).exactly(3).times.and_return(nil, 'main', 'main')
+        else
+          expect(Gitlab::Database).to receive(:db_config_share_with).twice.and_return(nil, 'main')
+        end
 
         expect { |b| described_class.each_database_connection(include_shared: false, &b) }
           .to yield_successive_args([ActiveRecord::Base.connection, 'main'])
diff --git a/spec/lib/gitlab/database/gitlab_schema_spec.rb b/spec/lib/gitlab/database/gitlab_schema_spec.rb
index a5a67c2c918..611b2fbad72 100644
--- a/spec/lib/gitlab/database/gitlab_schema_spec.rb
+++ b/spec/lib/gitlab/database/gitlab_schema_spec.rb
@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::GitlabSchema do
 
     it 'all tables have assigned a known gitlab_schema' do
       is_expected.to all(
-        match([be_a(String), be_in([:gitlab_shared, :gitlab_main, :gitlab_ci])])
+        match([be_a(String), be_in([:gitlab_internal, :gitlab_shared, :gitlab_main, :gitlab_ci])])
       )
     end
 
@@ -42,12 +42,12 @@ RSpec.describe Gitlab::Database::GitlabSchema do
     where(:name, :classification) do
       'ci_builds'                       | :gitlab_ci
       'my_schema.ci_builds'             | :gitlab_ci
-      'information_schema.columns'      | :gitlab_shared
+      'information_schema.columns'      | :gitlab_internal
       '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
+      'pg_attribute'                    | :gitlab_internal
       'my_other_table'                  | :undefined_my_other_table
     end
 
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index 77284b4d128..34370c9a21f 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -100,14 +100,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
         expect(config.pool_size).to eq(4)
       end
     end
-
-    it 'calls reuse_primary_connection!' do
-      expect_next_instance_of(described_class) do |subject|
-        expect(subject).to receive(:reuse_primary_connection!).and_call_original
-      end
-
-      described_class.for_model(model)
-    end
   end
 
   describe '#load_balancing_enabled?' do
@@ -203,61 +195,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
     end
   end
 
-  describe '#replica_db_config' do
+  describe '#db_config' do
     let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::ApplicationRecord') }
     let(:config) { described_class.for_model(model) }
 
     it 'returns exactly db_config' do
-      expect(config.replica_db_config).to eq(db_config)
-    end
-
-    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
-      it 'does not change replica_db_config' do
-        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
-
-        expect(config.replica_db_config).to eq(db_config)
-      end
-    end
-  end
-
-  describe 'reuse_primary_connection!' do
-    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::ApplicationRecord') }
-    let(:config) { described_class.for_model(model) }
-
-    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_* not configured' do
-      it 'the primary connection uses default specification' do
-        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil)
-
-        expect(config.primary_connection_specification_name).to eq('Ci::ApplicationRecord')
-      end
-    end
-
-    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' 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
-      it 'raises exception' do
-        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'unknown')
-
-        expect { config.reuse_primary_connection! }.to raise_error /Invalid value for/
-      end
+      expect(config.db_config).to eq(db_config)
     end
   end
 end
diff --git a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
index ee2718171c0..41312dbedd6 100644
--- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -79,42 +79,55 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
     end
   end
 
-  describe '.insert_all!' do
+  describe 'methods using exec_insert_all on the connection', :request_store do
+    let(:model_class) do
+      Class.new(ApplicationRecord) do
+        self.table_name = "_test_connection_proxy_insert_all"
+      end
+    end
+
+    let(:session) { Gitlab::Database::LoadBalancing::Session.new }
+
     before do
       ActiveRecord::Schema.define do
-        create_table :_test_connection_proxy_bulk_insert, force: true do |t|
-          t.string :name, null: true
+        create_table :_test_connection_proxy_insert_all, force: true do |t|
+          t.string :name, null: false
+          t.index :name, unique: true
         end
       end
+
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
     end
 
     after do
       ActiveRecord::Schema.define do
-        drop_table :_test_connection_proxy_bulk_insert, force: true
+        drop_table :_test_connection_proxy_insert_all, force: true
       end
     end
 
-    let(:model_class) do
-      Class.new(ApplicationRecord) do
-        self.table_name = "_test_connection_proxy_bulk_insert"
+    describe '#upsert' do
+      it 'upserts a record and marks the session to stick to the primary' do
+        expect { 2.times { model_class.upsert({ name: 'test' }, unique_by: :name) } }
+          .to change { model_class.count }.from(0).to(1)
+          .and change { session.use_primary? }.from(false).to(true)
       end
     end
 
-    it 'inserts data in bulk' do
-      expect(model_class).to receive(:connection)
-        .at_least(:once)
-        .and_return(proxy)
+    describe '#insert_all!' do
+      it 'inserts multiple records and marks the session to stick to the primary' do
+        expect { model_class.insert_all([{ name: 'one' }, { name: 'two' }]) }
+          .to change { model_class.count }.from(0).to(2)
+          .and change { session.use_primary? }.from(false).to(true)
+      end
+    end
 
-      expect(proxy).to receive(:write_using_load_balancer)
-        .at_least(:once)
-        .and_call_original
-
-      expect do
-        model_class.insert_all! [
-          { name: "item1" },
-          { name: "item2" }
-        ]
-      end.to change { model_class.count }.by(2)
+    describe '#insert' do
+      it 'inserts a single record and marks the session to stick to the primary' do
+        expect { model_class.insert({ name: 'single' }) }
+          .to change { model_class.count }.from(0).to(1)
+          .and change { session.use_primary? }.from(false).to(true)
+      end
     end
   end
 
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 3c7819c04b6..34eb64997c1 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -487,46 +487,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
   end
 
-  describe 'primary connection re-use', :reestablished_active_record_base, :add_ci_connection do
-    let(:model) { Ci::ApplicationRecord }
-
-    describe '#read' do
-      it 'returns ci replica connection' do
-        expect { |b| lb.read(&b) }.to yield_with_args do |args|
-          expect(args.pool.db_config.name).to eq('ci_replica')
-        end
-      end
-
-      context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
-        it 'returns ci replica connection' do
-          stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
-
-          expect { |b| lb.read(&b) }.to yield_with_args do |args|
-            expect(args.pool.db_config.name).to eq('ci_replica')
-          end
-        end
-      end
-    end
-
-    describe '#read_write' do
-      it 'returns Ci::ApplicationRecord connection' do
-        expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
-          expect(args.pool.db_config.name).to eq('ci')
-        end
-      end
-
-      context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
-        it 'returns ActiveRecord::Base connection' do
-          stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
-
-          expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
-            expect(args.pool.db_config.name).to eq('main')
-          end
-        end
-      end
-    end
-  end
-
   describe '#wal_diff' do
     it 'returns the diff between two write locations' do
       loc1 = lb.send(:get_write_location, lb.pool.connection)
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
index c44637b8d06..fa6d71bca7f 100644
--- a/spec/lib/gitlab/database/load_balancing/setup_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
@@ -122,123 +122,68 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
   context 'uses correct base models', :reestablished_active_record_base do
     using RSpec::Parameterized::TableSyntax
 
-    where do
-      {
-        "it picks a dedicated CI connection" => {
-          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
-          request_store_active: false,
-          ff_force_no_sharing_primary_model: false,
-          expectations: {
-            main: { read: 'main_replica', write: 'main' },
-            ci: { read: 'ci_replica', write: 'ci' }
-          }
-        },
-        "with re-use of primary connection it uses CI connection for reads" => {
-          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
-          request_store_active: false,
-          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 force_no_sharing_primary_model enabled with RequestStore it sticks FF and uses CI connection for reads and writes" => {
-          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
-          request_store_active: true,
-          ff_force_no_sharing_primary_model: true,
-          expectations: {
-            main: { read: 'main_replica', write: 'main' },
-            ci: { read: 'ci_replica', write: 'ci' }
-          }
-        },
-        "with re-use and FF force_no_sharing_primary_model enabled without RequestStore it doesn't use FF and uses CI connection for reads only" => {
-          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
-          request_store_active: true,
-          ff_force_no_sharing_primary_model: false,
-          expectations: {
-            main: { read: 'main_replica', write: 'main' },
-            ci: { read: 'ci_replica', write: 'main' }
-          }
-        }
-      }
-    end
-
-    with_them do
-      let(:ci_class) do
-        Class.new(ActiveRecord::Base) do
-          def self.name
-            'Ci::ApplicationRecordTemporary'
-          end
-
-          establish_connection ActiveRecord::DatabaseConfigurations::HashConfig.new(
-            Rails.env,
-            'ci',
-            ActiveRecord::Base.connection_db_config.configuration_hash
-          )
+    let(:ci_class) do
+      Class.new(ActiveRecord::Base) do
+        def self.name
+          'Ci::ApplicationRecordTemporary'
         end
-      end
 
-      let(:models) do
-        {
-          main: ActiveRecord::Base,
-          ci: ci_class
-        }
+        establish_connection ActiveRecord::DatabaseConfigurations::HashConfig.new(
+          Rails.env,
+          'ci',
+          ActiveRecord::Base.connection_db_config.configuration_hash
+        )
       end
+    end
 
-      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
-          end
-        else
-          example.run
-        end
-      end
+    let(:models) do
+      {
+        main: ActiveRecord::Base,
+        ci: ci_class
+      }
+    end
 
-      before do
-        allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
+    before do
+      allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
 
-        # Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects
-        allow_next_instance_of(described_class) do |setup|
-          allow(setup).to receive(:configure_connection)
+      # Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects
+      allow_next_instance_of(described_class) do |setup|
+        allow(setup).to receive(:configure_connection)
 
-          allow(setup).to receive(:setup_class_attribute) do |attribute, value|
-            allow(setup.model).to receive(attribute) { value }
-          end
+        allow(setup).to receive(:setup_class_attribute) do |attribute, value|
+          allow(setup.model).to receive(attribute) { value }
         end
+      end
 
-        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci)
-
-        # Make load balancer to force init with a dedicated replicas connections
-        models.each do |_, model|
-          described_class.new(model).tap do |subject|
-            subject.configuration.hosts = [subject.configuration.replica_db_config.host]
-            subject.setup
-          end
+      # Make load balancer to force init with a dedicated replicas connections
+      models.each do |_, model|
+        described_class.new(model).tap do |subject|
+          subject.configuration.hosts = [subject.configuration.db_config.host]
+          subject.setup
         end
       end
+    end
 
-      it 'results match expectations' do
-        result = models.transform_values do |model|
-          load_balancer = model.connection.instance_variable_get(:@load_balancer)
-
-          {
-            read: load_balancer.read { |connection| connection.pool.db_config.name },
-            write: load_balancer.read_write { |connection| connection.pool.db_config.name }
-          }
-        end
+    it 'results match expectations' do
+      result = models.transform_values do |model|
+        load_balancer = model.connection.instance_variable_get(:@load_balancer)
 
-        expect(result).to eq(expectations)
+        {
+          read: load_balancer.read { |connection| connection.pool.db_config.name },
+          write: load_balancer.read_write { |connection| connection.pool.db_config.name }
+        }
       end
 
-      it 'does return load_balancer assigned to a given connection' do
-        models.each do |name, model|
-          expect(model.load_balancer.name).to eq(name)
-          expect(model.sticking.instance_variable_get(:@load_balancer)).to eq(model.load_balancer)
-        end
+      expect(result).to eq({
+        main: { read: 'main_replica', write: 'main' },
+        ci: { read: 'ci_replica', write: 'ci' }
+      })
+    end
+
+    it 'does return load_balancer assigned to a given connection' do
+      models.each do |name, model|
+        expect(model.load_balancer.name).to eq(name)
+        expect(model.sticking.instance_variable_get(:@load_balancer)).to eq(model.load_balancer)
       end
     end
   end
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
index 9acf80e684f..b7915e6cf69 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -57,6 +57,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         run_middleware
 
         expect(job['wal_locations']).to be_nil
+        expect(job['wal_location_source']).to be_nil
       end
 
       include_examples 'job data consistency'
@@ -96,6 +97,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
           run_middleware
 
           expect(job['wal_locations']).to eq(expected_location)
+          expect(job['wal_location_source']).to eq(:replica)
         end
 
         include_examples 'job data consistency'
@@ -120,6 +122,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
           run_middleware
 
           expect(job['wal_locations']).to eq(expected_location)
+          expect(job['wal_location_source']).to eq(:primary)
         end
 
         include_examples 'job data consistency'
@@ -162,6 +165,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
           run_middleware
 
           expect(job['wal_locations']).to eq(wal_locations)
+          expect(job['wal_location_source']).to be_nil
         end
       end
 
diff --git a/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb b/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
new file mode 100644
index 00000000000..30e5fbbd803
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis, :delete do
+  let(:model) { ApplicationRecord }
+  let(:db_host) { model.connection_pool.db_config.host }
+
+  let(:test_table_name) { '_test_foo' }
+
+  before do
+    # Patch in our load balancer config, simply pointing at the test database twice
+    allow(Gitlab::Database::LoadBalancing::Configuration).to receive(:for_model) do |base_model|
+      Gitlab::Database::LoadBalancing::Configuration.new(base_model, [db_host, db_host])
+    end
+
+    Gitlab::Database::LoadBalancing::Setup.new(ApplicationRecord).setup
+
+    model.connection.execute(<<~SQL)
+      CREATE TABLE IF NOT EXISTS #{test_table_name} (id SERIAL PRIMARY KEY, value INTEGER)
+    SQL
+  end
+
+  after do
+    model.connection.execute(<<~SQL)
+      DROP TABLE IF EXISTS #{test_table_name}
+    SQL
+  end
+
+  def execute(conn)
+    conn.execute("INSERT INTO #{test_table_name} (value) VALUES (1)")
+    backend_pid = conn.execute("SELECT pg_backend_pid() AS pid").to_a.first['pid']
+
+    # This will result in a PG error, which is not raised.
+    # Instead, we retry the statement on a fresh connection (where the pid is different and it does nothing)
+    # and the load balancer continues with a fresh connection and no transaction if a transaction was open previously
+    conn.execute(<<~SQL)
+      SELECT CASE
+      WHEN pg_backend_pid() = #{backend_pid} THEN
+        pg_terminate_backend(#{backend_pid})
+      END
+    SQL
+
+    # This statement will execute on a new connection, and violate transaction semantics
+    # if we were in a transaction before
+    conn.execute("INSERT INTO #{test_table_name} (value) VALUES (2)")
+  end
+
+  it 'logs a warning when violating transaction semantics with writes' do
+    conn = model.connection
+
+    expect(::Gitlab::Database::LoadBalancing::Logger).to receive(:warn).with(hash_including(event: :transaction_leak))
+
+    conn.transaction do
+      expect(conn).to be_transaction_open
+
+      execute(conn)
+
+      expect(conn).not_to be_transaction_open
+    end
+
+    values = conn.execute("SELECT value FROM #{test_table_name}").to_a.map { |row| row['value'] }
+    expect(values).to contain_exactly(2) # Does not include 1 because the transaction was aborted and leaked
+  end
+
+  it 'does not log a warning when no transaction is open to be leaked' do
+    conn = model.connection
+
+    expect(::Gitlab::Database::LoadBalancing::Logger)
+      .not_to receive(:warn).with(hash_including(event: :transaction_leak))
+
+    expect(conn).not_to be_transaction_open
+
+    execute(conn)
+
+    expect(conn).not_to be_transaction_open
+
+    values = conn.execute("SELECT value FROM #{test_table_name}").to_a.map { |row| row['value'] }
+    expect(values).to contain_exactly(1, 2) # Includes both rows because there was no transaction to roll back
+  end
+end
diff --git a/spec/lib/gitlab/database/migration_helpers/announce_database_spec.rb b/spec/lib/gitlab/database/migration_helpers/announce_database_spec.rb
new file mode 100644
index 00000000000..57c51c9d9c2
--- /dev/null
+++ b/spec/lib/gitlab/database/migration_helpers/announce_database_spec.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::MigrationHelpers::AnnounceDatabase do
+  let(:migration) do
+    ActiveRecord::Migration.new('MyMigration', 1111).extend(described_class)
+  end
+
+  describe '#announce' do
+    it 'prefixes message with database name' do
+      expect { migration.announce('migrating') }.to output(/^main: == 1111 MyMigration: migrating/).to_stdout
+    end
+  end
+
+  describe '#say' do
+    it 'prefixes message with database name' do
+      expect { migration.say('transaction_open?()') }.to output(/^main: -- transaction_open?()/).to_stdout
+    end
+
+    it 'prefixes subitem message with database name' do
+      expect { migration.say('0.0000s', true) }.to output(/^main:    -> 0.0000s/).to_stdout
+    end
+  end
+
+  describe '#write' do
+    it 'does not prefix empty write' do
+      expect { migration.write }.to output(/^$/).to_stdout
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 04fe1fad10e..e09016b2b2b 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -2079,6 +2079,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         t.integer :other_id
         t.timestamps
       end
+
+      allow(model).to receive(:transaction_open?).and_return(false)
     end
 
     context 'when the target table does not exist' do
@@ -2191,6 +2193,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         t.timestamps
       end
 
+      allow(model).to receive(:transaction_open?).and_return(false)
+
       model.initialize_conversion_of_integer_to_bigint(table, columns, primary_key: primary_key)
       model.backfill_conversion_of_integer_to_bigint(table, columns, primary_key: primary_key)
     end
@@ -2242,10 +2246,20 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       }
     end
 
+    let(:migration_attributes) do
+      configuration.merge(gitlab_schema: Gitlab::Database.gitlab_schemas_for_connection(model.connection).first)
+    end
+
+    before do
+      allow(model).to receive(:transaction_open?).and_return(false)
+    end
+
     subject(:ensure_batched_background_migration_is_finished) { model.ensure_batched_background_migration_is_finished(**configuration) }
 
     it 'raises an error when migration exists and is not marked as finished' do
-      create(:batched_background_migration, :active, configuration)
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!).twice
+
+      create(:batched_background_migration, :active, migration_attributes)
 
       allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
         allow(runner).to receive(:finalize).with(job_class_name, table, column_name, job_arguments).and_return(false)
@@ -2255,7 +2269,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         .to raise_error "Expected batched background migration for the given configuration to be marked as 'finished', but it is 'active':" \
             "\t#{configuration}" \
             "\n\n" \
-            "Finalize it manually by running" \
+            "Finalize it manually by running the following command in a `bash` or `sh` shell:" \
             "\n\n" \
             "\tsudo gitlab-rake gitlab:background_migrations:finalize[CopyColumnUsingBackgroundMigrationJob,events,id,'[[\"id\"]\\,[\"id_convert_to_bigint\"]\\,null]']" \
             "\n\n" \
@@ -2265,13 +2279,19 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
 
     it 'does not raise error when migration exists and is marked as finished' do
-      create(:batched_background_migration, :finished, configuration)
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
+      create(:batched_background_migration, :finished, migration_attributes)
 
       expect { ensure_batched_background_migration_is_finished }
         .not_to raise_error
     end
 
     it 'logs a warning when migration does not exist' do
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
+      create(:batched_background_migration, :active, migration_attributes.merge(gitlab_schema: :gitlab_something_else))
+
       expect(Gitlab::AppLogger).to receive(:warn)
         .with("Could not find batched background migration for the given configuration: #{configuration}")
 
@@ -2280,6 +2300,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
 
     it 'finalizes the migration' do
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!).twice
+
       migration = create(:batched_background_migration, :active, configuration)
 
       allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
@@ -2291,6 +2313,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
     context 'when the flag finalize is false' do
       it 'does not finalize the migration' do
+        expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
         create(:batched_background_migration, :active, configuration)
 
         allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
@@ -3257,4 +3281,20 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       model.rename_constraint(:test_table, :fk_old_name, :fk_new_name)
     end
   end
+
+  describe '#drop_sequence' do
+    it "executes the statement to drop the sequence" do
+      expect(model).to receive(:execute).with /ALTER TABLE "test_table" ALTER COLUMN "test_column" DROP DEFAULT;\nDROP SEQUENCE IF EXISTS "test_table_id_seq"/
+
+      model.drop_sequence(:test_table, :test_column, :test_table_id_seq)
+    end
+  end
+
+  describe '#add_sequence' do
+    it "executes the statement to add the sequence" do
+      expect(model).to receive(:execute).with "CREATE SEQUENCE \"test_table_id_seq\" START 1;\nALTER TABLE \"test_table\" ALTER COLUMN \"test_column\" SET DEFAULT nextval(\'test_table_id_seq\')\n"
+
+      model.add_sequence(:test_table, :test_column, :test_table_id_seq, 1)
+    end
+  end
 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 b0caa21e01a..c423340a572 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -444,7 +444,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
           it 'does restore connection hierarchy' do
             expect_next_instances_of(job_class, 1..) do |job|
               expect(job).to receive(:perform) do
-                validate_connections!
+                validate_connections_stack!
               end
             end
 
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 d1a66036149..f3414727245 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
@@ -3,8 +3,14 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers do
+  let(:migration_class) do
+    Class.new(ActiveRecord::Migration[6.1])
+      .include(described_class)
+      .include(Gitlab::Database::Migrations::ReestablishedConnectionStack)
+  end
+
   let(:migration) do
-    ActiveRecord::Migration.new.extend(described_class)
+    migration_class.new
   end
 
   describe '#queue_batched_background_migration' do
@@ -12,6 +18,9 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
 
     before do
       allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
+      allow(migration).to receive(:transaction_open?).and_return(false)
     end
 
     context 'when such migration already exists' do
@@ -27,7 +36,8 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
           batch_class_name: 'MyBatchClass',
           batch_size: 200,
           sub_batch_size: 20,
-          job_arguments: [[:id], [:id_convert_to_bigint]]
+          job_arguments: [[:id], [:id_convert_to_bigint]],
+          gitlab_schema: :gitlab_ci
         )
 
         expect do
@@ -41,7 +51,8 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
             batch_max_value: 1000,
             batch_class_name: 'MyBatchClass',
             batch_size: 100,
-            sub_batch_size: 10)
+            sub_batch_size: 10,
+            gitlab_schema: :gitlab_ci)
         end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
       end
     end
@@ -60,7 +71,8 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
           batch_class_name: 'MyBatchClass',
           batch_size: 100,
           max_batch_size: 10000,
-          sub_batch_size: 10)
+          sub_batch_size: 10,
+          gitlab_schema: :gitlab_ci)
       end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
 
       expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
@@ -76,7 +88,8 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
         sub_batch_size: 10,
         job_arguments: %w[],
         status_name: :active,
-        total_tuple_count: pgclass_info.cardinality_estimate)
+        total_tuple_count: pgclass_info.cardinality_estimate,
+        gitlab_schema: 'gitlab_ci')
     end
 
     context 'when the job interval is lower than the minimum' do
@@ -160,6 +173,31 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
 
           expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_finished
         end
+
+        context 'when within transaction' do
+          before do
+            allow(migration).to receive(:transaction_open?).and_return(true)
+          end
+
+          it 'does raise an exception' do
+            expect { migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)}
+              .to raise_error /`queue_batched_background_migration` cannot be run inside a transaction./
+          end
+        end
+      end
+    end
+
+    context 'when gitlab_schema is not given' do
+      it 'fetches gitlab_schema from the migration context' do
+        expect(migration).to receive(:gitlab_schema_from_context).and_return(:gitlab_ci)
+
+        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.gitlab_schema).to eq('gitlab_ci')
       end
     end
   end
@@ -167,6 +205,12 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
   describe '#finalize_batched_background_migration' do
     let!(:batched_migration) { create(:batched_background_migration, job_class_name: 'MyClass', table_name: :projects, column_name: :id, job_arguments: []) }
 
+    before do
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
+      allow(migration).to receive(:transaction_open?).and_return(false)
+    end
+
     it 'finalizes the migration' do
       allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
         expect(runner).to receive(:finalize).with('MyClass', :projects, :id, [])
@@ -183,24 +227,162 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
       end
     end
 
-    context 'when uses a CI connection', :reestablished_active_record_base do
+    context 'when within transaction' do
       before do
-        skip_if_multiple_databases_not_setup
+        allow(migration).to receive(:transaction_open?).and_return(true)
+      end
 
-        ActiveRecord::Base.establish_connection(:ci) # rubocop:disable Database/EstablishConnection
+      it 'does raise an exception' do
+        expect { migration.finalize_batched_background_migration(job_class_name: 'MyJobClass', table_name: :projects, column_name: :id, job_arguments: []) }
+          .to raise_error /`finalize_batched_background_migration` cannot be run inside a transaction./
       end
+    end
 
-      it 'raises an exception' do
-        ci_migration = create(:batched_background_migration, :active)
+    context 'when running migration in reconfigured ActiveRecord::Base context' do
+      it_behaves_like 'reconfigures connection stack', 'ci' do
+        before do
+          create(:batched_background_migration,
+            job_class_name: 'Ci::MyClass',
+            table_name: :ci_builds,
+            column_name: :id,
+            job_arguments: [],
+            gitlab_schema: :gitlab_ci)
+        end
+
+        context 'when restrict_gitlab_migration is set to gitlab_ci' do
+          it 'finalizes the migration' do
+            migration_class.include(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema)
+            migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_ci
+
+            allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+              expect(runner).to receive(:finalize).with('Ci::MyClass', :ci_builds, :id, []) do
+                validate_connections_stack!
+              end
+            end
+
+            migration.finalize_batched_background_migration(
+              job_class_name: 'Ci::MyClass', table_name: :ci_builds, column_name: :id, job_arguments: [])
+          end
+        end
+
+        context 'when restrict_gitlab_migration is set to gitlab_main' do
+          it 'does not find any migrations' do
+            migration_class.include(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema)
+            migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_main
+
+            expect do
+              migration.finalize_batched_background_migration(
+                job_class_name: 'Ci::MyClass', table_name: :ci_builds, column_name: :id, job_arguments: [])
+            end.to raise_error /Could not find batched background migration/
+          end
+        end
+
+        context 'when no restrict is set' do
+          it 'does not find any migrations' do
+            expect do
+              migration.finalize_batched_background_migration(
+                job_class_name: 'Ci::MyClass', table_name: :ci_builds, column_name: :id, job_arguments: [])
+            end.to raise_error /Could not find batched background migration/
+          end
+        end
+      end
+    end
+
+    context 'when within transaction' do
+      before do
+        allow(migration).to receive(:transaction_open?).and_return(true)
+      end
+
+      it 'does raise an exception' do
+        expect { migration.finalize_batched_background_migration(job_class_name: 'MyJobClass', table_name: :projects, column_name: :id, job_arguments: []) }
+          .to raise_error /`finalize_batched_background_migration` cannot be run inside a transaction./
+      end
+    end
+  end
+
+  describe '#delete_batched_background_migration' do
+    let(:transaction_open) { false }
+
+    before do
+      expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
+
+      allow(migration).to receive(:transaction_open?).and_return(transaction_open)
+    end
+
+    context 'when migration exists' do
+      it 'deletes it' 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.finalize_batched_background_migration(
-            job_class_name: ci_migration.job_class_name,
-            table_name: ci_migration.table_name,
-            column_name: ci_migration.column_name,
-            job_arguments: ci_migration.job_arguments
-          )
-        end.to raise_error /is currently not supported when running in decomposed/
+          migration.delete_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            [[:id], [:id_convert_to_bigint]])
+        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.from(1).to(0)
+      end
+    end
+
+    context 'when migration does not exist' do
+      it 'does nothing' do
+        create(
+          :batched_background_migration,
+          job_class_name: 'SomeOtherJobClass',
+          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.delete_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            [[:id], [:id_convert_to_bigint]])
+        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
+      end
+    end
+
+    context 'when within transaction' do
+      let(:transaction_open) { true }
+
+      it 'raises an exception' do
+        expect { migration.delete_batched_background_migration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]]) }
+          .to raise_error /`#delete_batched_background_migration` cannot be run inside a transaction./
+      end
+    end
+  end
+
+  describe '#gitlab_schema_from_context' do
+    context 'when allowed_gitlab_schemas is not available' do
+      it 'defaults to :gitlab_main' do
+        expect(migration.gitlab_schema_from_context).to eq(:gitlab_main)
+      end
+    end
+
+    context 'when allowed_gitlab_schemas is available' do
+      it 'uses schema from allowed_gitlab_schema' do
+        expect(migration).to receive(:allowed_gitlab_schemas).and_return([:gitlab_ci])
+
+        expect(migration.gitlab_schema_from_context).to eq(:gitlab_ci)
       end
     end
   end
diff --git a/spec/lib/gitlab/database/migrations/reestablished_connection_stack_spec.rb b/spec/lib/gitlab/database/migrations/reestablished_connection_stack_spec.rb
index cfb308c63e4..d197f39be40 100644
--- a/spec/lib/gitlab/database/migrations/reestablished_connection_stack_spec.rb
+++ b/spec/lib/gitlab/database/migrations/reestablished_connection_stack_spec.rb
@@ -16,7 +16,7 @@ RSpec.describe Gitlab::Database::Migrations::ReestablishedConnectionStack do
         it_behaves_like "reconfigures connection stack", db_config_name do
           it 'does restore connection hierarchy' do
             model.with_restored_connection_stack do
-              validate_connections!
+              validate_connections_stack!
             end
           end
 
diff --git a/spec/lib/gitlab/database/migrations/test_batched_background_runner_spec.rb b/spec/lib/gitlab/database/migrations/test_batched_background_runner_spec.rb
index fbfff1268cc..2f3d44f6f8f 100644
--- a/spec/lib/gitlab/database/migrations/test_batched_background_runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/test_batched_background_runner_spec.rb
@@ -4,7 +4,6 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::TestBatchedBackgroundRunner, :freeze_time do
   include Gitlab::Database::MigrationHelpers
-  include Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers
   include Database::MigrationTestingHelpers
 
   let(:result_dir) { Dir.mktmpdir }
@@ -13,6 +12,10 @@ RSpec.describe Gitlab::Database::Migrations::TestBatchedBackgroundRunner, :freez
     FileUtils.rm_rf(result_dir)
   end
 
+  let(:migration) do
+    ActiveRecord::Migration.new.extend(Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers)
+  end
+
   let(:connection) { ApplicationRecord.connection }
 
   let(:table_name) { "_test_column_copying"}
@@ -26,11 +29,13 @@ RSpec.describe Gitlab::Database::Migrations::TestBatchedBackgroundRunner, :freez
 
       insert into #{table_name} (id) select i from generate_series(1, 1000) g(i);
     SQL
+
+    allow(migration).to receive(:transaction_open?).and_return(false)
   end
 
   context 'running a real background migration' do
     it 'runs sampled jobs from the batched background migration' do
-      queue_batched_background_migration('CopyColumnUsingBackgroundMigrationJob',
+      migration.queue_batched_background_migration('CopyColumnUsingBackgroundMigrationJob',
                                          table_name, :id,
                                          :id, :data,
                                          batch_size: 100,
@@ -46,7 +51,9 @@ RSpec.describe Gitlab::Database::Migrations::TestBatchedBackgroundRunner, :freez
     let(:migration_name) { 'TestBackgroundMigration' }
 
     before do
-      queue_batched_background_migration(migration_name, table_name, :id, job_interval: 5.minutes, batch_size: 100)
+      migration.queue_batched_background_migration(
+        migration_name, table_name, :id, job_interval: 5.minutes, batch_size: 100
+      )
     end
 
     it 'samples jobs' do
@@ -67,13 +74,13 @@ RSpec.describe Gitlab::Database::Migrations::TestBatchedBackgroundRunner, :freez
         travel 3.days
 
         new_migration = define_background_migration('NewMigration') { travel 1.second }
-        queue_batched_background_migration('NewMigration', table_name, :id,
+        migration.queue_batched_background_migration('NewMigration', table_name, :id,
                                            job_interval: 5.minutes,
                                            batch_size: 10,
                                            sub_batch_size: 5)
 
         other_new_migration = define_background_migration('NewMigration2') { travel 2.seconds }
-        queue_batched_background_migration('NewMigration2', table_name, :id,
+        migration.queue_batched_background_migration('NewMigration2', table_name, :id,
                                            job_interval: 5.minutes,
                                            batch_size: 10,
                                            sub_batch_size: 5)
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index 64dcdb9628a..dca4548a0a3 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -8,7 +8,11 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
   def has_partition(model, month)
     Gitlab::Database::PostgresPartition.for_parent_table(model.table_name).any? do |partition|
-      Gitlab::Database::Partitioning::TimePartition.from_sql(model.table_name, partition.name, partition.condition).from == month
+      Gitlab::Database::Partitioning::TimePartition.from_sql(
+        model.table_name,
+        partition.name,
+        partition.condition
+      ).from == month
     end
   end
 
@@ -16,14 +20,17 @@ 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: [], after_adding_partitions: nil) }
     let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "issues" }
+    let(:partitioning_strategy) do
+      double(missing_partitions: partitions, extra_partitions: [], after_adding_partitions: nil)
+    end
 
     before do
       allow(connection).to receive(:table_exists?).and_call_original
       allow(connection).to receive(:table_exists?).with(table).and_return(true)
       allow(connection).to receive(:execute).and_call_original
+      expect(partitioning_strategy).to receive(:validate_and_fix)
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -84,13 +91,16 @@ 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: [], after_adding_partitions: nil) }
     let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "foo" }
+    let(:partitioning_strategy) do
+      double(extra_partitions: extra_partitions, missing_partitions: [], after_adding_partitions: nil)
+    end
 
     before do
       allow(connection).to receive(:table_exists?).and_call_original
       allow(connection).to receive(:table_exists?).with(table).and_return(true)
+      expect(partitioning_strategy).to receive(:validate_and_fix)
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -107,6 +117,24 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
       sync_partitions
     end
+
+    it 'logs an error if the partitions are not detachable' do
+      allow(Gitlab::Database::PostgresForeignKey).to receive(:by_referenced_table_identifier).with("public.foo")
+        .and_return([double(name: "fk_1", constrained_table_identifier: "public.constrainted_table_1")])
+
+      expect(Gitlab::AppLogger).to receive(:error).with(
+        {
+          message: "Failed to create / detach partition(s)",
+          connection_name: "main",
+          exception_class: Gitlab::Database::Partitioning::PartitionManager::UnsafeToDetachPartitionError,
+          exception_message:
+            "Cannot detach foo1, it would block while checking foreign key fk_1 on public.constrainted_table_1",
+          table_name: "foo"
+        }
+      )
+
+      sync_partitions
+    end
   end
 
   describe '#detach_partitions' do
diff --git a/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
index 1cec0463055..d8b06ee1a5d 100644
--- a/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/sliding_list_strategy_spec.rb
@@ -37,12 +37,75 @@ RSpec.describe Gitlab::Database::Partitioning::SlidingListStrategy do
   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')
+        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 '#validate_and_fix' do
+    context 'feature flag is disabled' do
+      before do
+        stub_feature_flags(fix_sliding_list_partitioning: false)
+      end
+
+      it 'does not try to fix the default partition value' do
+        connection.change_column_default(model.table_name, strategy.partitioning_key, 3)
+        expect(strategy.model.connection).not_to receive(:change_column_default)
+        strategy.validate_and_fix
+      end
+    end
+
+    context 'feature flag is enabled' do
+      before do
+        stub_feature_flags(fix_sliding_list_partitioning: true)
+      end
+
+      it 'does not call change_column_default if the partitioning in a valid state' do
+        expect(strategy.model.connection).not_to receive(:change_column_default)
+
+        strategy.validate_and_fix
+      end
+
+      it 'calls change_column_default on partition_key with the most default partition number' do
+        connection.change_column_default(model.table_name, strategy.partitioning_key, 1)
+
+        expect(Gitlab::AppLogger).to receive(:warn).with(
+          message: 'Fixed default value of sliding_list_strategy partitioning_key',
+          connection_name: 'main',
+          old_value: 1,
+          new_value: 2,
+          table_name: table_name,
+          column: strategy.partitioning_key
+        )
+
+        expect(strategy.model.connection).to receive(:change_column_default).with(
+          model.table_name, strategy.partitioning_key, 2
+        ).and_call_original
+
+        strategy.validate_and_fix
+      end
+
+      it 'does not change the default column if it has been changed in the meanwhile by another process' do
+        expect(strategy).to receive(:current_default_value).and_return(1, 2)
+
+        expect(strategy.model.connection).not_to receive(:change_column_default)
+
+        expect(Gitlab::AppLogger).to receive(:warn).with(
+          message: 'Table partitions or partition key default value have been changed by another process',
+          table_name: table_name,
+          default_value: 2
+        )
+
+        strategy.validate_and_fix
+      end
+    end
+  end
+
   describe '#active_partition' do
     it 'is the partition with the largest value' do
       expect(strategy.active_partition.value).to eq(2)
@@ -157,6 +220,7 @@ RSpec.describe Gitlab::Database::Partitioning::SlidingListStrategy do
       end.not_to raise_error
     end
   end
+
   context 'redirecting inserts as the active partition changes' do
     let(:model) do
       Class.new(ApplicationRecord) do
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 0d687db0f96..62c5ead855a 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', :add_ci_connection, :request_store do
+  context 'properly observes all queries', :add_ci_connection do
     using RSpec::Parameterized::TableSyntax
 
     where do
@@ -28,8 +28,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
           expectations: {
             gitlab_schemas: "gitlab_main",
             db_config_name: "main"
-          },
-          setup: nil
+          }
         },
         "for query accessing gitlab_ci and gitlab_main" => {
           model: ApplicationRecord,
@@ -37,8 +36,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
           expectations: {
             gitlab_schemas: "gitlab_ci,gitlab_main",
             db_config_name: "main"
-          },
-          setup: nil
+          }
         },
         "for query accessing gitlab_ci and gitlab_main the gitlab_schemas is always ordered" => {
           model: ApplicationRecord,
@@ -46,8 +44,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
           expectations: {
             gitlab_schemas: "gitlab_ci,gitlab_main",
             db_config_name: "main"
-          },
-          setup: nil
+          }
         },
         "for query accessing CI database" => {
           model: Ci::ApplicationRecord,
@@ -56,62 +53,6 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
             gitlab_schemas: "gitlab_ci",
             db_config_name: "ci"
           }
-        },
-        "for query accessing CI database with re-use and disabled sharing" => {
-          model: Ci::ApplicationRecord,
-          sql: "SELECT 1 FROM ci_builds",
-          expectations: {
-            gitlab_schemas: "gitlab_ci",
-            db_config_name: "ci",
-            ci_dedicated_primary_connection: true
-          },
-          setup: ->(_) do
-            skip_if_multiple_databases_not_setup
-            stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
-            stub_feature_flags(force_no_sharing_primary_model: true)
-          end
-        },
-        "for query accessing CI database with re-use and enabled sharing" => {
-          model: Ci::ApplicationRecord,
-          sql: "SELECT 1 FROM ci_builds",
-          expectations: {
-            gitlab_schemas: "gitlab_ci",
-            db_config_name: "ci",
-            ci_dedicated_primary_connection: false
-          },
-          setup: ->(_) do
-            skip_if_multiple_databases_not_setup
-            stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
-            stub_feature_flags(force_no_sharing_primary_model: false)
-          end
-        },
-        "for query accessing CI database without re-use and disabled sharing" => {
-          model: Ci::ApplicationRecord,
-          sql: "SELECT 1 FROM ci_builds",
-          expectations: {
-            gitlab_schemas: "gitlab_ci",
-            db_config_name: "ci",
-            ci_dedicated_primary_connection: true
-          },
-          setup: ->(_) do
-            skip_if_multiple_databases_not_setup
-            stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil)
-            stub_feature_flags(force_no_sharing_primary_model: true)
-          end
-        },
-        "for query accessing CI database without re-use and enabled sharing" => {
-          model: Ci::ApplicationRecord,
-          sql: "SELECT 1 FROM ci_builds",
-          expectations: {
-            gitlab_schemas: "gitlab_ci",
-            db_config_name: "ci",
-            ci_dedicated_primary_connection: true
-          },
-          setup: ->(_) do
-            skip_if_multiple_databases_not_setup
-            stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil)
-            stub_feature_flags(force_no_sharing_primary_model: false)
-          end
         }
       }
     end
@@ -122,15 +63,11 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
       end
 
       it do
-        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil)
-
-        instance_eval(&setup) if setup
-
         allow(::Ci::ApplicationRecord.load_balancer).to receive(:configuration)
           .and_return(Gitlab::Database::LoadBalancing::Configuration.for_model(::Ci::ApplicationRecord))
 
         expect(described_class.schemas_metrics).to receive(:increment)
-          .with({ ci_dedicated_primary_connection: anything }.merge(expectations)).and_call_original
+          .with(expectations).and_call_original
 
         process_sql(model, sql)
       end
diff --git a/spec/lib/gitlab/database/shared_model_spec.rb b/spec/lib/gitlab/database/shared_model_spec.rb
index 574111f4c01..c88edc17817 100644
--- a/spec/lib/gitlab/database/shared_model_spec.rb
+++ b/spec/lib/gitlab/database/shared_model_spec.rb
@@ -27,6 +27,19 @@ RSpec.describe Gitlab::Database::SharedModel do
       end
     end
 
+    it 'raises an error if the connection does not include `:gitlab_shared` schema' do
+      allow(Gitlab::Database)
+        .to receive(:gitlab_schemas_for_connection)
+        .with(new_connection)
+        .and_return([:gitlab_main])
+
+      expect_original_connection_around do
+        expect do
+          described_class.using_connection(new_connection) {}
+        end.to raise_error(/Cannot set `SharedModel` to connection/)
+      end
+    end
+
     context 'when multiple connection overrides are nested', :aggregate_failures do
       let(:second_connection) { double('connection') }