Add latest changes from gitlab-org/gitlab@15-6-stable-eev15.6.0-rc42

30 files changed, 2192 insertions, 990 deletions

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 32746a46308..cc9f3d5b7f1 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
@@ -7,7 +7,15 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
 
   it { is_expected.to be_a Gitlab::Database::SharedModel }
 
-  it { expect(described_class::TIMEOUT_EXCEPTIONS).to match_array [ActiveRecord::StatementTimeout, ActiveRecord::ConnectionTimeoutError, ActiveRecord::AdapterTimeout, ActiveRecord::LockWaitTimeout] }
+  specify do
+    expect(described_class::TIMEOUT_EXCEPTIONS).to contain_exactly(
+      ActiveRecord::StatementTimeout,
+      ActiveRecord::ConnectionTimeoutError,
+      ActiveRecord::AdapterTimeout,
+      ActiveRecord::LockWaitTimeout,
+      ActiveRecord::QueryCanceled
+    )
+  end
 
   describe 'associations' do
     it { is_expected.to belong_to(:batched_migration).with_foreign_key(:batched_background_migration_id) }
@@ -272,7 +280,13 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
       context 'when is a timeout exception' do
         let(:exception) { ActiveRecord::StatementTimeout.new }
 
-        it { expect(subject).to be_truthy  }
+        it { expect(subject).to be_truthy }
+      end
+
+      context 'when is a QueryCanceled exception' do
+        let(:exception) { ActiveRecord::QueryCanceled.new }
+
+        it { expect(subject).to be_truthy }
       end
 
       context 'when is not a timeout exception' do
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 1ac9cbae036..31ae5e9b55d 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -211,6 +211,102 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
         expect(active_migration).to eq(migration3)
       end
     end
+
+    context 'when there are no active migrations available' do
+      it 'returns nil' do
+        expect(active_migration).to eq(nil)
+      end
+    end
+  end
+
+  describe '.find_executable' do
+    let(:connection) { Gitlab::Database.database_base_models[:main].connection }
+    let(:migration_id) { migration.id }
+
+    subject(:executable_migration) { described_class.find_executable(migration_id, connection: connection) }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        example.run
+      end
+    end
+
+    context 'when the migration does not exist' do
+      let(:migration_id) { non_existing_record_id }
+
+      it 'returns nil' do
+        expect(executable_migration).to be_nil
+      end
+    end
+
+    context 'when the migration is not active' do
+      let!(:migration) { create(:batched_background_migration, :finished) }
+
+      it 'returns nil' do
+        expect(executable_migration).to be_nil
+      end
+    end
+
+    context 'when the migration is on hold' do
+      let!(:migration) { create(:batched_background_migration, :active, on_hold_until: 10.minutes.from_now) }
+
+      it 'returns nil' do
+        expect(executable_migration).to be_nil
+      end
+    end
+
+    context 'when the migration is not available for the current connection' do
+      let!(:migration) { create(:batched_background_migration, :active, gitlab_schema: :gitlab_not_existing) }
+
+      it 'returns nil' do
+        expect(executable_migration).to be_nil
+      end
+    end
+
+    context 'when ther migration exists and is executable' do
+      let!(:migration) { create(:batched_background_migration, :active, gitlab_schema: :gitlab_main) }
+
+      it 'returns the migration' do
+        expect(executable_migration).to eq(migration)
+      end
+    end
+  end
+
+  describe '.active_migrations_distinct_on_table' do
+    let(:connection) { Gitlab::Database.database_base_models[:main].connection }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        example.run
+      end
+    end
+
+    it 'returns one pending executable migration per table' do
+      # non-active migration
+      create(:batched_background_migration, :finished)
+      # migration put on hold
+      create(:batched_background_migration, :active, on_hold_until: 10.minutes.from_now)
+      # migration not availab for the current connection
+      create(:batched_background_migration, :active, gitlab_schema: :gitlab_not_existing)
+      # active migration that is no longer on hold
+      migration_1 = create(:batched_background_migration, :active, table_name: :users, on_hold_until: 10.minutes.ago)
+      # another active migration for the same table
+      create(:batched_background_migration, :active, table_name: :users)
+      # active migration for different table
+      migration_2 = create(:batched_background_migration, :active, table_name: :projects)
+      # active migration for third table
+      create(:batched_background_migration, :active, table_name: :namespaces)
+
+      actual = described_class.active_migrations_distinct_on_table(connection: connection, limit: 2)
+
+      expect(actual).to eq([migration_1, migration_2])
+    end
+
+    it 'returns epmty collection when there are no pending executable migrations' do
+      actual = described_class.active_migrations_distinct_on_table(connection: connection, limit: 2)
+
+      expect(actual).to be_empty
+    end
   end
 
   describe '.created_after' do
diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index a87b0c1a3a8..852cc719d01 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -330,7 +330,7 @@ RSpec.describe Gitlab::Database::BatchCount do
     end
 
     it 'counts with "id" field' do
-      expect(described_class.batch_distinct_count(model, "#{column}")).to eq(2)
+      expect(described_class.batch_distinct_count(model, column.to_s)).to eq(2)
     end
 
     it 'counts with table.column field' do
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index 34370c9a21f..7dc2e0be3e5 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -23,7 +23,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
           record_type: 'A',
           interval: 60,
           disconnect_timeout: 120,
-          use_tcp: false
+          use_tcp: false,
+          max_replica_pools: nil
         )
         expect(config.pool_size).to eq(Gitlab::Database.default_pool_size)
       end
@@ -39,7 +40,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
             replica_check_interval: 3,
             hosts: %w[foo bar],
             discover: {
-              'record' => 'foo.example.com'
+              record: 'foo.example.com',
+              max_replica_pools: 5
             }
           }
         }
@@ -59,7 +61,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
           record_type: 'A',
           interval: 60,
           disconnect_timeout: 120,
-          use_tcp: false
+          use_tcp: false,
+          max_replica_pools: 5
         )
         expect(config.pool_size).to eq(4)
       end
@@ -95,7 +98,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
           record_type: 'A',
           interval: 60,
           disconnect_timeout: 120,
-          use_tcp: false
+          use_tcp: false,
+          max_replica_pools: nil
         )
         expect(config.pool_size).to eq(4)
       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 41312dbedd6..a2076f5b950 100644
--- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -53,7 +53,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   end
 
   Gitlab::Database::LoadBalancing::ConnectionProxy::NON_STICKY_READS.each do |name|
-    describe "#{name}" do
+    describe name.to_s do
       it 'runs the query on the replica' do
         expect(proxy).to receive(:read_using_load_balancer)
           .with(name, 'foo')
@@ -64,7 +64,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   end
 
   Gitlab::Database::LoadBalancing::ConnectionProxy::STICKY_WRITES.each do |name|
-    describe "#{name}" do
+    describe name.to_s do
       it 'runs the query on the primary and sticks to it' do
         session = Gitlab::Database::LoadBalancing::Session.new
 
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
new file mode 100644
index 00000000000..1a49aa2871f
--- /dev/null
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery/sampler_spec.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::Gitlab::Database::LoadBalancing::ServiceDiscovery::Sampler do
+  let(:sampler) { described_class.new(max_replica_pools: max_replica_pools, seed: 100) }
+  let(:max_replica_pools) { 3 }
+  let(:address_class) { ::Gitlab::Database::LoadBalancing::ServiceDiscovery::Address }
+  let(:addresses) do
+    [
+      address_class.new("127.0.0.1", 6432),
+      address_class.new("127.0.0.1", 6433),
+      address_class.new("127.0.0.1", 6434),
+      address_class.new("127.0.0.1", 6435),
+      address_class.new("127.0.0.2", 6432),
+      address_class.new("127.0.0.2", 6433),
+      address_class.new("127.0.0.2", 6434),
+      address_class.new("127.0.0.2", 6435)
+    ]
+  end
+
+  describe '#sample' do
+    it 'samples max_replica_pools addresses' do
+      expect(sampler.sample(addresses).count).to eq(max_replica_pools)
+    end
+
+    it 'samples random ports across all hosts' do
+      expect(sampler.sample(addresses)).to eq([
+                                                address_class.new("127.0.0.1", 6432),
+                                                address_class.new("127.0.0.2", 6435),
+                                                address_class.new("127.0.0.1", 6435)
+                                              ])
+    end
+
+    it 'returns the same answer for the same input when called multiple times' do
+      result = sampler.sample(addresses)
+      expect(sampler.sample(addresses)).to eq(result)
+      expect(sampler.sample(addresses)).to eq(result)
+    end
+
+    it 'gives a consistent answer regardless of input ordering' do
+      expect(sampler.sample(addresses.reverse)).to eq(sampler.sample(addresses))
+    end
+
+    it 'samples fairly across all hosts' do
+      # Choose a bunch of different seeds to prove that it always chooses 2
+      # different ports from each host when selecting 4
+      (1..10).each do |seed|
+        sampler = described_class.new(max_replica_pools: 4, seed: seed)
+
+        result = sampler.sample(addresses)
+
+        expect(result.count { |r| r.hostname == "127.0.0.1" }).to eq(2)
+        expect(result.count { |r| r.hostname == "127.0.0.2" }).to eq(2)
+      end
+    end
+
+    context 'when input is an empty array' do
+      it 'returns an empty array' do
+        expect(sampler.sample([])).to eq([])
+      end
+    end
+
+    context 'when there are less replicas than max_replica_pools' do
+      let(:max_replica_pools) { 100 }
+
+      it 'returns the same addresses' do
+        expect(sampler.sample(addresses)).to eq(addresses)
+      end
+    end
+
+    context 'when max_replica_pools is nil' do
+      let(:max_replica_pools) { nil }
+
+      it 'returns the same addresses' do
+        expect(sampler.sample(addresses)).to eq(addresses)
+      end
+    end
+  end
+end
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 f05910e5123..984d60e9962 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -231,10 +231,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         nameserver: 'localhost',
         port: 8600,
         record: 'foo',
-        record_type: record_type
+        record_type: record_type,
+        max_replica_pools: max_replica_pools
       )
     end
 
+    let(:max_replica_pools) { nil }
+
     let(:packet) { double(:packet, answer: [res1, res2]) }
 
     before do
@@ -266,24 +269,51 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
       let(:res1) { double(:resource, host: 'foo1.service.consul.', port: 5432, weight: 1, priority: 1, ttl: 90) }
       let(:res2) { double(:resource, host: 'foo2.service.consul.', port: 5433, weight: 1, priority: 1, ttl: 90) }
       let(:res3) { double(:resource, host: 'foo3.service.consul.', port: 5434, weight: 1, priority: 1, ttl: 90) }
-      let(:packet) { double(:packet, answer: [res1, res2, res3], additional: []) }
+      let(:res4) { double(:resource, host: 'foo4.service.consul.', port: 5432, weight: 1, priority: 1, ttl: 90) }
+      let(:packet) { double(:packet, answer: [res1, res2, res3, res4], additional: []) }
 
       before do
         expect_next_instance_of(Gitlab::Database::LoadBalancing::SrvResolver) do |resolver|
           allow(resolver).to receive(:address_for).with('foo1.service.consul.').and_return(IPAddr.new('255.255.255.0'))
           allow(resolver).to receive(:address_for).with('foo2.service.consul.').and_return(IPAddr.new('127.0.0.1'))
           allow(resolver).to receive(:address_for).with('foo3.service.consul.').and_return(nil)
+          allow(resolver).to receive(:address_for).with('foo4.service.consul.').and_return("127.0.0.2")
         end
       end
 
       it 'returns a TTL and ordered list of hosts' do
         addresses = [
           described_class::Address.new('127.0.0.1', 5433),
+          described_class::Address.new('127.0.0.2', 5432),
           described_class::Address.new('255.255.255.0', 5432)
         ]
 
         expect(service.addresses_from_dns).to eq([90, addresses])
       end
+
+      context 'when max_replica_pools is set' do
+        context 'when the number of addresses exceeds max_replica_pools' do
+          let(:max_replica_pools) { 2 }
+
+          it 'limits to max_replica_pools' do
+            expect(service.addresses_from_dns[1].count).to eq(2)
+          end
+        end
+
+        context 'when the number of addresses is less than max_replica_pools' do
+          let(:max_replica_pools) { 5 }
+
+          it 'returns all addresses' do
+            addresses = [
+              described_class::Address.new('127.0.0.1', 5433),
+              described_class::Address.new('127.0.0.2', 5432),
+              described_class::Address.new('255.255.255.0', 5432)
+            ]
+
+            expect(service.addresses_from_dns).to eq([90, addresses])
+          end
+        end
+      end
     end
 
     context 'when the resolver returns an empty response' do
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 88007de53d3..61b63016f1a 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
@@ -358,7 +358,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
   end
 
   def process_job(job)
-    Sidekiq::JobRetry.new.local(worker_class, job.to_json, 'default') do
+    Sidekiq::JobRetry.new(Sidekiq).local(worker_class, job.to_json, 'default') do
       worker_class.process_job(job)
     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
index 6026d979bcf..1eb077fe6ca 100644
--- a/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/transaction_leaking_spec.rb
@@ -4,18 +4,18 @@ require 'spec_helper'
 
 RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis, :delete do
   include StubENV
-  let(:model) { ApplicationRecord }
+  let(:model) { ActiveRecord::Base }
   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|
+    allow(Gitlab::Database::LoadBalancing::Configuration).to receive(:for_model).with(model) do |base_model|
       Gitlab::Database::LoadBalancing::Configuration.new(base_model, [db_host, db_host])
     end
 
-    Gitlab::Database::LoadBalancing::Setup.new(ApplicationRecord).setup
+    Gitlab::Database::LoadBalancing::Setup.new(model).setup
 
     model.connection.execute(<<~SQL)
       CREATE TABLE IF NOT EXISTS #{test_table_name} (id SERIAL PRIMARY KEY, value INTEGER)
@@ -30,6 +30,10 @@ RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis
     model.connection.execute(<<~SQL)
       DROP TABLE IF EXISTS #{test_table_name}
     SQL
+
+    # reset load balancing to original state
+    allow(Gitlab::Database::LoadBalancing::Configuration).to receive(:for_model).and_call_original
+    Gitlab::Database::LoadBalancing::Setup.new(model).setup
   end
 
   def execute(conn)
@@ -56,6 +60,7 @@ RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis
       conn = model.connection
 
       expect(::Gitlab::Database::LoadBalancing::Logger).to receive(:warn).with(hash_including(event: :transaction_leak))
+      expect(::Gitlab::Database::LoadBalancing::Logger).to receive(:warn).with(hash_including(event: :read_write_retry))
 
       conn.transaction do
         expect(conn).to be_transaction_open
@@ -74,6 +79,8 @@ RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis
 
       expect(::Gitlab::Database::LoadBalancing::Logger)
         .not_to receive(:warn).with(hash_including(event: :transaction_leak))
+      expect(::Gitlab::Database::LoadBalancing::Logger)
+        .to receive(:warn).with(hash_including(event: :read_write_retry))
 
       expect(conn).not_to be_transaction_open
 
@@ -105,6 +112,8 @@ RSpec.describe 'Load balancer behavior with errors inside a transaction', :redis
     it 'retries when not in a transaction' do
       expect(::Gitlab::Database::LoadBalancing::Logger)
         .not_to receive(:warn).with(hash_including(event: :transaction_leak))
+      expect(::Gitlab::Database::LoadBalancing::Logger)
+        .to receive(:warn).with(hash_including(event: :read_write_retry))
 
       expect { execute(model.connection) }.not_to raise_error
     end
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index 76dfaa74ae6..1c85abac91c 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -468,9 +468,10 @@ RSpec.describe Gitlab::Database::LoadBalancing, :suppress_gitlab_schemas_validat
           payload = event.payload
 
           assert =
-            if payload[:name] == 'SCHEMA'
+            case payload[:name]
+            when 'SCHEMA'
               false
-            elsif payload[:name] == 'SQL' # Custom query
+            when 'SQL' # Custom query
               true
             else
               keywords = %w[_test_load_balancing_test]
diff --git a/spec/lib/gitlab/database/migration_helpers/v2_spec.rb b/spec/lib/gitlab/database/migration_helpers/v2_spec.rb
index 2055dc33d48..0d75094a2fd 100644
--- a/spec/lib/gitlab/database/migration_helpers/v2_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers/v2_spec.rb
@@ -35,15 +35,15 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
       end
     end
 
-    context 'when the existing column has a default value' do
+    context 'when the existing column has a default function' do
       before do
-        migration.change_column_default :_test_table, existing_column, 'default value'
+        migration.change_column_default :_test_table, existing_column, -> { 'now()' }
       end
 
       it 'raises an error' do
         expect do
           migration.public_send(operation, :_test_table, :original, :renamed)
-        end.to raise_error("#{operation} does not currently support columns with default values")
+        end.to raise_error("#{operation} does not currently support columns with default functions")
       end
     end
 
@@ -67,6 +67,94 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
       end
     end
 
+    context 'when the existing column has a default value' do
+      before do
+        migration.change_column_default :_test_table, existing_column, 'default value'
+      end
+
+      it 'creates the renamed column, syncing existing data' do
+        existing_record_1 = model.create!(status: 0, existing_column => 'existing')
+        existing_record_2 = model.create!(status: 0)
+
+        migration.send(operation, :_test_table, :original, :renamed)
+        model.reset_column_information
+
+        expect(migration.column_exists?(:_test_table, added_column)).to eq(true)
+
+        expect(existing_record_1.reload).to have_attributes(status: 0, original: 'existing', renamed: 'existing')
+        expect(existing_record_2.reload).to have_attributes(status: 0, original: 'default value', renamed: 'default value')
+      end
+
+      it 'installs triggers to sync new data' do
+        migration.public_send(operation, :_test_table, :original, :renamed)
+        model.reset_column_information
+
+        new_record_1 = model.create!(status: 1, original: 'first')
+        new_record_2 = model.create!(status: 1, renamed: 'second')
+        new_record_3 = model.create!(status: 1)
+        new_record_4 = model.create!(status: 1)
+
+        expect(new_record_1.reload).to have_attributes(status: 1, original: 'first', renamed: 'first')
+        expect(new_record_2.reload).to have_attributes(status: 1, original: 'second', renamed: 'second')
+        expect(new_record_3.reload).to have_attributes(status: 1, original: 'default value', renamed: 'default value')
+        expect(new_record_4.reload).to have_attributes(status: 1, original: 'default value', renamed: 'default value')
+
+        new_record_1.update!(original: 'updated')
+        new_record_2.update!(renamed: nil)
+        new_record_3.update!(renamed: 'update renamed')
+        new_record_4.update!(original: 'update original')
+
+        expect(new_record_1.reload).to have_attributes(status: 1, original: 'updated', renamed: 'updated')
+        expect(new_record_2.reload).to have_attributes(status: 1, original: nil, renamed: nil)
+        expect(new_record_3.reload).to have_attributes(status: 1, original: 'update renamed', renamed: 'update renamed')
+        expect(new_record_4.reload).to have_attributes(status: 1, original: 'update original', renamed: 'update original')
+      end
+    end
+
+    context 'when the existing column has a default value that evaluates to NULL' do
+      before do
+        migration.change_column_default :_test_table, existing_column, -> { "('test' || null)" }
+      end
+
+      it 'creates the renamed column, syncing existing data' do
+        existing_record_1 = model.create!(status: 0, existing_column => 'existing')
+        existing_record_2 = model.create!(status: 0)
+
+        migration.send(operation, :_test_table, :original, :renamed)
+        model.reset_column_information
+
+        expect(migration.column_exists?(:_test_table, added_column)).to eq(true)
+
+        expect(existing_record_1.reload).to have_attributes(status: 0, original: 'existing', renamed: 'existing')
+        expect(existing_record_2.reload).to have_attributes(status: 0, original: nil, renamed: nil)
+      end
+
+      it 'installs triggers to sync new data' do
+        migration.public_send(operation, :_test_table, :original, :renamed)
+        model.reset_column_information
+
+        new_record_1 = model.create!(status: 1, original: 'first')
+        new_record_2 = model.create!(status: 1, renamed: 'second')
+        new_record_3 = model.create!(status: 1)
+        new_record_4 = model.create!(status: 1)
+
+        expect(new_record_1.reload).to have_attributes(status: 1, original: 'first', renamed: 'first')
+        expect(new_record_2.reload).to have_attributes(status: 1, original: 'second', renamed: 'second')
+        expect(new_record_3.reload).to have_attributes(status: 1, original: nil, renamed: nil)
+        expect(new_record_4.reload).to have_attributes(status: 1, original: nil, renamed: nil)
+
+        new_record_1.update!(original: 'updated')
+        new_record_2.update!(renamed: nil)
+        new_record_3.update!(renamed: 'update renamed')
+        new_record_4.update!(original: 'update original')
+
+        expect(new_record_1.reload).to have_attributes(status: 1, original: 'updated', renamed: 'updated')
+        expect(new_record_2.reload).to have_attributes(status: 1, original: nil, renamed: nil)
+        expect(new_record_3.reload).to have_attributes(status: 1, original: 'update renamed', renamed: 'update renamed')
+        expect(new_record_4.reload).to have_attributes(status: 1, original: 'update original', renamed: 'update original')
+      end
+    end
+
     it 'creates the renamed column, syncing existing data' do
       existing_record_1 = model.create!(status: 0, existing_column => 'existing')
       existing_record_2 = model.create!(status: 0, existing_column => nil)
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index bcdd5646994..65fbc8d9935 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -389,6 +389,40 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
         model.add_concurrent_index(:users, :foo)
       end
+
+      context 'when targeting a partition table' do
+        let(:schema) { 'public' }
+        let(:name) { '_test_partition_01' }
+        let(:identifier) { "#{schema}.#{name}" }
+
+        before do
+          model.execute(<<~SQL)
+            CREATE TABLE public._test_partitioned_table (
+              id serial NOT NULL,
+              partition_id serial NOT NULL,
+              PRIMARY KEY (id, partition_id)
+            ) PARTITION BY LIST(partition_id);
+
+            CREATE TABLE #{identifier} PARTITION OF public._test_partitioned_table
+            FOR VALUES IN (1);
+          SQL
+        end
+
+        context 'when allow_partition is true' do
+          it 'creates the index concurrently' do
+            expect(model).to receive(:add_index).with(:_test_partition_01, :foo, algorithm: :concurrently)
+
+            model.add_concurrent_index(:_test_partition_01, :foo, allow_partition: true)
+          end
+        end
+
+        context 'when allow_partition is not provided' do
+          it 'raises ArgumentError' do
+            expect { model.add_concurrent_index(:_test_partition_01, :foo) }
+              .to raise_error(ArgumentError, /use add_concurrent_partitioned_index/)
+          end
+        end
+      end
     end
 
     context 'inside a transaction' do
@@ -435,6 +469,37 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           model.remove_concurrent_index(:users, :foo)
         end
 
+        context 'when targeting a partition table' do
+          let(:schema) { 'public' }
+          let(:partition_table_name) { '_test_partition_01' }
+          let(:identifier) { "#{schema}.#{partition_table_name}" }
+          let(:index_name) { '_test_partitioned_index' }
+          let(:partition_index_name) { '_test_partition_01_partition_id_idx' }
+          let(:column_name) { 'partition_id' }
+
+          before do
+            model.execute(<<~SQL)
+              CREATE TABLE public._test_partitioned_table (
+                id serial NOT NULL,
+                partition_id serial NOT NULL,
+                PRIMARY KEY (id, partition_id)
+              ) PARTITION BY LIST(partition_id);
+
+              CREATE INDEX #{index_name} ON public._test_partitioned_table(#{column_name});
+
+              CREATE TABLE #{identifier} PARTITION OF public._test_partitioned_table
+              FOR VALUES IN (1);
+            SQL
+          end
+
+          context 'when dropping an index on the partition table' do
+            it 'raises ArgumentError' do
+              expect { model.remove_concurrent_index(partition_table_name, column_name) }
+                .to raise_error(ArgumentError, /use remove_concurrent_partitioned_index_by_name/)
+            end
+          end
+        end
+
         describe 'by index name' do
           before do
             allow(model).to receive(:index_exists_by_name?).with(:users, "index_x_by_y").and_return(true)
@@ -476,6 +541,36 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
             model.remove_concurrent_index_by_name(:users, "index_x_by_y")
           end
+
+          context 'when targeting a partition table' do
+            let(:schema) { 'public' }
+            let(:partition_table_name) { '_test_partition_01' }
+            let(:identifier) { "#{schema}.#{partition_table_name}" }
+            let(:index_name) { '_test_partitioned_index' }
+            let(:partition_index_name) { '_test_partition_01_partition_id_idx' }
+
+            before do
+              model.execute(<<~SQL)
+                CREATE TABLE public._test_partitioned_table (
+                  id serial NOT NULL,
+                  partition_id serial NOT NULL,
+                  PRIMARY KEY (id, partition_id)
+                ) PARTITION BY LIST(partition_id);
+
+                CREATE INDEX #{index_name} ON public._test_partitioned_table(partition_id);
+
+                CREATE TABLE #{identifier} PARTITION OF public._test_partitioned_table
+                FOR VALUES IN (1);
+              SQL
+            end
+
+            context 'when dropping an index on the partition table' do
+              it 'raises ArgumentError' do
+                expect { model.remove_concurrent_index_by_name(partition_table_name, partition_index_name) }
+                  .to raise_error(ArgumentError, /use remove_concurrent_partitioned_index_by_name/)
+              end
+            end
+          end
         end
       end
     end
@@ -1006,88 +1101,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
-  describe '#disable_statement_timeout' do
-    it 'disables statement timeouts to current transaction only' do
-      expect(model).to receive(:execute).with('SET LOCAL statement_timeout TO 0')
-
-      model.disable_statement_timeout
-    end
-
-    # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
-    context 'with real environment', :delete do
-      before do
-        model.execute("SET statement_timeout TO '20000'")
-      end
-
-      after do
-        model.execute('RESET statement_timeout')
-      end
-
-      it 'defines statement to 0 only for current transaction' do
-        expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
-
-        model.connection.transaction do
-          model.disable_statement_timeout
-          expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
-        end
-
-        expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
-      end
-
-      context 'when passing a blocks' do
-        it 'disables statement timeouts on session level and executes the block' do
-          expect(model).to receive(:execute).with('SET statement_timeout TO 0')
-          expect(model).to receive(:execute).with('RESET statement_timeout').at_least(:once)
-
-          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
-        end
-
-        # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
-        context 'with real environment', :delete do
-          before do
-            model.execute("SET statement_timeout TO '20000'")
-          end
-
-          after do
-            model.execute('RESET statement_timeout')
-          end
-
-          it 'defines statement to 0 for any code run inside the block' do
-            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
-
-            model.disable_statement_timeout do
-              model.connection.transaction do
-                expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
-              end
-
-              expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
-            end
-          end
-        end
-      end
-    end
-
-    # This spec runs without an enclosing transaction (:delete truncation method for db_cleaner)
-    context 'when the statement_timeout is already disabled', :delete do
-      before do
-        ActiveRecord::Migration.connection.execute('SET statement_timeout TO 0')
-      end
-
-      after do
-        # Use ActiveRecord::Migration.connection instead of model.execute
-        # so that this call is not counted below
-        ActiveRecord::Migration.connection.execute('RESET statement_timeout')
-      end
-
-      it 'yields control without disabling the timeout or resetting' do
-        expect(model).not_to receive(:execute).with('SET statement_timeout TO 0')
-        expect(model).not_to receive(:execute).with('RESET statement_timeout')
-
-        expect { |block| model.disable_statement_timeout(&block) }.to yield_control
-      end
-    end
-  end
-
   describe '#true_value' do
     it 'returns the appropriate value' do
       expect(model.true_value).to eq("'t'")
@@ -2006,8 +2019,116 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       end
     end
 
+    let(:same_queue_different_worker) do
+      Class.new do
+        include Sidekiq::Worker
+
+        sidekiq_options queue: 'test'
+
+        def self.name
+          'SameQueueDifferentWorkerClass'
+        end
+      end
+    end
+
+    let(:unrelated_worker) do
+      Class.new do
+        include Sidekiq::Worker
+
+        sidekiq_options queue: 'unrelated'
+
+        def self.name
+          'UnrelatedWorkerClass'
+        end
+      end
+    end
+
     before do
       stub_const(worker.name, worker)
+      stub_const(unrelated_worker.name, unrelated_worker)
+      stub_const(same_queue_different_worker.name, same_queue_different_worker)
+    end
+
+    describe '#sidekiq_remove_jobs', :clean_gitlab_redis_queues do
+      def clear_queues
+        Sidekiq::Queue.new('test').clear
+        Sidekiq::Queue.new('unrelated').clear
+        Sidekiq::RetrySet.new.clear
+        Sidekiq::ScheduledSet.new.clear
+      end
+
+      around do |example|
+        clear_queues
+        Sidekiq::Testing.disable!(&example)
+        clear_queues
+      end
+
+      it "removes all related job instances from the job class's queue" do
+        worker.perform_async
+        same_queue_different_worker.perform_async
+        unrelated_worker.perform_async
+
+        queue_we_care_about = Sidekiq::Queue.new(worker.queue)
+        unrelated_queue = Sidekiq::Queue.new(unrelated_worker.queue)
+
+        expect(queue_we_care_about.size).to eq(2)
+        expect(unrelated_queue.size).to eq(1)
+
+        model.sidekiq_remove_jobs(job_klass: worker)
+
+        expect(queue_we_care_about.size).to eq(1)
+        expect(queue_we_care_about.map(&:klass)).not_to include(worker.name)
+        expect(queue_we_care_about.map(&:klass)).to include(
+          same_queue_different_worker.name
+        )
+        expect(unrelated_queue.size).to eq(1)
+      end
+
+      context 'when job instances are in the scheduled set' do
+        it 'removes all related job instances from the scheduled set' do
+          worker.perform_in(1.hour)
+          unrelated_worker.perform_in(1.hour)
+
+          scheduled = Sidekiq::ScheduledSet.new
+
+          expect(scheduled.size).to eq(2)
+          expect(scheduled.map(&:klass)).to include(
+            worker.name,
+            unrelated_worker.name
+          )
+
+          model.sidekiq_remove_jobs(job_klass: worker)
+
+          expect(scheduled.size).to eq(1)
+          expect(scheduled.map(&:klass)).not_to include(worker.name)
+          expect(scheduled.map(&:klass)).to include(unrelated_worker.name)
+        end
+      end
+
+      context 'when job instances are in the retry set' do
+        include_context 'when handling retried jobs'
+
+        it 'removes all related job instances from the retry set' do
+          retry_in(worker, 1.hour)
+          retry_in(worker, 2.hours)
+          retry_in(worker, 3.hours)
+          retry_in(unrelated_worker, 4.hours)
+
+          retries = Sidekiq::RetrySet.new
+
+          expect(retries.size).to eq(4)
+          expect(retries.map(&:klass)).to include(
+            worker.name,
+            unrelated_worker.name
+          )
+
+          model.sidekiq_remove_jobs(job_klass: worker)
+
+          expect(retries.size).to eq(1)
+          expect(retries.map(&:klass)).not_to include(worker.name)
+          expect(retries.map(&:klass)).to include(unrelated_worker.name)
+        end
+      end
     end
 
     describe '#sidekiq_queue_length' do
@@ -2031,7 +2152,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       end
     end
 
-    describe '#migrate_sidekiq_queue' do
+    describe '#sidekiq_queue_migrate' do
       it 'migrates jobs from one sidekiq queue to another' do
         Sidekiq::Testing.disable! do
           worker.perform_async('Something', [1])
@@ -2071,6 +2192,110 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#convert_to_type_column' do
+    it 'returns the name of the temporary column used to convert to bigint' do
+      expect(model.convert_to_type_column(:id, :int, :bigint)).to eq('id_convert_int_to_bigint')
+    end
+
+    it 'returns the name of the temporary column used to convert to uuid' do
+      expect(model.convert_to_type_column(:uuid, :string, :uuid)).to eq('uuid_convert_string_to_uuid')
+    end
+  end
+
+  describe '#create_temporary_columns_and_triggers' do
+    let(:table) { :test_table }
+    let(:column) { :id }
+    let(:mappings) do
+      {
+        id: {
+          from_type: :int,
+          to_type: :bigint
+        }
+      }
+    end
+
+    let(:old_bigint_column_naming) { false }
+
+    subject do
+      model.create_temporary_columns_and_triggers(
+        table,
+        mappings,
+        old_bigint_column_naming: old_bigint_column_naming
+      )
+    end
+
+    before do
+      model.create_table table, id: false do |t|
+        t.integer :id, primary_key: true
+        t.integer :non_nullable_column, null: false
+        t.integer :nullable_column
+        t.timestamps
+      end
+    end
+
+    context 'when no mappings are provided' do
+      let(:mappings) { nil }
+
+      it 'raises an error' do
+        expect { subject }.to raise_error("No mappings for column conversion provided")
+      end
+    end
+
+    context 'when any of the mappings does not have the required keys' do
+      let(:mappings) do
+        {
+          id: {
+            from_type: :int
+          }
+        }
+      end
+
+      it 'raises an error' do
+        expect { subject }.to raise_error("Some mappings don't have required keys provided")
+      end
+    end
+
+    context 'when the target table does not exist' do
+      it 'raises an error' do
+        expect { model.create_temporary_columns_and_triggers(:non_existent_table, mappings) }.to raise_error("Table non_existent_table does not exist")
+      end
+    end
+
+    context 'when the column to migrate does not exist' do
+      let(:missing_column) { :test }
+      let(:mappings) do
+        {
+          missing_column => {
+            from_type: :int,
+            to_type: :bigint
+          }
+        }
+      end
+
+      it 'raises an error' do
+        expect { subject }.to raise_error("Column #{missing_column} does not exist on #{table}")
+      end
+    end
+
+    context 'when old_bigint_column_naming is true' do
+      let(:old_bigint_column_naming) { true }
+
+      it 'calls convert_to_bigint_column' do
+        expect(model).to receive(:convert_to_bigint_column).with(:id).and_return("id_convert_to_bigint")
+
+        subject
+      end
+    end
+
+    context 'when old_bigint_column_naming is false' do
+      it 'calls convert_to_type_column' do
+        expect(model).to receive(:convert_to_type_column).with(:id, :int, :bigint).and_return("id_convert_to_bigint")
+
+        subject
+      end
+    end
+  end
+
   describe '#initialize_conversion_of_integer_to_bigint' do
     let(:table) { :test_table }
     let(:column) { :id }
@@ -2227,7 +2452,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       let(:columns) { :id }
 
       it 'removes column, trigger, and function' do
-        temporary_column = model.convert_to_bigint_column(:id)
+        temporary_column = model.convert_to_bigint_column(columns)
         trigger_name = model.rename_trigger_name(table, :id, temporary_column)
 
         model.revert_initialize_conversion_of_integer_to_bigint(table, columns)
@@ -2420,101 +2645,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
-  describe '#ensure_batched_background_migration_is_finished' do
-    let(:job_class_name) { 'CopyColumnUsingBackgroundMigrationJob' }
-    let(:table) { :events }
-    let(:column_name) { :id }
-    let(:job_arguments) { [["id"], ["id_convert_to_bigint"], nil] }
-
-    let(:configuration) do
-      {
-        job_class_name: job_class_name,
-        table_name: table,
-        column_name: column_name,
-        job_arguments: job_arguments
-      }
-    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
-      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)
-      end
-
-      expect { ensure_batched_background_migration_is_finished }
-        .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 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" \
-            "For more information, check the documentation" \
-            "\n\n" \
-            "\thttps://docs.gitlab.com/ee/user/admin_area/monitoring/background_migrations.html#database-migrations-failing-because-of-batched-background-migration-not-finished"
-    end
-
-    it 'does not raise error when migration exists and is marked as finished' do
-      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}")
-
-      expect { ensure_batched_background_migration_is_finished }
-        .not_to raise_error
-    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|
-        expect(runner).to receive(:finalize).with(job_class_name, table, column_name, job_arguments).and_return(migration.finish!)
-      end
-
-      ensure_batched_background_migration_is_finished
-    end
-
-    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|
-          expect(runner).not_to receive(:finalize).with(job_class_name, table, column_name, job_arguments)
-        end
-
-        expect { model.ensure_batched_background_migration_is_finished(**configuration.merge(finalize: false)) }.to raise_error(RuntimeError)
-      end
-    end
-  end
-
   describe '#index_exists_by_name?' do
     it 'returns true if an index exists' do
       ActiveRecord::Migration.connection.execute(
@@ -2621,48 +2751,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
-  describe '#with_lock_retries' do
-    let(:buffer) { StringIO.new }
-    let(:in_memory_logger) { Gitlab::JsonLogger.new(buffer) }
-    let(:env) { { 'DISABLE_LOCK_RETRIES' => 'true' } }
-
-    it 'sets the migration class name in the logs' do
-      model.with_lock_retries(env: env, logger: in_memory_logger) {}
-
-      buffer.rewind
-      expect(buffer.read).to include("\"class\":\"#{model.class}\"")
-    end
-
-    where(raise_on_exhaustion: [true, false])
-
-    with_them do
-      it 'sets raise_on_exhaustion as requested' do
-        with_lock_retries = double
-        expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
-        expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: raise_on_exhaustion)
-
-        model.with_lock_retries(env: env, logger: in_memory_logger, raise_on_exhaustion: raise_on_exhaustion) {}
-      end
-    end
-
-    it 'does not raise on exhaustion by default' do
-      with_lock_retries = double
-      expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
-      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
-
-      model.with_lock_retries(env: env, logger: in_memory_logger) {}
-    end
-
-    it 'defaults to allowing subtransactions' do
-      with_lock_retries = double
-
-      expect(Gitlab::Database::WithLockRetries).to receive(:new).with(hash_including(allow_savepoints: true)).and_return(with_lock_retries)
-      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
-
-      model.with_lock_retries(env: env, logger: in_memory_logger) {}
-    end
-  end
-
   describe '#backfill_iids' do
     include MigrationsHelpers
 
@@ -2778,720 +2866,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
-  describe '#check_constraint_name' do
-    it 'returns a valid constraint name' do
-      name = model.check_constraint_name(:this_is_a_very_long_table_name,
-                                         :with_a_very_long_column_name,
-                                         :with_a_very_long_type)
-
-      expect(name).to be_an_instance_of(String)
-      expect(name).to start_with('check_')
-      expect(name.length).to eq(16)
-    end
-  end
-
-  describe '#check_constraint_exists?' do
-    before do
-      ActiveRecord::Migration.connection.execute(
-        'ALTER TABLE projects ADD CONSTRAINT check_1 CHECK (char_length(path) <= 5) NOT VALID'
-      )
-
-      ActiveRecord::Migration.connection.execute(
-        'CREATE SCHEMA new_test_schema'
-      )
-
-      ActiveRecord::Migration.connection.execute(
-        'CREATE TABLE new_test_schema.projects (id integer, name character varying)'
-      )
-
-      ActiveRecord::Migration.connection.execute(
-        'ALTER TABLE new_test_schema.projects ADD CONSTRAINT check_2 CHECK (char_length(name) <= 5)'
-      )
-    end
-
-    it 'returns true if a constraint exists' do
-      expect(model.check_constraint_exists?(:projects, 'check_1'))
-        .to be_truthy
-    end
-
-    it 'returns false if a constraint does not exist' do
-      expect(model.check_constraint_exists?(:projects, 'this_does_not_exist'))
-        .to be_falsy
-    end
-
-    it 'returns false if a constraint with the same name exists in another table' do
-      expect(model.check_constraint_exists?(:users, 'check_1'))
-        .to be_falsy
-    end
-
-    it 'returns false if a constraint with the same name exists for the same table in another schema' do
-      expect(model.check_constraint_exists?(:projects, 'check_2'))
-        .to be_falsy
-    end
-  end
-
-  describe '#add_check_constraint' do
-    before do
-      allow(model).to receive(:check_constraint_exists?).and_return(false)
-    end
-
-    context 'constraint name validation' do
-      it 'raises an error when too long' do
-        expect do
-          model.add_check_constraint(
-            :test_table,
-            'name IS NOT NULL',
-            'a' * (Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH + 1)
-          )
-        end.to raise_error(RuntimeError)
-      end
-
-      it 'does not raise error when the length is acceptable' do
-        constraint_name = 'a' * Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH
-
-        expect(model).to receive(:transaction_open?).and_return(false)
-        expect(model).to receive(:check_constraint_exists?).and_return(false)
-        expect(model).to receive(:with_lock_retries).and_call_original
-        expect(model).to receive(:execute).with(/ADD CONSTRAINT/)
-
-        model.add_check_constraint(
-          :test_table,
-          'name IS NOT NULL',
-          constraint_name,
-          validate: false
-        )
-      end
-    end
-
-    context 'inside a transaction' do
-      it 'raises an error' do
-        expect(model).to receive(:transaction_open?).and_return(true)
-
-        expect do
-          model.add_check_constraint(
-            :test_table,
-            'name IS NOT NULL',
-            'check_name_not_null'
-          )
-        end.to raise_error(RuntimeError)
-      end
-    end
-
-    context 'outside a transaction' do
-      before do
-        allow(model).to receive(:transaction_open?).and_return(false)
-      end
-
-      context 'when the constraint is already defined in the database' do
-        it 'does not create a constraint' do
-          expect(model).to receive(:check_constraint_exists?)
-                       .with(:test_table, 'check_name_not_null')
-                       .and_return(true)
-
-          expect(model).not_to receive(:execute).with(/ADD CONSTRAINT/)
-
-          # setting validate: false to only focus on the ADD CONSTRAINT command
-          model.add_check_constraint(
-            :test_table,
-            'name IS NOT NULL',
-            'check_name_not_null',
-            validate: false
-          )
-        end
-      end
-
-      context 'when the constraint is not defined in the database' do
-        it 'creates the constraint' do
-          expect(model).to receive(:with_lock_retries).and_call_original
-          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
-
-          # setting validate: false to only focus on the ADD CONSTRAINT command
-          model.add_check_constraint(
-            :test_table,
-            'char_length(name) <= 255',
-            'check_name_not_null',
-            validate: false
-          )
-        end
-      end
-
-      context 'when validate is not provided' do
-        it 'performs validation' do
-          expect(model).to receive(:check_constraint_exists?)
-                       .with(:test_table, 'check_name_not_null')
-                       .and_return(false).exactly(1)
-
-          expect(model).to receive(:disable_statement_timeout).and_call_original
-          expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
-          expect(model).to receive(:with_lock_retries).and_call_original
-          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
-
-          # we need the check constraint to exist so that the validation proceeds
-          expect(model).to receive(:check_constraint_exists?)
-                       .with(:test_table, 'check_name_not_null')
-                       .and_return(true).exactly(1)
-
-          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
-
-          model.add_check_constraint(
-            :test_table,
-            'char_length(name) <= 255',
-            'check_name_not_null'
-          )
-        end
-      end
-
-      context 'when validate is provided with a falsey value' do
-        it 'skips validation' do
-          expect(model).not_to receive(:disable_statement_timeout)
-          expect(model).to receive(:with_lock_retries).and_call_original
-          expect(model).to receive(:execute).with(/ADD CONSTRAINT/)
-          expect(model).not_to receive(:execute).with(/VALIDATE CONSTRAINT/)
-
-          model.add_check_constraint(
-            :test_table,
-            'char_length(name) <= 255',
-            'check_name_not_null',
-            validate: false
-          )
-        end
-      end
-
-      context 'when validate is provided with a truthy value' do
-        it 'performs validation' do
-          expect(model).to receive(:check_constraint_exists?)
-                       .with(:test_table, 'check_name_not_null')
-                       .and_return(false).exactly(1)
-
-          expect(model).to receive(:disable_statement_timeout).and_call_original
-          expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
-          expect(model).to receive(:with_lock_retries).and_call_original
-          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
-
-          expect(model).to receive(:check_constraint_exists?)
-                       .with(:test_table, 'check_name_not_null')
-                       .and_return(true).exactly(1)
-
-          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
-
-          model.add_check_constraint(
-            :test_table,
-            'char_length(name) <= 255',
-            'check_name_not_null',
-            validate: true
-          )
-        end
-      end
-    end
-  end
-
-  describe '#validate_check_constraint' do
-    context 'when the constraint does not exist' do
-      it 'raises an error' do
-        error_message = /Could not find check constraint "check_1" on table "test_table"/
-
-        expect(model).to receive(:check_constraint_exists?).and_return(false)
-
-        expect do
-          model.validate_check_constraint(:test_table, 'check_1')
-        end.to raise_error(RuntimeError, error_message)
-      end
-    end
-
-    context 'when the constraint exists' do
-      it 'performs validation' do
-        validate_sql = /ALTER TABLE test_table VALIDATE CONSTRAINT check_name/
-
-        expect(model).to receive(:check_constraint_exists?).and_return(true)
-        expect(model).to receive(:disable_statement_timeout).and_call_original
-        expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-        expect(model).to receive(:execute).with(/SET statement_timeout TO/)
-        expect(model).to receive(:execute).ordered.with(validate_sql)
-        expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
-
-        model.validate_check_constraint(:test_table, 'check_name')
-      end
-    end
-  end
-
-  describe '#remove_check_constraint' do
-    before do
-      allow(model).to receive(:transaction_open?).and_return(false)
-    end
-
-    it 'removes the constraint' do
-      drop_sql = /ALTER TABLE test_table\s+DROP CONSTRAINT IF EXISTS check_name/
-
-      expect(model).to receive(:with_lock_retries).and_call_original
-      expect(model).to receive(:execute).with(drop_sql)
-
-      model.remove_check_constraint(:test_table, 'check_name')
-    end
-  end
-
-  describe '#copy_check_constraints' do
-    context 'inside a transaction' do
-      it 'raises an error' do
-        expect(model).to receive(:transaction_open?).and_return(true)
-
-        expect do
-          model.copy_check_constraints(:test_table, :old_column, :new_column)
-        end.to raise_error(RuntimeError)
-      end
-    end
-
-    context 'outside a transaction' do
-      before do
-        allow(model).to receive(:transaction_open?).and_return(false)
-        allow(model).to receive(:column_exists?).and_return(true)
-      end
-
-      let(:old_column_constraints) do
-        [
-          {
-            'schema_name' => 'public',
-            'table_name' => 'test_table',
-            'column_name' => 'old_column',
-            'constraint_name' => 'check_d7d49d475d',
-            'constraint_def' => 'CHECK ((old_column IS NOT NULL))'
-          },
-          {
-            'schema_name' => 'public',
-            'table_name' => 'test_table',
-            'column_name' => 'old_column',
-            'constraint_name' => 'check_48560e521e',
-            'constraint_def' => 'CHECK ((char_length(old_column) <= 255))'
-          },
-          {
-            'schema_name' => 'public',
-            'table_name' => 'test_table',
-            'column_name' => 'old_column',
-            'constraint_name' => 'custom_check_constraint',
-            'constraint_def' => 'CHECK (((old_column IS NOT NULL) AND (another_column IS NULL)))'
-          },
-          {
-            'schema_name' => 'public',
-            'table_name' => 'test_table',
-            'column_name' => 'old_column',
-            'constraint_name' => 'not_valid_check_constraint',
-            'constraint_def' => 'CHECK ((old_column IS NOT NULL)) NOT VALID'
-          }
-        ]
-      end
-
-      it 'copies check constraints from one column to another' do
-        allow(model).to receive(:check_constraints_for)
-        .with(:test_table, :old_column, schema: nil)
-          .and_return(old_column_constraints)
-
-        allow(model).to receive(:not_null_constraint_name).with(:test_table, :new_column)
-          .and_return('check_1')
-
-        allow(model).to receive(:text_limit_name).with(:test_table, :new_column)
-          .and_return('check_2')
-
-        allow(model).to receive(:check_constraint_name)
-          .with(:test_table, :new_column, 'copy_check_constraint')
-          .and_return('check_3')
-
-        expect(model).to receive(:add_check_constraint)
-          .with(
-            :test_table,
-            '(new_column IS NOT NULL)',
-            'check_1',
-            validate: true
-          ).once
-
-        expect(model).to receive(:add_check_constraint)
-          .with(
-            :test_table,
-            '(char_length(new_column) <= 255)',
-            'check_2',
-            validate: true
-          ).once
-
-        expect(model).to receive(:add_check_constraint)
-          .with(
-            :test_table,
-            '((new_column IS NOT NULL) AND (another_column IS NULL))',
-            'check_3',
-            validate: true
-          ).once
-
-        expect(model).to receive(:add_check_constraint)
-          .with(
-            :test_table,
-            '(new_column IS NOT NULL)',
-            'check_1',
-            validate: false
-          ).once
-
-        model.copy_check_constraints(:test_table, :old_column, :new_column)
-      end
-
-      it 'does nothing if there are no constraints defined for the old column' do
-        allow(model).to receive(:check_constraints_for)
-        .with(:test_table, :old_column, schema: nil)
-          .and_return([])
-
-        expect(model).not_to receive(:add_check_constraint)
-
-        model.copy_check_constraints(:test_table, :old_column, :new_column)
-      end
-
-      it 'raises an error when the orginating column does not exist' do
-        allow(model).to receive(:column_exists?).with(:test_table, :old_column).and_return(false)
-
-        error_message = /Column old_column does not exist on test_table/
-
-        expect do
-          model.copy_check_constraints(:test_table, :old_column, :new_column)
-        end.to raise_error(RuntimeError, error_message)
-      end
-
-      it 'raises an error when the target column does not exist' do
-        allow(model).to receive(:column_exists?).with(:test_table, :new_column).and_return(false)
-
-        error_message = /Column new_column does not exist on test_table/
-
-        expect do
-          model.copy_check_constraints(:test_table, :old_column, :new_column)
-        end.to raise_error(RuntimeError, error_message)
-      end
-    end
-  end
-
-  describe '#add_text_limit' do
-    context 'when it is called with the default options' do
-      it 'calls add_check_constraint with an infered constraint name and validate: true' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'max_length')
-        check = "char_length(name) <= 255"
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:add_check_constraint)
-                     .with(:test_table, check, constraint_name, validate: true)
-
-        model.add_text_limit(:test_table, :name, 255)
-      end
-    end
-
-    context 'when all parameters are provided' do
-      it 'calls add_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_limit'
-        check = "char_length(name) <= 255"
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:add_check_constraint)
-                     .with(:test_table, check, constraint_name, validate: false)
-
-        model.add_text_limit(
-          :test_table,
-          :name,
-          255,
-          constraint_name: constraint_name,
-          validate: false
-        )
-      end
-    end
-  end
-
-  describe '#validate_text_limit' do
-    context 'when constraint_name is not provided' do
-      it 'calls validate_check_constraint with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'max_length')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:validate_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.validate_text_limit(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls validate_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_limit'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:validate_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.validate_text_limit(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#remove_text_limit' do
-    context 'when constraint_name is not provided' do
-      it 'calls remove_check_constraint with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'max_length')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:remove_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.remove_text_limit(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls remove_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_limit'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:remove_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.remove_text_limit(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#check_text_limit_exists?' do
-    context 'when constraint_name is not provided' do
-      it 'calls check_constraint_exists? with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'max_length')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:check_constraint_exists?)
-                     .with(:test_table, constraint_name)
-
-        model.check_text_limit_exists?(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls check_constraint_exists? with the correct parameters' do
-        constraint_name = 'check_name_limit'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:check_constraint_exists?)
-                     .with(:test_table, constraint_name)
-
-        model.check_text_limit_exists?(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#add_not_null_constraint' do
-    context 'when it is called with the default options' do
-      it 'calls add_check_constraint with an infered constraint name and validate: true' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'not_null')
-        check = "name IS NOT NULL"
-
-        expect(model).to receive(:column_is_nullable?).and_return(true)
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:add_check_constraint)
-                     .with(:test_table, check, constraint_name, validate: true)
-
-        model.add_not_null_constraint(:test_table, :name)
-      end
-    end
-
-    context 'when all parameters are provided' do
-      it 'calls add_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_not_null'
-        check = "name IS NOT NULL"
-
-        expect(model).to receive(:column_is_nullable?).and_return(true)
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:add_check_constraint)
-                     .with(:test_table, check, constraint_name, validate: false)
-
-        model.add_not_null_constraint(
-          :test_table,
-          :name,
-          constraint_name: constraint_name,
-          validate: false
-        )
-      end
-    end
-
-    context 'when the column is defined as NOT NULL' do
-      it 'does not add a check constraint' do
-        expect(model).to receive(:column_is_nullable?).and_return(false)
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).not_to receive(:add_check_constraint)
-
-        model.add_not_null_constraint(:test_table, :name)
-      end
-    end
-  end
-
-  describe '#validate_not_null_constraint' do
-    context 'when constraint_name is not provided' do
-      it 'calls validate_check_constraint with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'not_null')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:validate_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.validate_not_null_constraint(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls validate_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_not_null'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:validate_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.validate_not_null_constraint(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#remove_not_null_constraint' do
-    context 'when constraint_name is not provided' do
-      it 'calls remove_check_constraint with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'not_null')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:remove_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.remove_not_null_constraint(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls remove_check_constraint with the correct parameters' do
-        constraint_name = 'check_name_not_null'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:remove_check_constraint)
-                     .with(:test_table, constraint_name)
-
-        model.remove_not_null_constraint(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#check_not_null_constraint_exists?' do
-    context 'when constraint_name is not provided' do
-      it 'calls check_constraint_exists? with an infered constraint name' do
-        constraint_name = model.check_constraint_name(:test_table,
-                                                      :name,
-                                                      'not_null')
-
-        expect(model).to receive(:check_constraint_name).and_call_original
-        expect(model).to receive(:check_constraint_exists?)
-                     .with(:test_table, constraint_name)
-
-        model.check_not_null_constraint_exists?(:test_table, :name)
-      end
-    end
-
-    context 'when constraint_name is provided' do
-      it 'calls check_constraint_exists? with the correct parameters' do
-        constraint_name = 'check_name_not_null'
-
-        expect(model).not_to receive(:check_constraint_name)
-        expect(model).to receive(:check_constraint_exists?)
-                     .with(:test_table, constraint_name)
-
-        model.check_not_null_constraint_exists?(:test_table, :name, constraint_name: constraint_name)
-      end
-    end
-  end
-
-  describe '#create_extension' do
-    subject { model.create_extension(extension) }
-
-    let(:extension) { :btree_gist }
-
-    it 'executes CREATE EXTENSION statement' do
-      expect(model).to receive(:execute).with(/CREATE EXTENSION IF NOT EXISTS #{extension}/)
-
-      subject
-    end
-
-    context 'without proper permissions' do
-      before do
-        allow(model).to receive(:execute)
-          .with(/CREATE EXTENSION IF NOT EXISTS #{extension}/)
-          .and_raise(ActiveRecord::StatementInvalid, 'InsufficientPrivilege: permission denied')
-      end
-
-      it 'raises an exception and prints an error message' do
-        expect { subject }
-          .to output(/user is not allowed/).to_stderr
-          .and raise_error(ActiveRecord::StatementInvalid, /InsufficientPrivilege/)
-      end
-    end
-  end
-
-  describe '#drop_extension' do
-    subject { model.drop_extension(extension) }
-
-    let(:extension) { 'btree_gist' }
-
-    it 'executes CREATE EXTENSION statement' do
-      expect(model).to receive(:execute).with(/DROP EXTENSION IF EXISTS #{extension}/)
-
-      subject
-    end
-
-    context 'without proper permissions' do
-      before do
-        allow(model).to receive(:execute)
-          .with(/DROP EXTENSION IF EXISTS #{extension}/)
-          .and_raise(ActiveRecord::StatementInvalid, 'InsufficientPrivilege: permission denied')
-      end
-
-      it 'raises an exception and prints an error message' do
-        expect { subject }
-          .to output(/user is not allowed/).to_stderr
-          .and raise_error(ActiveRecord::StatementInvalid, /InsufficientPrivilege/)
-      end
-    end
-  end
-
-  describe '#rename_constraint' do
-    it "executes the statement to rename constraint" do
-      expect(model).to receive(:execute).with /ALTER TABLE "test_table"\nRENAME CONSTRAINT "fk_old_name" TO "fk_new_name"/
-
-      model.rename_constraint(:test_table, :fk_old_name, :fk_new_name)
-    end
-  end
-
-  describe '#drop_constraint' do
-    it "executes the statement to drop the constraint" do
-      expect(model).to receive(:execute).with("ALTER TABLE \"test_table\" DROP CONSTRAINT \"constraint_name\" CASCADE\n")
-
-      model.drop_constraint(:test_table, :constraint_name, cascade: true)
-    end
-
-    context 'when cascade option is false' do
-      it "executes the statement to drop the constraint without cascade" do
-        expect(model).to receive(:execute).with("ALTER TABLE \"test_table\" DROP CONSTRAINT \"constraint_name\" \n")
-
-        model.drop_constraint(:test_table, :constraint_name, cascade: false)
-      end
-    end
-  end
-
   describe '#add_primary_key_using_index' do
     it "executes the statement to add the primary key" do
       expect(model).to receive(:execute).with /ALTER TABLE "test_table" ADD CONSTRAINT "old_name" PRIMARY KEY USING INDEX "new_name"/
@@ -3558,4 +2932,36 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       model.add_sequence(:test_table, :test_column, :test_table_id_seq, 1)
     end
   end
+
+  describe "#partition?" do
+    subject { model.partition?(table_name) }
+
+    let(:table_name) { 'ci_builds_metadata' }
+
+    context "when a partition table exist" do
+      context 'when the view postgres_partitions exists' do
+        it 'calls the view', :aggregate_failures do
+          expect(Gitlab::Database::PostgresPartition).to receive(:partition_exists?).with(table_name).and_call_original
+          expect(subject).to be_truthy
+        end
+      end
+
+      context 'when the view postgres_partitions does not exist' do
+        before do
+          allow(model).to receive(:view_exists?).and_return(false)
+        end
+
+        it 'does not call the view', :aggregate_failures do
+          expect(Gitlab::Database::PostgresPartition).to receive(:legacy_partition_exists?).with(table_name).and_call_original
+          expect(subject).to be_truthy
+        end
+      end
+    end
+
+    context "when a partition table does not exist" do
+      let(:table_name) { 'partition_does_not_exist' }
+
+      it { is_expected.to be_falsey }
+    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 f21f1ac5e52..d4fff947c29 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -14,9 +14,6 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
   shared_examples_for 'helpers that enqueue background migrations' do |worker_class, connection_class, tracking_database|
     before do
       allow(model).to receive(:tracking_database).and_return(tracking_database)
-
-      # Due to lib/gitlab/database/load_balancing/configuration.rb:92 requiring RequestStore
-      # we cannot use stub_feature_flags(force_no_sharing_primary_model: true)
       allow(connection_class.connection.load_balancer.configuration)
         .to receive(:use_dedicated_connection?).and_return(true)
 
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 a2f6e6b43ed..3e249b14f2e 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
@@ -425,4 +425,99 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
       end
     end
   end
+
+  describe '#ensure_batched_background_migration_is_finished' do
+    let(:job_class_name) { 'CopyColumnUsingBackgroundMigrationJob' }
+    let(:table) { :events }
+    let(:column_name) { :id }
+    let(:job_arguments) { [["id"], ["id_convert_to_bigint"], nil] }
+
+    let(:configuration) do
+      {
+        job_class_name: job_class_name,
+        table_name: table,
+        column_name: column_name,
+        job_arguments: job_arguments
+      }
+    end
+
+    let(:migration_attributes) do
+      configuration.merge(gitlab_schema: Gitlab::Database.gitlab_schemas_for_connection(migration.connection).first)
+    end
+
+    before do
+      allow(migration).to receive(:transaction_open?).and_return(false)
+    end
+
+    subject(:ensure_batched_background_migration_is_finished) { migration.ensure_batched_background_migration_is_finished(**configuration) }
+
+    it 'raises an error when migration exists and is not marked as finished' do
+      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)
+      end
+
+      expect { ensure_batched_background_migration_is_finished }
+        .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 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" \
+            "For more information, check the documentation" \
+            "\n\n" \
+            "\thttps://docs.gitlab.com/ee/user/admin_area/monitoring/background_migrations.html#database-migrations-failing-because-of-batched-background-migration-not-finished"
+    end
+
+    it 'does not raise error when migration exists and is marked as finished' do
+      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}")
+
+      expect { ensure_batched_background_migration_is_finished }
+        .not_to raise_error
+    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|
+        expect(runner).to receive(:finalize).with(job_class_name, table, column_name, job_arguments).and_return(migration.finish!)
+      end
+
+      ensure_batched_background_migration_is_finished
+    end
+
+    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|
+          expect(runner).not_to receive(:finalize).with(job_class_name, table, column_name, job_arguments)
+        end
+
+        expect { migration.ensure_batched_background_migration_is_finished(**configuration.merge(finalize: false)) }.to raise_error(RuntimeError)
+      end
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/migrations/constraints_helpers_spec.rb b/spec/lib/gitlab/database/migrations/constraints_helpers_spec.rb
new file mode 100644
index 00000000000..6848fc85aa1
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/constraints_helpers_spec.rb
@@ -0,0 +1,679 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::ConstraintsHelpers do
+  let(:model) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  before do
+    allow(model).to receive(:puts)
+  end
+
+  describe '#check_constraint_name' do
+    it 'returns a valid constraint name' do
+      name = model.check_constraint_name(:this_is_a_very_long_table_name,
+                                         :with_a_very_long_column_name,
+                                         :with_a_very_long_type)
+
+      expect(name).to be_an_instance_of(String)
+      expect(name).to start_with('check_')
+      expect(name.length).to eq(16)
+    end
+  end
+
+  describe '#check_constraint_exists?' do
+    before do
+      ActiveRecord::Migration.connection.execute(
+        'ALTER TABLE projects ADD CONSTRAINT check_1 CHECK (char_length(path) <= 5) NOT VALID'
+      )
+
+      ActiveRecord::Migration.connection.execute(
+        'CREATE SCHEMA new_test_schema'
+      )
+
+      ActiveRecord::Migration.connection.execute(
+        'CREATE TABLE new_test_schema.projects (id integer, name character varying)'
+      )
+
+      ActiveRecord::Migration.connection.execute(
+        'ALTER TABLE new_test_schema.projects ADD CONSTRAINT check_2 CHECK (char_length(name) <= 5)'
+      )
+    end
+
+    it 'returns true if a constraint exists' do
+      expect(model)
+        .to be_check_constraint_exists(:projects, 'check_1')
+    end
+
+    it 'returns false if a constraint does not exist' do
+      expect(model)
+        .not_to be_check_constraint_exists(:projects, 'this_does_not_exist')
+    end
+
+    it 'returns false if a constraint with the same name exists in another table' do
+      expect(model)
+        .not_to be_check_constraint_exists(:users, 'check_1')
+    end
+
+    it 'returns false if a constraint with the same name exists for the same table in another schema' do
+      expect(model)
+        .not_to be_check_constraint_exists(:projects, 'check_2')
+    end
+  end
+
+  describe '#add_check_constraint' do
+    before do
+      allow(model).to receive(:check_constraint_exists?).and_return(false)
+    end
+
+    context 'when constraint name validation' do
+      it 'raises an error when too long' do
+        expect do
+          model.add_check_constraint(
+            :test_table,
+            'name IS NOT NULL',
+            'a' * (Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH + 1)
+          )
+        end.to raise_error(RuntimeError)
+      end
+
+      it 'does not raise error when the length is acceptable' do
+        constraint_name = 'a' * Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH
+
+        expect(model).to receive(:transaction_open?).and_return(false)
+        expect(model).to receive(:check_constraint_exists?).and_return(false)
+        expect(model).to receive(:with_lock_retries).and_call_original
+        expect(model).to receive(:execute).with(/ADD CONSTRAINT/)
+
+        model.add_check_constraint(
+          :test_table,
+          'name IS NOT NULL',
+          constraint_name,
+          validate: false
+        )
+      end
+    end
+
+    context 'when inside a transaction' do
+      it 'raises an error' do
+        expect(model).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          model.add_check_constraint(
+            :test_table,
+            'name IS NOT NULL',
+            'check_name_not_null'
+          )
+        end.to raise_error(RuntimeError)
+      end
+    end
+
+    context 'when outside a transaction' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      context 'when the constraint is already defined in the database' do
+        it 'does not create a constraint' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true)
+
+          expect(model).not_to receive(:execute).with(/ADD CONSTRAINT/)
+
+          # setting validate: false to only focus on the ADD CONSTRAINT command
+          model.add_check_constraint(
+            :test_table,
+            'name IS NOT NULL',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when the constraint is not defined in the database' do
+        it 'creates the constraint' do
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          # setting validate: false to only focus on the ADD CONSTRAINT command
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when validate is not provided' do
+        it 'performs validation' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(false).exactly(1)
+
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:statement_timeout_disabled?).and_return(false)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          # we need the check constraint to exist so that the validation proceeds
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true).exactly(1)
+
+          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null'
+          )
+        end
+      end
+
+      context 'when validate is provided with a falsey value' do
+        it 'skips validation' do
+          expect(model).not_to receive(:disable_statement_timeout)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT/)
+          expect(model).not_to receive(:execute).with(/VALIDATE CONSTRAINT/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when validate is provided with a truthy value' do
+        it 'performs validation' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(false).exactly(1)
+
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:statement_timeout_disabled?).and_return(false)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true).exactly(1)
+
+          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: true
+          )
+        end
+      end
+    end
+  end
+
+  describe '#validate_check_constraint' do
+    context 'when the constraint does not exist' do
+      it 'raises an error' do
+        error_message = /Could not find check constraint "check_1" on table "test_table"/
+
+        expect(model).to receive(:check_constraint_exists?).and_return(false)
+
+        expect do
+          model.validate_check_constraint(:test_table, 'check_1')
+        end.to raise_error(RuntimeError, error_message)
+      end
+    end
+
+    context 'when the constraint exists' do
+      it 'performs validation' do
+        validate_sql = /ALTER TABLE test_table VALIDATE CONSTRAINT check_name/
+
+        expect(model).to receive(:check_constraint_exists?).and_return(true)
+        expect(model).to receive(:disable_statement_timeout).and_call_original
+        expect(model).to receive(:statement_timeout_disabled?).and_return(false)
+        expect(model).to receive(:execute).with(/SET statement_timeout TO/)
+        expect(model).to receive(:execute).ordered.with(validate_sql)
+        expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
+
+        model.validate_check_constraint(:test_table, 'check_name')
+      end
+    end
+  end
+
+  describe '#remove_check_constraint' do
+    before do
+      allow(model).to receive(:transaction_open?).and_return(false)
+    end
+
+    it 'removes the constraint' do
+      drop_sql = /ALTER TABLE test_table\s+DROP CONSTRAINT IF EXISTS check_name/
+
+      expect(model).to receive(:with_lock_retries).and_call_original
+      expect(model).to receive(:execute).with(drop_sql)
+
+      model.remove_check_constraint(:test_table, 'check_name')
+    end
+  end
+
+  describe '#copy_check_constraints' do
+    context 'when inside a transaction' do
+      it 'raises an error' do
+        expect(model).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          model.copy_check_constraints(:test_table, :old_column, :new_column)
+        end.to raise_error(RuntimeError)
+      end
+    end
+
+    context 'when outside a transaction' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+        allow(model).to receive(:column_exists?).and_return(true)
+      end
+
+      let(:old_column_constraints) do
+        [
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'check_d7d49d475d',
+            'constraint_def' => 'CHECK ((old_column IS NOT NULL))'
+          },
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'check_48560e521e',
+            'constraint_def' => 'CHECK ((char_length(old_column) <= 255))'
+          },
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'custom_check_constraint',
+            'constraint_def' => 'CHECK (((old_column IS NOT NULL) AND (another_column IS NULL)))'
+          },
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'not_valid_check_constraint',
+            'constraint_def' => 'CHECK ((old_column IS NOT NULL)) NOT VALID'
+          }
+        ]
+      end
+
+      it 'copies check constraints from one column to another' do
+        allow(model).to receive(:check_constraints_for)
+        .with(:test_table, :old_column, schema: nil)
+          .and_return(old_column_constraints)
+
+        allow(model).to receive(:not_null_constraint_name).with(:test_table, :new_column)
+          .and_return('check_1')
+
+        allow(model).to receive(:text_limit_name).with(:test_table, :new_column)
+          .and_return('check_2')
+
+        allow(model).to receive(:check_constraint_name)
+          .with(:test_table, :new_column, 'copy_check_constraint')
+          .and_return('check_3')
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '(new_column IS NOT NULL)',
+            'check_1',
+            validate: true
+          ).once
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '(char_length(new_column) <= 255)',
+            'check_2',
+            validate: true
+          ).once
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '((new_column IS NOT NULL) AND (another_column IS NULL))',
+            'check_3',
+            validate: true
+          ).once
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '(new_column IS NOT NULL)',
+            'check_1',
+            validate: false
+          ).once
+
+        model.copy_check_constraints(:test_table, :old_column, :new_column)
+      end
+
+      it 'does nothing if there are no constraints defined for the old column' do
+        allow(model).to receive(:check_constraints_for)
+        .with(:test_table, :old_column, schema: nil)
+          .and_return([])
+
+        expect(model).not_to receive(:add_check_constraint)
+
+        model.copy_check_constraints(:test_table, :old_column, :new_column)
+      end
+
+      it 'raises an error when the orginating column does not exist' do
+        allow(model).to receive(:column_exists?).with(:test_table, :old_column).and_return(false)
+
+        error_message = /Column old_column does not exist on test_table/
+
+        expect do
+          model.copy_check_constraints(:test_table, :old_column, :new_column)
+        end.to raise_error(RuntimeError, error_message)
+      end
+
+      it 'raises an error when the target column does not exist' do
+        allow(model).to receive(:column_exists?).with(:test_table, :new_column).and_return(false)
+
+        error_message = /Column new_column does not exist on test_table/
+
+        expect do
+          model.copy_check_constraints(:test_table, :old_column, :new_column)
+        end.to raise_error(RuntimeError, error_message)
+      end
+    end
+  end
+
+  describe '#add_text_limit' do
+    context 'when it is called with the default options' do
+      it 'calls add_check_constraint with an infered constraint name and validate: true' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+        check = "char_length(name) <= 255"
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: true)
+
+        model.add_text_limit(:test_table, :name, 255)
+      end
+    end
+
+    context 'when all parameters are provided' do
+      it 'calls add_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+        check = "char_length(name) <= 255"
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: false)
+
+        model.add_text_limit(
+          :test_table,
+          :name,
+          255,
+          constraint_name: constraint_name,
+          validate: false
+        )
+      end
+    end
+  end
+
+  describe '#validate_text_limit' do
+    context 'when constraint_name is not provided' do
+      it 'calls validate_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_text_limit(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls validate_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_text_limit(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#remove_text_limit' do
+    context 'when constraint_name is not provided' do
+      it 'calls remove_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_text_limit(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls remove_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_text_limit(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#check_text_limit_exists?' do
+    context 'when constraint_name is not provided' do
+      it 'calls check_constraint_exists? with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_text_limit_exists?(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls check_constraint_exists? with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_text_limit_exists?(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#add_not_null_constraint' do
+    context 'when it is called with the default options' do
+      it 'calls add_check_constraint with an infered constraint name and validate: true' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'not_null')
+        check = "name IS NOT NULL"
+
+        expect(model).to receive(:column_is_nullable?).and_return(true)
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: true)
+
+        model.add_not_null_constraint(:test_table, :name)
+      end
+    end
+
+    context 'when all parameters are provided' do
+      it 'calls add_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_not_null'
+        check = "name IS NOT NULL"
+
+        expect(model).to receive(:column_is_nullable?).and_return(true)
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: false)
+
+        model.add_not_null_constraint(
+          :test_table,
+          :name,
+          constraint_name: constraint_name,
+          validate: false
+        )
+      end
+    end
+
+    context 'when the column is defined as NOT NULL' do
+      it 'does not add a check constraint' do
+        expect(model).to receive(:column_is_nullable?).and_return(false)
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).not_to receive(:add_check_constraint)
+
+        model.add_not_null_constraint(:test_table, :name)
+      end
+    end
+  end
+
+  describe '#validate_not_null_constraint' do
+    context 'when constraint_name is not provided' do
+      it 'calls validate_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'not_null')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_not_null_constraint(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls validate_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_not_null'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_not_null_constraint(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#remove_not_null_constraint' do
+    context 'when constraint_name is not provided' do
+      it 'calls remove_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'not_null')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_not_null_constraint(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls remove_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_not_null'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_not_null_constraint(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#check_not_null_constraint_exists?' do
+    context 'when constraint_name is not provided' do
+      it 'calls check_constraint_exists? with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'not_null')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_not_null_constraint_exists?(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls check_constraint_exists? with the correct parameters' do
+        constraint_name = 'check_name_not_null'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_not_null_constraint_exists?(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#rename_constraint' do
+    it "executes the statement to rename constraint" do
+      expect(model).to receive(:execute).with(
+        /ALTER TABLE "test_table"\nRENAME CONSTRAINT "fk_old_name" TO "fk_new_name"/
+      )
+
+      model.rename_constraint(:test_table, :fk_old_name, :fk_new_name)
+    end
+  end
+
+  describe '#drop_constraint' do
+    it "executes the statement to drop the constraint" do
+      expect(model).to receive(:execute).with(
+        "ALTER TABLE \"test_table\" DROP CONSTRAINT \"constraint_name\" CASCADE\n"
+      )
+
+      model.drop_constraint(:test_table, :constraint_name, cascade: true)
+    end
+
+    context 'when cascade option is false' do
+      it "executes the statement to drop the constraint without cascade" do
+        expect(model).to receive(:execute).with("ALTER TABLE \"test_table\" DROP CONSTRAINT \"constraint_name\" \n")
+
+        model.drop_constraint(:test_table, :constraint_name, cascade: false)
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migrations/extension_helpers_spec.rb b/spec/lib/gitlab/database/migrations/extension_helpers_spec.rb
new file mode 100644
index 00000000000..fb29e06bc01
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/extension_helpers_spec.rb
@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::ExtensionHelpers do
+  let(:model) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  before do
+    allow(model).to receive(:puts)
+  end
+
+  describe '#create_extension' do
+    subject { model.create_extension(extension) }
+
+    let(:extension) { :btree_gist }
+
+    it 'executes CREATE EXTENSION statement' do
+      expect(model).to receive(:execute).with(/CREATE EXTENSION IF NOT EXISTS #{extension}/)
+
+      subject
+    end
+
+    context 'without proper permissions' do
+      before do
+        allow(model).to receive(:execute)
+          .with(/CREATE EXTENSION IF NOT EXISTS #{extension}/)
+          .and_raise(ActiveRecord::StatementInvalid, 'InsufficientPrivilege: permission denied')
+      end
+
+      it 'raises an exception and prints an error message' do
+        expect { subject }
+          .to output(/user is not allowed/).to_stderr
+          .and raise_error(ActiveRecord::StatementInvalid, /InsufficientPrivilege/)
+      end
+    end
+  end
+
+  describe '#drop_extension' do
+    subject { model.drop_extension(extension) }
+
+    let(:extension) { 'btree_gist' }
+
+    it 'executes CREATE EXTENSION statement' do
+      expect(model).to receive(:execute).with(/DROP EXTENSION IF EXISTS #{extension}/)
+
+      subject
+    end
+
+    context 'without proper permissions' do
+      before do
+        allow(model).to receive(:execute)
+          .with(/DROP EXTENSION IF EXISTS #{extension}/)
+          .and_raise(ActiveRecord::StatementInvalid, 'InsufficientPrivilege: permission denied')
+      end
+
+      it 'raises an exception and prints an error message' do
+        expect { subject }
+          .to output(/user is not allowed/).to_stderr
+          .and raise_error(ActiveRecord::StatementInvalid, /InsufficientPrivilege/)
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migrations/lock_retries_helpers_spec.rb b/spec/lib/gitlab/database/migrations/lock_retries_helpers_spec.rb
new file mode 100644
index 00000000000..a8739f6758f
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/lock_retries_helpers_spec.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::LockRetriesHelpers do
+  let(:model) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  describe '#with_lock_retries' do
+    let(:buffer) { StringIO.new }
+    let(:in_memory_logger) { Gitlab::JsonLogger.new(buffer) }
+    let(:env) { { 'DISABLE_LOCK_RETRIES' => 'true' } }
+
+    it 'sets the migration class name in the logs' do
+      model.with_lock_retries(env: env, logger: in_memory_logger) {}
+
+      buffer.rewind
+      expect(buffer.read).to include("\"class\":\"#{model.class}\"")
+    end
+
+    where(raise_on_exhaustion: [true, false])
+
+    with_them do
+      it 'sets raise_on_exhaustion as requested' do
+        with_lock_retries = double
+        expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+        expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: raise_on_exhaustion)
+
+        model.with_lock_retries(env: env, logger: in_memory_logger, raise_on_exhaustion: raise_on_exhaustion) {}
+      end
+    end
+
+    it 'does not raise on exhaustion by default' do
+      with_lock_retries = double
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) {}
+    end
+
+    it 'defaults to allowing subtransactions' do
+      with_lock_retries = double
+
+      expect(Gitlab::Database::WithLockRetries)
+        .to receive(:new).with(hash_including(allow_savepoints: true)).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) {}
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migrations/runner_spec.rb b/spec/lib/gitlab/database/migrations/runner_spec.rb
index f364ebfa522..bd382547689 100644
--- a/spec/lib/gitlab/database/migrations/runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/runner_spec.rb
@@ -2,26 +2,65 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Runner, :reestablished_active_record_base do
-  include Database::MultipleDatabases
-
   let(:base_result_dir) { Pathname.new(Dir.mktmpdir) }
 
   let(:migration_runs) { [] } # This list gets populated as the runner tries to run migrations
 
   # Tests depend on all of these lists being sorted in the order migrations would be applied
-  let(:applied_migrations_other_branches) { [double(ActiveRecord::Migration, version: 1, name: 'migration_complete_other_branch')] }
+  let(:applied_migrations_other_branches) do
+    [
+      double(
+        ActiveRecord::Migration,
+        version: 1,
+        name: 'migration_complete_other_branch',
+        filename: 'db/migrate/1_migration_complete_other_branch.rb'
+      )
+    ]
+  end
 
   let(:applied_migrations_this_branch) do
     [
-      double(ActiveRecord::Migration, version: 2, name: 'older_migration_complete_this_branch'),
-      double(ActiveRecord::Migration, version: 3, name: 'newer_migration_complete_this_branch')
+      double(
+        ActiveRecord::Migration,
+        version: 2,
+        name: 'older_migration_complete_this_branch',
+        filename: 'db/migrate/2_older_migration_complete_this_branch.rb'
+      ),
+      double(
+        ActiveRecord::Migration,
+        version: 3,
+        name: 'post_migration_complete_this_branch',
+        filename: 'db/post_migrate/3_post_migration_complete_this_branch.rb'
+      ),
+      double(
+        ActiveRecord::Migration,
+        version: 4,
+        name: 'newer_migration_complete_this_branch',
+        filename: 'db/migrate/4_newer_migration_complete_this_branch.rb'
+      )
     ].sort_by(&:version)
   end
 
   let(:pending_migrations) do
     [
-      double(ActiveRecord::Migration, version: 4, name: 'older_migration_pending'),
-      double(ActiveRecord::Migration, version: 5, name: 'newer_migration_pending')
+      double(
+        ActiveRecord::Migration,
+        version: 5,
+        name: 'older_migration_pending',
+        filename: 'db/migrate/5_older_migration_pending.rb'
+      ),
+      double(
+        ActiveRecord::Migration,
+        version: 6,
+        name: 'post_migration_pending',
+        filename: 'db/post_migrate/6_post_migration_pending.rb'
+      ),
+      double(
+        ActiveRecord::Migration,
+        version: 7,
+        name: 'newer_migration_pending',
+        filename: 'db/migrate/7_newer_migration_pending.rb'
+      )
     ].sort_by(&:version)
   end
 
@@ -87,11 +126,11 @@ RSpec.describe Gitlab::Database::Migrations::Runner, :reestablished_active_recor
       context 'running migrations' do
         subject(:up) { described_class.up(database: database, legacy_mode: legacy_mode) }
 
-        it 'runs the unapplied migrations in version order', :aggregate_failures do
+        it 'runs the unapplied migrations in regular/post order, then version order', :aggregate_failures do
           up.run
 
-          expect(migration_runs.map(&:dir)).to match_array([:up, :up])
-          expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
+          expect(migration_runs.map(&:dir)).to match_array([:up, :up, :up])
+          expect(migration_runs.map(&:version_to_migrate)).to eq([5, 7, 6])
         end
 
         it 'writes a metadata file with the current schema version and database name' do
@@ -130,8 +169,8 @@ RSpec.describe Gitlab::Database::Migrations::Runner, :reestablished_active_recor
         it 'runs the applied migrations for the current branch in reverse order', :aggregate_failures do
           down.run
 
-          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))
+          expect(migration_runs.map(&:dir)).to match_array([:down, :down, :down])
+          expect(migration_runs.map(&:version_to_migrate)).to eq([3, 4, 2])
         end
       end
 
diff --git a/spec/lib/gitlab/database/migrations/timeout_helpers_spec.rb b/spec/lib/gitlab/database/migrations/timeout_helpers_spec.rb
new file mode 100644
index 00000000000..d35211af680
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/timeout_helpers_spec.rb
@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::TimeoutHelpers do
+  let(:model) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  describe '#disable_statement_timeout' do
+    it 'disables statement timeouts to current transaction only' do
+      expect(model).to receive(:execute).with('SET LOCAL statement_timeout TO 0')
+
+      model.disable_statement_timeout
+    end
+
+    # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
+    context 'with real environment', :delete do
+      before do
+        model.execute("SET statement_timeout TO '20000'")
+      end
+
+      after do
+        model.execute('RESET statement_timeout')
+      end
+
+      it 'defines statement to 0 only for current transaction' do
+        expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+
+        model.connection.transaction do
+          model.disable_statement_timeout
+          expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+        end
+
+        expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+      end
+
+      context 'when passing a blocks' do
+        it 'disables statement timeouts on session level and executes the block' do
+          expect(model).to receive(:execute).with('SET statement_timeout TO 0')
+          expect(model).to receive(:execute).with('RESET statement_timeout').at_least(:once)
+
+          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
+        end
+
+        # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
+        context 'with real environment', :delete do
+          before do
+            model.execute("SET statement_timeout TO '20000'")
+          end
+
+          after do
+            model.execute('RESET statement_timeout')
+          end
+
+          it 'defines statement to 0 for any code run inside the block' do
+            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+
+            model.disable_statement_timeout do
+              model.connection.transaction do
+                expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+              end
+
+              expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+            end
+          end
+        end
+      end
+    end
+
+    # This spec runs without an enclosing transaction (:delete truncation method for db_cleaner)
+    context 'when the statement_timeout is already disabled', :delete do
+      before do
+        ActiveRecord::Migration.connection.execute('SET statement_timeout TO 0')
+      end
+
+      after do
+        # Use ActiveRecord::Migration.connection instead of model.execute
+        # so that this call is not counted below
+        ActiveRecord::Migration.connection.execute('RESET statement_timeout')
+      end
+
+      it 'yields control without disabling the timeout or resetting' do
+        expect(model).not_to receive(:execute).with('SET statement_timeout TO 0')
+        expect(model).not_to receive(:execute).with('RESET statement_timeout')
+
+        expect { |block| model.disable_statement_timeout(&block) }.to yield_control
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb b/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
index 0e804b4feac..cd3a94f5737 100644
--- a/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
@@ -16,6 +16,7 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
   let(:referenced_table_name) { '_test_referenced_table' }
   let(:other_referenced_table_name) { '_test_other_referenced_table' }
   let(:parent_table_name) { "#{table_name}_parent" }
+  let(:lock_tables) { [] }
 
   let(:model) { define_batchable_model(table_name, connection: connection) }
 
@@ -27,7 +28,8 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
       table_name: table_name,
       partitioning_column: partitioning_column,
       parent_table_name: parent_table_name,
-      zero_partition_value: partitioning_default
+      zero_partition_value: partitioning_default,
+      lock_tables: lock_tables
     )
   end
 
@@ -168,6 +170,16 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
       end
     end
 
+    context 'with locking tables' do
+      let(:lock_tables) { [table_name] }
+
+      it 'locks the table' do
+        recorder = ActiveRecord::QueryRecorder.new { partition }
+
+        expect(recorder.log).to include(/LOCK "_test_table_to_partition" IN ACCESS EXCLUSIVE MODE/)
+      end
+    end
+
     context 'when an error occurs during the conversion' do
       def fail_first_time
         # We can't directly use a boolean here, as we need something that will be passed by-reference to the proc
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 2ef873e8adb..336dec3a8a0 100644
--- a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
@@ -92,11 +92,11 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
 
       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|
-            expect(partition_name).to eq('test_partition')
-            expect(foreign_key_exists_by_name(partition_name, 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_falsey
+          expect(dropper).to receive(:drop_detached_partition).and_wrap_original do |drop_method, partition|
+            expect(partition.table_name).to eq('test_partition')
+            expect(foreign_key_exists_by_name(partition.table_name, 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_falsey
 
-            drop_method.call(partition_name)
+            drop_method.call(partition)
           end
 
           expect(foreign_key_exists_by_name('test_partition', 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_truthy
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb
index 7465f69b87c..a81c8a5a49c 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb
@@ -65,8 +65,11 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
       end
 
       def expect_add_concurrent_index_and_call_original(table, column, index)
-        expect(migration).to receive(:add_concurrent_index).ordered.with(table, column, { name: index })
-          .and_wrap_original { |_, table, column, options| connection.add_index(table, column, **options) }
+        expect(migration).to receive(:add_concurrent_index).ordered.with(table, column, { name: index, allow_partition: true })
+          .and_wrap_original do |_, table, column, options|
+            options.delete(:allow_partition)
+            connection.add_index(table, column, **options)
+          end
       end
     end
 
@@ -91,7 +94,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
 
       it 'forwards them to the index helper methods', :aggregate_failures do
         expect(migration).to receive(:add_concurrent_index)
-          .with(partition1_identifier, column_name, { name: partition1_index, where: 'x > 0', unique: true })
+          .with(partition1_identifier, column_name, { name: partition1_index, where: 'x > 0', unique: true, allow_partition: true })
 
         expect(migration).to receive(:add_index)
           .with(table_name, column_name, { name: index_name, where: 'x > 0', unique: true })
@@ -231,4 +234,165 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
       end
     end
   end
+
+  describe '#indexes_by_definition_for_table' do
+    context 'when a partitioned table has indexes' do
+      subject do
+        migration.indexes_by_definition_for_table(table_name)
+      end
+
+      before do
+        connection.execute(<<~SQL)
+          CREATE INDEX #{index_name} ON #{table_name} (#{column_name});
+        SQL
+      end
+
+      it 'captures partitioned index names by index definition' do
+        expect(subject).to match(a_hash_including({ "CREATE _ btree (#{column_name})" => index_name }))
+      end
+    end
+
+    context 'when a non-partitioned table has indexes' do
+      let(:regular_table_name) { '_test_regular_table' }
+      let(:regular_index_name) { '_test_regular_index_name' }
+
+      subject do
+        migration.indexes_by_definition_for_table(regular_table_name)
+      end
+
+      before do
+        connection.execute(<<~SQL)
+          CREATE TABLE #{regular_table_name} (
+            #{column_name} timestamptz NOT NULL
+          );
+
+          CREATE INDEX #{regular_index_name} ON #{regular_table_name} (#{column_name});
+        SQL
+      end
+
+      it 'captures index names by index definition' do
+        expect(subject).to match(a_hash_including({ "CREATE _ btree (#{column_name})" => regular_index_name }))
+      end
+    end
+
+    context 'when a non-partitioned table has duplicate indexes' do
+      let(:regular_table_name) { '_test_regular_table' }
+      let(:regular_index_name) { '_test_regular_index_name' }
+      let(:duplicate_index_name) { '_test_duplicate_index_name' }
+
+      subject do
+        migration.indexes_by_definition_for_table(regular_table_name)
+      end
+
+      before do
+        connection.execute(<<~SQL)
+          CREATE TABLE #{regular_table_name} (
+            #{column_name} timestamptz NOT NULL
+          );
+
+          CREATE INDEX #{regular_index_name} ON #{regular_table_name} (#{column_name});
+          CREATE INDEX #{duplicate_index_name} ON #{regular_table_name} (#{column_name});
+        SQL
+      end
+
+      it 'raises an error' do
+        expect { subject }.to raise_error { described_class::DuplicatedIndexesError }
+      end
+    end
+  end
+
+  describe '#rename_indexes_for_table' do
+    let(:original_table_name) { '_test_rename_indexes_table' }
+    let(:first_partition_name) { '_test_rename_indexes_table_1' }
+    let(:transient_table_name) { '_test_rename_indexes_table_child' }
+    let(:custom_column_name) { 'created_at' }
+    let(:generated_column_name) { 'updated_at' }
+    let(:custom_index_name) { 'index_test_rename_indexes_table_on_created_at' }
+    let(:custom_index_name_regenerated) { '_test_rename_indexes_table_created_at_idx' }
+    let(:generated_index_name) { '_test_rename_indexes_table_updated_at_idx' }
+    let(:generated_index_name_collided) { '_test_rename_indexes_table_updated_at_idx1' }
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE #{original_table_name} (
+          #{custom_column_name} timestamptz NOT NULL,
+          #{generated_column_name} timestamptz NOT NULL
+        );
+
+        CREATE INDEX #{custom_index_name} ON #{original_table_name} (#{custom_column_name});
+        CREATE INDEX ON #{original_table_name} (#{generated_column_name});
+      SQL
+    end
+
+    context 'when changing a table within the current schema' do
+      let!(:identifiers) { migration.indexes_by_definition_for_table(original_table_name) }
+
+      before do
+        connection.execute(<<~SQL)
+          ALTER TABLE #{original_table_name} RENAME TO #{first_partition_name};
+          CREATE TABLE #{original_table_name} (LIKE #{first_partition_name} INCLUDING ALL);
+          DROP TABLE #{first_partition_name};
+        SQL
+      end
+
+      it 'maps index names after they are changed' do
+        migration.rename_indexes_for_table(original_table_name, identifiers)
+
+        expect_index_to_exist(custom_index_name)
+        expect_index_to_exist(generated_index_name)
+      end
+
+      it 'does not rename an index which does not exist in the to_hash' do
+        partial_identifiers = identifiers.reject { |_, name| name == custom_index_name }
+
+        migration.rename_indexes_for_table(original_table_name, partial_identifiers)
+
+        expect_index_not_to_exist(custom_index_name)
+        expect_index_to_exist(generated_index_name)
+      end
+    end
+
+    context 'when partitioning an existing table' do
+      before do
+        connection.execute(<<~SQL)
+          /* Create new parent table */
+          CREATE TABLE #{first_partition_name} (LIKE #{original_table_name} INCLUDING ALL);
+        SQL
+      end
+
+      it 'renames indexes across schemas' do
+        # Capture index names generated by postgres
+        generated_index_names = migration.indexes_by_definition_for_table(first_partition_name)
+
+        # Capture index names from original table
+        original_index_names = migration.indexes_by_definition_for_table(original_table_name)
+
+        connection.execute(<<~SQL)
+          /* Rename original table out of the way */
+          ALTER TABLE #{original_table_name} RENAME TO #{transient_table_name};
+
+          /* Rename new parent table to original name */
+          ALTER TABLE #{first_partition_name} RENAME TO #{original_table_name};
+
+          /* Move original table to gitlab_partitions_dynamic schema */
+          ALTER TABLE #{transient_table_name} SET SCHEMA #{partition_schema};
+
+          /* Rename original table to be the first partition */
+          ALTER TABLE #{partition_schema}.#{transient_table_name} RENAME TO #{first_partition_name};
+        SQL
+
+        # Apply index names generated by postgres to first partition
+        migration.rename_indexes_for_table(first_partition_name, generated_index_names, schema_name: partition_schema)
+
+        expect_index_to_exist('_test_rename_indexes_table_1_created_at_idx')
+        expect_index_to_exist('_test_rename_indexes_table_1_updated_at_idx')
+
+        # Apply index names from original table to new parent table
+        migration.rename_indexes_for_table(original_table_name, original_index_names)
+
+        expect_index_to_exist(custom_index_name)
+        expect_index_to_exist(generated_index_name)
+      end
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
index 8bb9ad2737a..e76b1da3834 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
@@ -43,6 +43,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
 
   context 'list partitioning conversion helpers' do
     shared_examples_for 'delegates to ConvertTableToFirstListPartition' do
+      let(:extra_options) { {} }
       it 'throws an error if in a transaction' do
         allow(migration).to receive(:transaction_open?).and_return(true)
         expect { migrate }.to raise_error(/cannot be run inside a transaction/)
@@ -54,7 +55,8 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
                                 table_name: source_table,
                                 parent_table_name: partitioned_table,
                                 partitioning_column: partition_column,
-                                zero_partition_value: min_date) do |converter|
+                                zero_partition_value: min_date,
+                                **extra_options) do |converter|
           expect(converter).to receive(expected_method)
         end
 
@@ -64,12 +66,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
 
     describe '#convert_table_to_first_list_partition' do
       it_behaves_like 'delegates to ConvertTableToFirstListPartition' do
+        let(:lock_tables) { [source_table] }
+        let(:extra_options) { { lock_tables: lock_tables } }
         let(:expected_method) { :partition }
         let(:migrate) do
           migration.convert_table_to_first_list_partition(table_name: source_table,
                                                           partitioning_column: partition_column,
                                                           parent_table_name: partitioned_table,
-                                                          initial_partitioning_value: min_date)
+                                                          initial_partitioning_value: min_date,
+                                                          lock_tables: lock_tables)
         end
       end
     end
diff --git a/spec/lib/gitlab/database/postgres_partition_spec.rb b/spec/lib/gitlab/database/postgres_partition_spec.rb
index 5a44090d5ae..14a4d405621 100644
--- a/spec/lib/gitlab/database/postgres_partition_spec.rb
+++ b/spec/lib/gitlab/database/postgres_partition_spec.rb
@@ -72,4 +72,36 @@ RSpec.describe Gitlab::Database::PostgresPartition, type: :model do
       expect(find(identifier).condition).to eq("FOR VALUES FROM ('2020-01-01 00:00:00+00') TO ('2020-02-01 00:00:00+00')")
     end
   end
+
+  describe '.partition_exists?' do
+    subject { described_class.partition_exists?(table_name) }
+
+    context 'when the partition exists' do
+      let(:table_name) { "ci_builds_metadata" }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'when the partition does not exist' do
+      let(:table_name) { 'partition_does_not_exist' }
+
+      it { is_expected.to be_falsey }
+    end
+  end
+
+  describe '.legacy_partition_exists?' do
+    subject { described_class.legacy_partition_exists?(table_name) }
+
+    context 'when the partition exists' do
+      let(:table_name) { "ci_builds_metadata" }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'when the partition does not exist' do
+      let(:table_name) { 'partition_does_not_exist' }
+
+      it { is_expected.to be_falsey }
+    end
+  end
 end
diff --git a/spec/lib/gitlab/database/query_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzer_spec.rb
index 0b849063562..6dc9ffc4aba 100644
--- a/spec/lib/gitlab/database/query_analyzer_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzer_spec.rb
@@ -10,6 +10,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   before do
     allow(described_class.instance).to receive(:all_analyzers).and_return([analyzer, disabled_analyzer])
     allow(analyzer).to receive(:enabled?).and_return(true)
+    allow(analyzer).to receive(:raw?).and_return(false)
     allow(analyzer).to receive(:suppressed?).and_return(false)
     allow(analyzer).to receive(:begin!)
     allow(analyzer).to receive(:end!)
@@ -181,6 +182,13 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
       expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
     end
 
+    it 'does call analyze with raw sql when raw? is true' do
+      expect(analyzer).to receive(:raw?).and_return(true)
+      expect(analyzer).to receive(:analyze).with('SELECT 1 FROM projects')
+
+      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/ci/partitioning_id_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzers/ci/partitioning_id_analyzer_spec.rb
new file mode 100644
index 00000000000..0fe19041b6d
--- /dev/null
+++ b/spec/lib/gitlab/database/query_analyzers/ci/partitioning_id_analyzer_spec.rb
@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::QueryAnalyzers::Ci::PartitioningIdAnalyzer, query_analyzers: false do
+  let(:analyzer) { described_class }
+
+  before do
+    allow(Gitlab::Database::QueryAnalyzer.instance).to receive(:all_analyzers).and_return([analyzer])
+  end
+
+  context 'when ci_partitioning_analyze_queries_partition_id_check is disabled' do
+    before do
+      stub_feature_flags(ci_partitioning_analyze_queries_partition_id_check: false)
+    end
+
+    it 'does not analyze the query' do
+      expect(analyzer).not_to receive(:analyze)
+
+      process_sql(Ci::BuildMetadata, "SELECT 1 FROM ci_builds_metadata")
+    end
+  end
+
+  context 'when ci_partitioning_analyze_queries_partition_id_check is enabled' do
+    context 'when querying a routing table' do
+      shared_examples 'a good query' do |sql|
+        it 'does not raise error' do
+          expect { process_sql(Ci::BuildMetadata, sql) }.not_to raise_error
+        end
+      end
+
+      shared_examples 'a bad query' do |sql|
+        it 'raises PartitionIdMissingError' do
+          expect { process_sql(Ci::BuildMetadata, sql) }.to raise_error(described_class::PartitionIdMissingError)
+        end
+      end
+
+      context 'when partition_id is present' do
+        context 'when selecting data' do
+          it_behaves_like 'a good query', 'SELECT * FROM p_ci_builds_metadata WHERE partition_id = 100'
+        end
+
+        context 'with a join query' do
+          sql = <<~SQL
+            SELECT ci_builds.id
+            FROM p_ci_builds
+            JOIN p_ci_builds_metadata ON p_ci_builds_metadata.build_id = ci_builds.id
+            WHERE ci_builds.type = 'Ci::Build'
+              AND ci_builds.partition_id = 100
+              AND (NOT p_ci_builds_metadata.id IN
+                    (SELECT p_ci_builds_metadata.id
+                      FROM p_ci_builds_metadata
+                      WHERE p_ci_builds_metadata.build_id = ci_builds.id
+                        AND p_ci_builds_metadata.interruptible = TRUE
+                        AND p_ci_builds_metadata.partition_id = 100 ));
+          SQL
+
+          it_behaves_like 'a good query', sql
+        end
+
+        context 'when removing data' do
+          it_behaves_like 'a good query', 'DELETE FROM p_ci_builds_metadata WHERE partition_id = 100'
+        end
+
+        context 'when updating data' do
+          sql = 'UPDATE p_ci_builds_metadata SET interruptible = false WHERE partition_id = 100'
+
+          it_behaves_like 'a good query', sql
+        end
+
+        context 'when inserting a record' do
+          it_behaves_like 'a good query', 'INSERT INTO p_ci_builds_metadata (id, partition_id) VALUES(1, 1)'
+        end
+      end
+
+      context 'when partition_id is missing' do
+        context 'when inserting a record' do
+          it_behaves_like 'a bad query', 'INSERT INTO p_ci_builds_metadata (id) VALUES(1)'
+        end
+
+        context 'when selecting data' do
+          it_behaves_like 'a bad query', 'SELECT * FROM p_ci_builds_metadata WHERE id = 1'
+        end
+
+        context 'when removing data' do
+          it_behaves_like 'a bad query', 'DELETE FROM p_ci_builds_metadata WHERE id = 1'
+        end
+
+        context 'when updating data' do
+          it_behaves_like 'a bad query', 'UPDATE p_ci_builds_metadata SET interruptible = false WHERE id = 1'
+        end
+
+        context 'with a join query' do
+          sql = <<~SQL
+            SELECT ci_builds.id
+            FROM ci_builds
+            JOIN p_ci_builds_metadata ON p_ci_builds_metadata.build_id = ci_builds.id
+            WHERE ci_builds.type = 'Ci::Build'
+              AND ci_builds.partition_id = 100
+              AND (NOT p_ci_builds_metadata.id IN
+                    (SELECT p_ci_builds_metadata.id
+                      FROM p_ci_builds_metadata
+                      WHERE p_ci_builds_metadata.build_id = ci_builds.id
+                        AND p_ci_builds_metadata.interruptible = TRUE ));
+          SQL
+
+          it_behaves_like 'a bad query', sql
+        end
+      end
+    end
+  end
+
+  private
+
+  def process_sql(model, sql)
+    Gitlab::Database::QueryAnalyzer.instance.within do
+      # Skip load balancer and retrieve connection assigned to model
+      Gitlab::Database::QueryAnalyzer.instance.send(:process_sql, sql, model.retrieve_connection)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/query_analyzers/ci/partitioning_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzers/ci/partitioning_routing_analyzer_spec.rb
index ef7c7965c09..1f86c2ccbb0 100644
--- a/spec/lib/gitlab/database/query_analyzers/ci/partitioning_analyzer_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzers/ci/partitioning_routing_analyzer_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::QueryAnalyzers::Ci::PartitioningAnalyzer, query_analyzers: false do
+RSpec.describe Gitlab::Database::QueryAnalyzers::Ci::PartitioningRoutingAnalyzer, query_analyzers: false do
   let(:analyzer) { described_class }
 
   before do
@@ -54,15 +54,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::Ci::PartitioningAnalyzer, query
 
     context 'when analyzing non targeted table' do
       it 'does not raise error' do
-        expect { process_sql(Ci::BuildMetadata, "SELECT 1 FROM projects") }
-          .not_to raise_error
-      end
-    end
-
-    context 'when querying a routing table' do
-      it 'does not raise error' do
-        expect { process_sql(Ci::BuildMetadata, "SELECT 1 FROM p_ci_builds_metadata") }
-          .not_to raise_error
+        expect { process_sql(Ci::BuildMetadata, "SELECT 1 FROM projects") }.not_to raise_error
       end
     end
   end
diff --git a/spec/lib/gitlab/database/query_analyzers/query_recorder_spec.rb b/spec/lib/gitlab/database/query_analyzers/query_recorder_spec.rb
new file mode 100644
index 00000000000..ec01ae623ae
--- /dev/null
+++ b/spec/lib/gitlab/database/query_analyzers/query_recorder_spec.rb
@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::QueryAnalyzers::QueryRecorder, query_analyzers: false do
+  # We keep only the QueryRecorder analyzer running
+  around do |example|
+    described_class.with_suppressed(false) do
+      example.run
+    end
+  end
+
+  context 'when analyzer is enabled for tests' do
+    let(:query) { 'SELECT 1 FROM projects' }
+    let(:log_path) { Rails.root.join(described_class::LOG_FILE) }
+
+    before do
+      stub_env('CI', 'true')
+
+      # This is needed to be able to stub_env the CI variable
+      ::Gitlab::Database::QueryAnalyzer.instance.begin!([described_class])
+    end
+
+    after do
+      ::Gitlab::Database::QueryAnalyzer.instance.end!([described_class])
+    end
+
+    it 'logs queries to a file' do
+      allow(FileUtils).to receive(:mkdir_p)
+        .with(File.dirname(log_path))
+      expect(File).to receive(:write)
+        .with(log_path, /^{"sql":"#{query}/, mode: 'a')
+      expect(described_class).to receive(:analyze).with(/^#{query}/).and_call_original
+
+      expect { ApplicationRecord.connection.execute(query) }.not_to raise_error
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/tables_truncate_spec.rb b/spec/lib/gitlab/database/tables_truncate_spec.rb
index 01af9efd782..4f68cd93a8e 100644
--- a/spec/lib/gitlab/database/tables_truncate_spec.rb
+++ b/spec/lib/gitlab/database/tables_truncate_spec.rb
@@ -233,6 +233,26 @@ RSpec.describe Gitlab::Database::TablesTruncate, :reestablished_active_record_ba
     it_behaves_like 'truncating legacy tables on a database'
   end
 
+  context 'when running with multiple shared databases' do
+    before do
+      skip_if_multiple_databases_not_setup
+      ci_db_config = Ci::ApplicationRecord.connection_db_config
+      allow(::Gitlab::Database).to receive(:db_config_share_with).with(ci_db_config).and_return('main')
+    end
+
+    it 'raises an error when truncating the main database that it is a single database setup' do
+      expect do
+        described_class.new(database_name: 'main', min_batch_size: min_batch_size).execute
+      end.to raise_error(/Cannot truncate legacy tables in single-db setup/)
+    end
+
+    it 'raises an error when truncating the ci database that it is a single database setup' do
+      expect do
+        described_class.new(database_name: 'ci', min_batch_size: min_batch_size).execute
+      end.to raise_error(/Cannot truncate legacy tables in single-db setup/)
+    end
+  end
+
   context 'when running in a single database mode' do
     before do
       skip_if_multiple_databases_are_setup
diff --git a/spec/lib/gitlab/database/type/symbolized_jsonb_spec.rb b/spec/lib/gitlab/database/type/symbolized_jsonb_spec.rb
new file mode 100644
index 00000000000..a8401667b34
--- /dev/null
+++ b/spec/lib/gitlab/database/type/symbolized_jsonb_spec.rb
@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Type::SymbolizedJsonb do
+  let(:type) { described_class.new }
+
+  describe '#deserialize' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { type.deserialize(json) }
+
+    where(:json, :value) do
+      nil                                   | nil
+      '{"key":"value"}'                     | { key: 'value' }
+      '{"key":[1,2,3]}'                     | { key: [1, 2, 3] }
+      '{"key":{"subkey":"value"}}'          | { key: { subkey: 'value' } }
+      '{"key":{"a":[{"b":"c"},{"d":"e"}]}}' | { key: { a: [{ b: 'c' }, { d: 'e' }] } }
+    end
+
+    with_them do
+      it { is_expected.to match(value) }
+    end
+  end
+
+  context 'when used by a model' do
+    let(:model) do
+      Class.new(ApplicationRecord) do
+        self.table_name = :_test_symbolized_jsonb
+
+        attribute :options, :sym_jsonb
+      end
+    end
+
+    let(:record) do
+      model.create!(name: 'test', options: { key: 'value' })
+    end
+
+    before do
+      ApplicationRecord.connection.execute(<<~SQL)
+        CREATE TABLE _test_symbolized_jsonb(
+          id serial NOT NULL PRIMARY KEY,
+          name text,
+          options jsonb);
+      SQL
+
+      model.reset_column_information
+    end
+
+    it { expect(record.options).to match({ key: 'value' }) }
+
+    it 'ignores changes to other attributes' do
+      record.name = 'other test'
+
+      expect(record.changes).to match('name' => ['test', 'other test'])
+    end
+
+    it 'tracks changes to options' do
+      record.options = { key: 'other value' }
+
+      expect(record.changes).to match('options' => [{ 'key' => 'value' }, { 'key' => 'other value' }])
+    end
+  end
+end