Add latest changes from gitlab-org/gitlab@14-9-stable-eev14.9.0-rc42

24 files changed, 1331 insertions, 149 deletions

diff --git a/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb
index eb16a8ccfa5..9ba3dad72b3 100644
--- a/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/async_indexes/migration_helpers_spec.rb
@@ -16,45 +16,29 @@ RSpec.describe Gitlab::Database::AsyncIndexes::MigrationHelpers do
   describe '#unprepare_async_index' do
     let!(:async_index) { create(:postgres_async_index, name: index_name) }
 
-    context 'when the flag is enabled' do
-      before do
-        stub_feature_flags(database_async_index_creation: true)
-      end
+    it 'destroys the record' do
+      expect do
+        migration.unprepare_async_index(table_name, 'id')
+      end.to change { index_model.where(name: index_name).count }.by(-1)
+    end
+
+    context 'when an explicit name is given' do
+      let(:index_name) { 'my_test_async_index' }
 
       it 'destroys the record' do
         expect do
-          migration.unprepare_async_index(table_name, 'id')
+          migration.unprepare_async_index(table_name, 'id', name: index_name)
         end.to change { index_model.where(name: index_name).count }.by(-1)
       end
-
-      context 'when an explicit name is given' do
-        let(:index_name) { 'my_test_async_index' }
-
-        it 'destroys the record' do
-          expect do
-            migration.unprepare_async_index(table_name, 'id', name: index_name)
-          end.to change { index_model.where(name: index_name).count }.by(-1)
-        end
-      end
-
-      context 'when the async index table does not exist' do
-        it 'does not raise an error' do
-          connection.drop_table(:postgres_async_indexes)
-
-          expect(index_model).not_to receive(:find_by)
-
-          expect { migration.unprepare_async_index(table_name, 'id') }.not_to raise_error
-        end
-      end
     end
 
-    context 'when the feature flag is disabled' do
-      it 'does not destroy the record' do
-        stub_feature_flags(database_async_index_creation: false)
+    context 'when the async index table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_indexes)
 
-        expect do
-          migration.unprepare_async_index(table_name, 'id')
-        end.not_to change { index_model.where(name: index_name).count }
+        expect(index_model).not_to receive(:find_by)
+
+        expect { migration.unprepare_async_index(table_name, 'id') }.not_to raise_error
       end
     end
   end
@@ -63,35 +47,19 @@ RSpec.describe Gitlab::Database::AsyncIndexes::MigrationHelpers do
     let(:index_name) { "index_#{table_name}_on_id" }
     let!(:async_index) { create(:postgres_async_index, name: index_name) }
 
-    context 'when the flag is enabled' do
-      before do
-        stub_feature_flags(database_async_index_creation: true)
-      end
-
-      it 'destroys the record' do
-        expect do
-          migration.unprepare_async_index_by_name(table_name, index_name)
-        end.to change { index_model.where(name: index_name).count }.by(-1)
-      end
-
-      context 'when the async index table does not exist' do
-        it 'does not raise an error' do
-          connection.drop_table(:postgres_async_indexes)
-
-          expect(index_model).not_to receive(:find_by)
-
-          expect { migration.unprepare_async_index_by_name(table_name, index_name) }.not_to raise_error
-        end
-      end
+    it 'destroys the record' do
+      expect do
+        migration.unprepare_async_index_by_name(table_name, index_name)
+      end.to change { index_model.where(name: index_name).count }.by(-1)
     end
 
-    context 'when the feature flag is disabled' do
-      it 'does not destroy the record' do
-        stub_feature_flags(database_async_index_creation: false)
+    context 'when the async index table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_indexes)
 
-        expect do
-          migration.unprepare_async_index_by_name(table_name, index_name)
-        end.not_to change { index_model.where(name: index_name).count }
+        expect(index_model).not_to receive(:find_by)
+
+        expect { migration.unprepare_async_index_by_name(table_name, index_name) }.not_to raise_error
       end
     end
   end
@@ -101,14 +69,23 @@ RSpec.describe Gitlab::Database::AsyncIndexes::MigrationHelpers do
       connection.create_table(table_name)
     end
 
-    context 'when the feature flag is enabled' do
-      before do
-        stub_feature_flags(database_async_index_creation: true)
-      end
+    it 'creates the record for the async index' do
+      expect do
+        migration.prepare_async_index(table_name, 'id')
+      end.to change { index_model.where(name: index_name).count }.by(1)
+
+      record = index_model.find_by(name: index_name)
 
-      it 'creates the record for the async index' do
+      expect(record.table_name).to eq(table_name)
+      expect(record.definition).to match(/CREATE INDEX CONCURRENTLY "#{index_name}"/)
+    end
+
+    context 'when an explicit name is given' do
+      let(:index_name) { 'my_async_index_name' }
+
+      it 'creates the record with the given name' do
         expect do
-          migration.prepare_async_index(table_name, 'id')
+          migration.prepare_async_index(table_name, 'id', name: index_name)
         end.to change { index_model.where(name: index_name).count }.by(1)
 
         record = index_model.find_by(name: index_name)
@@ -116,77 +93,52 @@ RSpec.describe Gitlab::Database::AsyncIndexes::MigrationHelpers do
         expect(record.table_name).to eq(table_name)
         expect(record.definition).to match(/CREATE INDEX CONCURRENTLY "#{index_name}"/)
       end
+    end
 
-      context 'when an explicit name is given' do
-        let(:index_name) { 'my_async_index_name' }
-
-        it 'creates the record with the given name' do
-          expect do
-            migration.prepare_async_index(table_name, 'id', name: index_name)
-          end.to change { index_model.where(name: index_name).count }.by(1)
-
-          record = index_model.find_by(name: index_name)
+    context 'when the index already exists' do
+      it 'does not create the record' do
+        connection.add_index(table_name, 'id', name: index_name)
 
-          expect(record.table_name).to eq(table_name)
-          expect(record.definition).to match(/CREATE INDEX CONCURRENTLY "#{index_name}"/)
-        end
+        expect do
+          migration.prepare_async_index(table_name, 'id')
+        end.not_to change { index_model.where(name: index_name).count }
       end
+    end
 
-      context 'when the index already exists' do
-        it 'does not create the record' do
-          connection.add_index(table_name, 'id', name: index_name)
+    context 'when the record already exists' do
+      it 'does attempt to create the record' do
+        create(:postgres_async_index, table_name: table_name, name: index_name)
 
-          expect do
-            migration.prepare_async_index(table_name, 'id')
-          end.not_to change { index_model.where(name: index_name).count }
-        end
+        expect do
+          migration.prepare_async_index(table_name, 'id')
+        end.not_to change { index_model.where(name: index_name).count }
       end
 
-      context 'when the record already exists' do
-        it 'does attempt to create the record' do
-          create(:postgres_async_index, table_name: table_name, name: index_name)
-
-          expect do
-            migration.prepare_async_index(table_name, 'id')
-          end.not_to change { index_model.where(name: index_name).count }
-        end
-
-        it 'updates definition if changed' do
-          index = create(:postgres_async_index, table_name: table_name, name: index_name, definition: '...')
-
-          expect do
-            migration.prepare_async_index(table_name, 'id', name: index_name)
-          end.to change { index.reload.definition }
-        end
+      it 'updates definition if changed' do
+        index = create(:postgres_async_index, table_name: table_name, name: index_name, definition: '...')
 
-        it 'does not update definition if not changed' do
-          definition = "CREATE INDEX CONCURRENTLY \"index_#{table_name}_on_id\" ON \"#{table_name}\" (\"id\")"
-          index = create(:postgres_async_index, table_name: table_name, name: index_name, definition: definition)
-
-          expect do
-            migration.prepare_async_index(table_name, 'id', name: index_name)
-          end.not_to change { index.reload.updated_at }
-        end
+        expect do
+          migration.prepare_async_index(table_name, 'id', name: index_name)
+        end.to change { index.reload.definition }
       end
 
-      context 'when the async index table does not exist' do
-        it 'does not raise an error' do
-          connection.drop_table(:postgres_async_indexes)
-
-          expect(index_model).not_to receive(:safe_find_or_create_by!)
+      it 'does not update definition if not changed' do
+        definition = "CREATE INDEX CONCURRENTLY \"index_#{table_name}_on_id\" ON \"#{table_name}\" (\"id\")"
+        index = create(:postgres_async_index, table_name: table_name, name: index_name, definition: definition)
 
-          expect { migration.prepare_async_index(table_name, 'id') }.not_to raise_error
-        end
+        expect do
+          migration.prepare_async_index(table_name, 'id', name: index_name)
+        end.not_to change { index.reload.updated_at }
       end
     end
 
-    context 'when the feature flag is disabled' do
-      it 'does not create the record' do
-        stub_feature_flags(database_async_index_creation: false)
+    context 'when the async index table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_indexes)
 
-        expect do
-          migration.prepare_async_index(table_name, 'id')
-        end.not_to change { index_model.where(name: index_name).count }
+        expect(index_model).not_to receive(:safe_find_or_create_by!)
+
+        expect { migration.prepare_async_index(table_name, 'id') }.not_to raise_error
       end
     end
   end
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 7338ea657b9..8c663ff9f8a 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_spec.rb
@@ -5,6 +5,10 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model do
   it_behaves_like 'having unique enum values'
 
+  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] }
+
   describe 'associations' do
     it { is_expected.to belong_to(:batched_migration).with_foreign_key(:batched_background_migration_id) }
     it { is_expected.to have_many(:batched_job_transition_logs).with_foreign_key(:batched_background_migration_job_id) }
@@ -13,6 +17,8 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
   describe 'state machine' do
     let_it_be(:job) { create(:batched_background_migration_job, :failed) }
 
+    it { expect(described_class.state_machine.states.map(&:name)).to eql(%i(pending running failed succeeded)) }
+
     context 'when a job is running' do
       it 'logs the transition' do
         expect(Gitlab::AppLogger).to receive(:info).with( { batched_job_id: job.id, message: 'BatchedJob transition', new_state: :running, previous_state: :failed } )
@@ -45,6 +51,51 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
       end
     end
 
+    context 'when a job fails the number of max times' do
+      let(:max_times) { described_class::MAX_ATTEMPTS }
+      let!(:job) { create(:batched_background_migration_job, :running, batch_size: 10, min_value: 6, max_value: 15, attempts: max_times) }
+
+      context 'when job can be split' do
+        let(:exception) { ActiveRecord::StatementTimeout.new('Timeout!') }
+
+        before do
+          allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+            allow(batch_class).to receive(:next_batch).and_return([6, 10])
+          end
+        end
+
+        it 'splits the job into two retriable jobs' do
+          expect { job.failure!(error: exception) }.to change { job.batched_migration.batched_jobs.retriable.count }.from(0).to(2)
+        end
+      end
+
+      context 'when the job cannot be split' do
+        let(:exception) { ActiveRecord::StatementTimeout.new('Timeout!') }
+        let(:max_times) { described_class::MAX_ATTEMPTS }
+        let!(:job) { create(:batched_background_migration_job, :running, batch_size: 50, sub_batch_size: 20, min_value: 6, max_value: 15, attempts: max_times) }
+        let(:error_message) { 'Job cannot be split further' }
+        let(:split_and_retry_exception) { Gitlab::Database::BackgroundMigration::SplitAndRetryError.new(error_message) }
+
+        before do
+          allow(job).to receive(:split_and_retry!).and_raise(split_and_retry_exception)
+        end
+
+        it 'does not split the job' do
+          expect { job.failure!(error: exception) }.not_to change { job.batched_migration.batched_jobs.retriable.count }
+        end
+
+        it 'keeps the same job attributes' do
+          expect { job.failure!(error: exception) }.not_to change { job }
+        end
+
+        it 'logs the error' do
+          expect(Gitlab::AppLogger).to receive(:error).with( { message: error_message, batched_job_id: job.id } )
+
+          job.failure!(error: exception)
+        end
+      end
+    end
+
     context 'when a job fails' do
       let(:job) { create(:batched_background_migration_job, :running) }
 
@@ -145,6 +196,49 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
     end
   end
 
+  describe '#can_split?' do
+    subject { job.can_split?(exception) }
+
+    context 'when the number of attempts is greater than the limit and the batch_size is greater than the sub_batch_size' do
+      let(:job) { create(:batched_background_migration_job, :failed, batch_size: 4, sub_batch_size: 2, attempts: described_class::MAX_ATTEMPTS + 1) }
+
+      context 'when is a timeout exception' do
+        let(:exception) { ActiveRecord::StatementTimeout.new }
+
+        it { expect(subject).to be_truthy  }
+      end
+
+      context 'when is not a timeout exception' do
+        let(:exception) { RuntimeError.new }
+
+        it { expect(subject).to be_falsey }
+      end
+    end
+
+    context 'when the number of attempts is lower than the limit and the batch_size is greater than the sub_batch_size' do
+      let(:job) { create(:batched_background_migration_job, :failed, batch_size: 4, sub_batch_size: 2, attempts: described_class::MAX_ATTEMPTS - 1) }
+
+      context 'when is a timeout exception' do
+        let(:exception) { ActiveRecord::StatementTimeout.new }
+
+        it { expect(subject).to be_falsey }
+      end
+
+      context 'when is not a timeout exception' do
+        let(:exception) { RuntimeError.new }
+
+        it { expect(subject).to be_falsey }
+      end
+    end
+
+    context 'when the batch_size is lower than the sub_batch_size' do
+      let(:job) { create(:batched_background_migration_job, :failed, batch_size: 2, sub_batch_size: 4) }
+      let(:exception) { ActiveRecord::StatementTimeout.new }
+
+      it { expect(subject).to be_falsey }
+    end
+  end
+
   describe '#time_efficiency' do
     subject { job.time_efficiency }
 
@@ -197,15 +291,17 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
   end
 
   describe '#split_and_retry!' do
-    let!(:job) { create(:batched_background_migration_job, :failed, batch_size: 10, min_value: 6, max_value: 15, attempts: 3) }
+    let_it_be(:migration) { create(:batched_background_migration, table_name: :events) }
+    let_it_be(:job) { create(:batched_background_migration_job, :failed, batched_migration: migration, batch_size: 10, min_value: 6, max_value: 15, attempts: 3) }
+    let_it_be(:project) { create(:project) }
 
-    context 'when job can be split' do
-      before do
-        allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
-          allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
-        end
+    before_all do
+      (6..16).each do |id|
+        create(:event, id: id, project: project)
       end
+    end
 
+    context 'when job can be split' do
       it 'sets the correct attributes' do
         expect { job.split_and_retry! }.to change { described_class.count }.by(1)
 
@@ -261,9 +357,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
 
     context 'when computed midpoint is larger than the max value of the batch' do
       before do
-        allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
-          allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 16])
-        end
+        Event.where(id: 6..12).delete_all
       end
 
       it 'lowers the batch size and resets the number of attempts' do
diff --git a/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb b/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb
index c42a0fc5e05..59f4f40c0ef 100644
--- a/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_job_transition_log_spec.rb
@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJobTransitionLog, type: :model do
+  it { is_expected.to be_a Gitlab::Database::SharedModel }
+
   describe 'associations' do
     it { is_expected.to belong_to(:batched_job).with_foreign_key(:batched_background_migration_job_id) }
   end
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
index bb2c6b9a3ae..124d204cb62 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
@@ -428,4 +428,27 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       end
     end
   end
+
+  describe '.finalize' do
+    context 'when the connection is passed' do
+      let(:connection) { double('connection') }
+
+      let(:table_name) { :_test_batched_migrations_test_table }
+      let(:column_name) { :some_id }
+      let(:job_arguments) { [:some, :other, :arguments] }
+      let(:batched_migration) { create(:batched_background_migration, table_name: table_name, column_name: column_name) }
+
+      it 'initializes the object with the given connection' do
+        expect(described_class).to receive(:new).with(connection: connection).and_call_original
+
+        described_class.finalize(
+          batched_migration.job_class_name,
+          table_name,
+          column_name,
+          job_arguments,
+          connection: connection
+        )
+      end
+    end
+  end
 end
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 ea4ba4dd137..803123e8e34 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -5,6 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :model do
   it_behaves_like 'having unique enum values'
 
+  it { is_expected.to be_a Gitlab::Database::SharedModel }
+
   describe 'associations' do
     it { is_expected.to have_many(:batched_jobs).with_foreign_key(:batched_background_migration_id) }
 
@@ -272,7 +274,13 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
       before do
         allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
-          allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
+          allow(batch_class).to receive(:next_batch).with(
+            anything,
+            anything,
+            batch_min_value: 6,
+            batch_size: 5,
+            job_arguments: batched_migration.job_arguments
+          ).and_return([6, 10])
         end
       end
 
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
index 4f5536d8771..d6c984c7adb 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb
@@ -193,6 +193,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
 
     it_behaves_like 'an error is raised', RuntimeError.new('Something broke!')
     it_behaves_like 'an error is raised', SignalException.new('SIGTERM')
+    it_behaves_like 'an error is raised', ActiveRecord::StatementTimeout.new('Timeout!')
   end
 
   context 'when the batched background migration does not inherit from BaseJob' do
diff --git a/spec/lib/gitlab/database/each_database_spec.rb b/spec/lib/gitlab/database/each_database_spec.rb
index d526b3bc1ac..d46c1ca8681 100644
--- a/spec/lib/gitlab/database/each_database_spec.rb
+++ b/spec/lib/gitlab/database/each_database_spec.rb
@@ -3,13 +3,13 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::EachDatabase do
-  describe '.each_database_connection' do
+  describe '.each_database_connection', :add_ci_connection do
     before do
       allow(Gitlab::Database).to receive(:database_base_models)
         .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
     end
 
-    it 'yields each connection after connecting SharedModel', :add_ci_connection do
+    it 'yields each connection after connecting SharedModel' do
       expect(Gitlab::Database::SharedModel).to receive(:using_connection)
         .with(ActiveRecord::Base.connection).ordered.and_yield
 
@@ -22,6 +22,42 @@ RSpec.describe Gitlab::Database::EachDatabase do
           [Ci::ApplicationRecord.connection, 'ci']
         )
     end
+
+    context 'when only certain databases are selected' do
+      it 'yields the selected connections after connecting SharedModel' do
+        expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+          .with(Ci::ApplicationRecord.connection).ordered.and_yield
+
+        expect { |b| described_class.each_database_connection(only: 'ci', &b) }
+          .to yield_successive_args([Ci::ApplicationRecord.connection, 'ci'])
+      end
+
+      context 'when the selected names are passed as symbols' do
+        it 'yields the selected connections after connecting SharedModel' do
+          expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+            .with(Ci::ApplicationRecord.connection).ordered.and_yield
+
+          expect { |b| described_class.each_database_connection(only: :ci, &b) }
+            .to yield_successive_args([Ci::ApplicationRecord.connection, 'ci'])
+        end
+      end
+
+      context 'when the selected names are invalid' do
+        it 'does not yield any connections' do
+          expect do |b|
+            described_class.each_database_connection(only: :notvalid, &b)
+          rescue ArgumentError => e
+            expect(e.message).to match(/notvalid is not a valid database name/)
+          end.not_to yield_control
+        end
+
+        it 'raises an error' do
+          expect do
+            described_class.each_database_connection(only: :notvalid) {}
+          end.to raise_error(ArgumentError, /notvalid is not a valid database name/)
+        end
+      end
+    end
   end
 
   describe '.each_model_connection' do
@@ -69,8 +105,8 @@ RSpec.describe Gitlab::Database::EachDatabase do
         allow(main_model).to receive(:connection).and_return(main_connection)
         allow(ci_model).to receive(:connection).and_return(ci_connection)
 
-        allow(main_connection).to receive_message_chain('pool.db_config.name').and_return('main')
-        allow(ci_connection).to receive_message_chain('pool.db_config.name').and_return('ci')
+        allow(main_model).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(ci_model).to receive_message_chain('connection_db_config.name').and_return('ci')
       end
 
       it 'yields each model after connecting SharedModel' do
@@ -81,10 +117,44 @@ RSpec.describe Gitlab::Database::EachDatabase do
       end
     end
 
-    def expect_yielded_models(models_to_iterate, expected_values)
+    context 'when the database connections are limited by the only_on option' do
+      let(:shared_model) { Class.new(Gitlab::Database::SharedModel) }
+      let(:main_model) { Class.new(ActiveRecord::Base) }
+      let(:ci_model) { Class.new(Ci::ApplicationRecord) }
+
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models)
+          .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
+
+        allow(main_model).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(ci_model).to receive_message_chain('connection_db_config.name').and_return('ci')
+      end
+
+      context 'when a single name is passed in' do
+        it 'yields models only connected to the given database' do
+          expect_yielded_models([main_model, ci_model, shared_model], [
+            { model: ci_model, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+            { model: shared_model, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ], only_on: 'ci')
+        end
+      end
+
+      context 'when a list of names are passed in' do
+        it 'yields models only connected to the given databases' do
+          expect_yielded_models([main_model, ci_model, shared_model], [
+            { model: main_model, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: ci_model, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+            { model: shared_model, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: shared_model, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ], only_on: %i[main ci])
+        end
+      end
+    end
+
+    def expect_yielded_models(models_to_iterate, expected_values, only_on: nil)
       times_yielded = 0
 
-      described_class.each_model_connection(models_to_iterate) do |model, name|
+      described_class.each_model_connection(models_to_iterate, only_on: only_on) do |model, name|
         expected = expected_values[times_yielded]
 
         expect(model).to be(expected[:model])
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index e87c9c20707..77284b4d128 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -7,13 +7,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration, :request_store do
   let(:db_config) { ActiveRecord::DatabaseConfigurations::HashConfig.new('test', 'ci', configuration_hash) }
   let(:model) { double(:model, connection_db_config: db_config) }
 
-  before do
-    # It's confusing to think about these specs with this enabled by default so
-    # we make it disabled by default and just write the specific spec for when
-    # it's enabled
-    stub_feature_flags(force_no_sharing_primary_model: false)
-  end
-
   describe '.for_model' do
     context 'when load balancing is not configured' do
       it 'uses the default settings' do
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
index 20519a759b2..4d565ce137a 100644
--- a/spec/lib/gitlab/database/load_balancing/setup_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb
@@ -274,6 +274,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
       end
 
       before do
+        allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
+
         # Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects
         allow_next_instance_of(described_class) do |setup|
           allow(setup).to receive(:configure_connection)
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index 45878b2e266..f320fe0276f 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -92,6 +92,18 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         end
       end
 
+      context 'when an invalid connection is used' do
+        it 'returns :unknown' do
+          expect(described_class.db_role_for_connection(:invalid)).to eq(:unknown)
+        end
+      end
+
+      context 'when a null connection is used' do
+        it 'returns :unknown' do
+          expect(described_class.db_role_for_connection(nil)).to eq(:unknown)
+        end
+      end
+
       context 'when a read connection is used' do
         it 'returns :replica' do
           load_balancer.read do |connection|
diff --git a/spec/lib/gitlab/database/migration_helpers/restrict_gitlab_schema_spec.rb b/spec/lib/gitlab/database/migration_helpers/restrict_gitlab_schema_spec.rb
new file mode 100644
index 00000000000..ad9a3a6e257
--- /dev/null
+++ b/spec/lib/gitlab/database/migration_helpers/restrict_gitlab_schema_spec.rb
@@ -0,0 +1,561 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::MigrationHelpers::RestrictGitlabSchema, query_analyzers: false, stub_feature_flags: false do
+  let(:schema_class) { Class.new(Gitlab::Database::Migration[1.0]).include(described_class) }
+
+  describe '#restrict_gitlab_migration' do
+    it 'invalid schema raises exception' do
+      expect { schema_class.restrict_gitlab_migration gitlab_schema: :gitlab_non_exisiting }
+        .to raise_error /Unknown 'gitlab_schema:/
+    end
+
+    it 'does configure allowed_gitlab_schema' do
+      schema_class.restrict_gitlab_migration gitlab_schema: :gitlab_main
+
+      expect(schema_class.allowed_gitlab_schemas).to eq(%i[gitlab_main])
+    end
+  end
+
+  context 'when executing migrations' do
+    using RSpec::Parameterized::TableSyntax
+
+    where do
+      {
+        "does create table in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              create_table :_test_table do |t|
+                t.references :project, foreign_key: true, null: false
+                t.timestamps_with_timezone null: false
+              end
+            end
+          end,
+          query_matcher: /CREATE TABLE "_test_table"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add column to projects in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              add_column :projects, :__test_column, :integer
+            end
+          end,
+          query_matcher: /ALTER TABLE "projects" ADD "__test_column" integer/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add column to ci_builds in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              add_column :ci_builds, :__test_column, :integer
+            end
+          end,
+          query_matcher: /ALTER TABLE "ci_builds" ADD "__test_column" integer/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add index to projects in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              # Due to running in transactin we cannot use `add_concurrent_index`
+              add_index :projects, :hidden
+            end
+          end,
+          query_matcher: /CREATE INDEX/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add index to ci_builds in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              # Due to running in transactin we cannot use `add_concurrent_index`
+              add_index :ci_builds, :tag, where: "type = 'Ci::Build'", name: 'index_ci_builds_on_tag_and_type_eq_ci_build'
+            end
+          end,
+          query_matcher: /CREATE INDEX/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does create trigger in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::SchemaHelpers
+
+            def up
+              create_trigger_function('_test_trigger_function', replace: true) do
+                <<~SQL
+                  RETURN NULL;
+                SQL
+              end
+            end
+
+            def down
+              drop_function('_test_trigger_function')
+            end
+          end,
+          query_matcher: /CREATE OR REPLACE FUNCTION/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does create schema in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::SchemaHelpers
+
+            def up
+              execute("create schema __test_schema")
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /create schema __test_schema/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_main: {
+              # This is not properly detected today since there are no helpers
+              # available to consider this as a DDL type of change
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does attach loose foreign key trigger in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::MigrationHelpers::LooseForeignKeyHelpers
+
+            enable_lock_retries!
+
+            def up
+              track_record_deletions(:audit_events)
+            end
+
+            def down
+              untrack_record_deletions(:audit_events)
+            end
+          end,
+          query_matcher: /CREATE TRIGGER/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does insert into software_licenses" => {
+          migration: ->(klass) do
+            def up
+              software_license_class.create!(name: 'aaa')
+            end
+
+            def down
+              software_license_class.where(name: 'aaa').delete_all
+            end
+
+            def software_license_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'software_licenses'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "software_licenses"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing tables outside of gitlab_main" => {
+          migration: ->(klass) do
+            def up
+              ci_instance_variables_class.create!(variable_type: 1, key: 'aaa')
+            end
+
+            def down
+              ci_instance_variables_class.delete_all
+            end
+
+            def ci_instance_variables_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'ci_instance_variables'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "ci_instance_variables"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :dml_access_denied,
+              ci: :skipped
+            }
+          }
+        },
+        "does allow modifying gitlab_shared" => {
+          migration: ->(klass) do
+            def up
+              detached_partitions_class.create!(drop_after: Time.current, table_name: '_test_table')
+            end
+
+            def down
+            end
+
+            def detached_partitions_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'detached_partitions'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "detached_partitions"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_main: {
+              # TBD: This allow to selectively modify shared tables in context of a specific DB only
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does update data in batches of gitlab_main, but skips gitlab_ci" => {
+          migration: ->(klass) do
+            def up
+              update_column_in_batches(:projects, :archived, true) do |table, query|
+                query.where(table[:archived].eq(false)) # rubocop:disable CodeReuse/ActiveRecord
+              end
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          query_matcher: /FROM "projects"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does not allow executing mixed DDL and DML migrations" => {
+          migration: ->(klass) do
+            def up
+              execute('UPDATE projects SET hidden=false')
+              add_index(:projects, :hidden, name: 'test_index')
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does schedule background migrations on gitlab_main" => {
+          migration: ->(klass) do
+            def up
+              queue_background_migration_jobs_by_range_at_intervals(
+                define_batchable_model('vulnerability_occurrences'),
+                'RemoveDuplicateVulnerabilitiesFindings',
+                2.minutes.to_i,
+                batch_size: 5_000
+              )
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          query_matcher: /FROM "vulnerability_occurrences"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does support prepare_async_index" => {
+          migration: ->(klass) do
+            def up
+              prepare_async_index :projects, :hidden,
+                name: :index_projects_on_hidden
+            end
+
+            def down
+              unprepare_async_index_by_name :projects, :index_projects_on_hidden
+            end
+          end,
+          query_matcher: /INSERT INTO "postgres_async_indexes"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :dml_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing current settings" => {
+          migration: ->(klass) do
+            def up
+              ApplicationSetting.last
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /FROM "application_settings"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing feature flags" => {
+          migration: ->(klass) do
+            def up
+              Feature.enabled?(:redis_hll_tracking, type: :ops, default_enabled: :yaml)
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /FROM "features"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        }
+      }
+    end
+
+    with_them do
+      let(:migration_class) { Class.new(schema_class, &migration) }
+
+      Gitlab::Database.database_base_models.each do |db_config_name, model|
+        context "for db_config_name=#{db_config_name}" do
+          around do |example|
+            with_reestablished_active_record_base do
+              reconfigure_db_connection(model: ActiveRecord::Base, config_model: model)
+
+              example.run
+            end
+          end
+
+          before do
+            allow_next_instance_of(migration_class) do |migration|
+              allow(migration).to receive(:transaction_open?).and_return(false)
+            end
+          end
+
+          %i[no_gitlab_schema gitlab_schema_gitlab_main gitlab_schema_gitlab_shared].each do |restrict_gitlab_migration|
+            context "while restrict_gitlab_migration=#{restrict_gitlab_migration}" do
+              it "does run migrate :up and :down" do
+                expected_result = expected.fetch(restrict_gitlab_migration)[db_config_name.to_sym]
+                skip "not configured" unless expected_result
+
+                case restrict_gitlab_migration
+                when :no_gitlab_schema
+                  # no-op
+                when :gitlab_schema_gitlab_main
+                  migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_main
+                when :gitlab_schema_gitlab_shared
+                  migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_shared
+                end
+
+                # In some cases (for :down) we ignore error and expect no other errors
+                case expected_result
+                when :success
+                  expect { migration_class.migrate(:up) }.to make_queries_matching(query_matcher)
+                  expect { migration_class.migrate(:down) }.not_to make_queries_matching(query_matcher)
+
+                when :dml_not_allowed
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLNotAllowedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLNotAllowedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :dml_access_denied
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLAccessDeniedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLAccessDeniedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :ddl_not_allowed
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DDLNotAllowedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DDLNotAllowedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :skipped
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema::MigrationSkippedError)
+                  expect { migration_class.migrate(:down) }.to raise_error(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema::MigrationSkippedError)
+                end
+              end
+            end
+          end
+
+          def ignore_error(error)
+            yield
+          rescue error
+          end
+        end
+      end
+    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 96dc3a0fc28..e64f5807385 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -164,11 +164,19 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
           end
         end
 
-        context "when the primary_column_name is not an integer" do
+        context 'when the primary_column_name is a string' do
+          it 'does not raise error' do
+            expect do
+              model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :name_regex)
+            end.not_to raise_error
+          end
+        end
+
+        context "when the primary_column_name is not an integer or a string" do
           it 'raises error' do
             expect do
               model.queue_background_migration_jobs_by_range_at_intervals(ContainerExpirationPolicy, 'FooJob', 10.minutes, primary_column_name: :enabled)
-            end.to raise_error(StandardError, /is not an integer column/)
+            end.to raise_error(StandardError, /is not an integer or string column/)
           end
         end
 
diff --git a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
index a757cac0a2a..35e4cef6da5 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
@@ -9,7 +9,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
   let(:query) { "select date_trunc('day', $1::timestamptz) + $2 * (interval '1 hour')" }
   let(:query_binds) { [Time.current, 3] }
   let(:directory_path) { Dir.mktmpdir }
-  let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-query-details.json" }
+  let(:log_file) { "#{directory_path}/query-details.json" }
   let(:query_details) { Gitlab::Json.parse(File.read(log_file)) }
   let(:migration_version) { 20210422152437 }
   let(:migration_name) { 'test' }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
index eb66972e5ab..34678b77a0f 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
@@ -18,7 +18,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
   it 'writes a file with the query log' do
     observe
 
-    expect(File.read("#{directory_path}/#{migration_version}_#{migration_name}.log")).to include(query)
+    expect(File.read("#{directory_path}/migration.log")).to include(query)
   end
 
   it 'does not change the default logger' do
diff --git a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
index f433e25b2ba..51b19e7f2da 100644
--- a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
   let(:connection) { ActiveRecord::Migration.connection }
   let(:observation) { Gitlab::Database::Migrations::Observation.new(version: migration_version, name: migration_name) }
   let(:directory_path) { Dir.mktmpdir }
-  let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
+  let(:log_file) { "#{directory_path}/transaction-duration.json" }
   let(:transaction_duration) { Gitlab::Json.parse(File.read(log_file)) }
   let(:migration_version) { 20210422152437 }
   let(:migration_name) { 'test' }
diff --git a/spec/lib/gitlab/database/migrations/runner_spec.rb b/spec/lib/gitlab/database/migrations/runner_spec.rb
index 7dc965c84fa..84482e6b450 100644
--- a/spec/lib/gitlab/database/migrations/runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/runner_spec.rb
@@ -79,6 +79,15 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
         expect(migration_runs.map(&:dir)).to match_array([:up, :up])
         expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
       end
+
+      it 'writes a metadata file with the current schema version' do
+        up.run
+
+        metadata_file = result_dir.join('up', described_class::METADATA_FILENAME)
+        expect(metadata_file.exist?).to be_truthy
+        metadata = Gitlab::Json.parse(File.read(metadata_file))
+        expect(metadata).to match('version' => described_class::SCHEMA_VERSION)
+      end
     end
   end
 
@@ -105,5 +114,14 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
         expect(migration_runs.map(&:version_to_migrate)).to eq(applied_migrations_this_branch.reverse.map(&:version))
       end
     end
+
+    it 'writes a metadata file with the current schema version' do
+      down.run
+
+      metadata_file = result_dir.join('down', described_class::METADATA_FILENAME)
+      expect(metadata_file.exist?).to be_truthy
+      metadata = Gitlab::Json.parse(File.read(metadata_file))
+      expect(metadata).to match('version' => described_class::SCHEMA_VERSION)
+    end
   end
 end
diff --git a/spec/lib/gitlab/database/migrations/test_background_runner_spec.rb b/spec/lib/gitlab/database/migrations/test_background_runner_spec.rb
new file mode 100644
index 00000000000..c6fe88a7c2d
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/test_background_runner_spec.rb
@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
+  include Gitlab::Database::Migrations::BackgroundMigrationHelpers
+
+  # In order to test the interaction between queueing sidekiq jobs and seeing those jobs in queues,
+  # we need to disable sidekiq's testing mode and actually send our jobs to redis
+  around do |ex|
+    Sidekiq::Testing.disable! { ex.run }
+  end
+
+  context 'without jobs to run' do
+    it 'returns immediately' do
+      runner = described_class.new
+      expect(runner).not_to receive(:run_job)
+      described_class.new.run_jobs(for_duration: 1.second)
+    end
+  end
+
+  context 'with jobs to run' do
+    let(:migration_name) { 'TestBackgroundMigration' }
+
+    before do
+      (1..5).each do |i|
+        migrate_in(i.minutes, migration_name, [i])
+      end
+    end
+
+    context 'finding pending background jobs' do
+      it 'finds all the migrations' do
+        expect(described_class.new.traditional_background_migrations.to_a.size).to eq(5)
+      end
+    end
+
+    context 'running migrations', :freeze_time do
+      def define_background_migration(name)
+        klass = Class.new do
+          # Can't simply def perform here as we won't have access to the block,
+          # similarly can't define_method(:perform, &block) here as it would change the block receiver
+          define_method(:perform) { |*args| yield(*args) }
+        end
+        stub_const("Gitlab::BackgroundMigration::#{name}", klass)
+        klass
+      end
+
+      def expect_migration_call_counts(migrations_to_calls)
+        migrations_to_calls.each do |migration, calls|
+          expect_next_instances_of(migration, calls) do |m|
+            expect(m).to receive(:perform).and_call_original
+          end
+        end
+      end
+
+      it 'runs the migration class correctly' do
+        calls = []
+        define_background_migration(migration_name) do |i|
+          calls << i
+        end
+        described_class.new.run_jobs(for_duration: 1.second) # Any time would work here as we do not advance time
+        expect(calls).to contain_exactly(1, 2, 3, 4, 5)
+      end
+
+      it 'runs the migration for a uniform amount of time' do
+        migration = define_background_migration(migration_name) do |i|
+          travel(1.minute)
+        end
+
+        expect_migration_call_counts(migration => 3)
+
+        described_class.new.run_jobs(for_duration: 3.minutes)
+      end
+
+      context 'with multiple migrations to run' do
+        let(:other_migration_name) { 'OtherBackgroundMigration' }
+
+        before do
+          (1..5).each do |i|
+            migrate_in(i.minutes, other_migration_name, [i])
+          end
+        end
+
+        it 'splits the time between migrations when all migrations use all their time' do
+          migration = define_background_migration(migration_name) do |i|
+            travel(1.minute)
+          end
+
+          other_migration = define_background_migration(other_migration_name) do |i|
+            travel(2.minutes)
+          end
+
+          expect_migration_call_counts(
+            migration => 2, # 1 minute jobs for 90 seconds, can finish the first and start the second
+            other_migration => 1 # 2 minute jobs for 90 seconds, past deadline after a single job
+          )
+
+          described_class.new.run_jobs(for_duration: 3.minutes)
+        end
+
+        it 'does not give leftover time to extra migrations' do
+          # This is currently implemented this way for simplicity, but it could make sense to change this behavior.
+
+          migration = define_background_migration(migration_name) do
+            travel(1.second)
+          end
+          other_migration = define_background_migration(other_migration_name) do
+            travel(1.minute)
+          end
+          expect_migration_call_counts(
+            migration => 5,
+            other_migration => 2
+          )
+
+          described_class.new.run_jobs(for_duration: 3.minutes)
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/partitioning_spec.rb b/spec/lib/gitlab/database/partitioning_spec.rb
index 154cc2b7972..7c69f639aab 100644
--- a/spec/lib/gitlab/database/partitioning_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_spec.rb
@@ -109,6 +109,20 @@ RSpec.describe Gitlab::Database::Partitioning do
           .and change { find_partitions(table_names.last).size }.from(0)
       end
     end
+
+    context 'when only a specific database is requested' do
+      before do
+        allow(models.first).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(models.last).to receive_message_chain('connection_db_config.name').and_return('ci')
+      end
+
+      it 'manages partitions for models for the given database', :aggregate_failures do
+        expect { described_class.sync_partitions(models, only_on: 'ci') }
+          .to change { find_partitions(table_names.last).size }.from(0)
+
+        expect(find_partitions(table_names.first).size).to eq(0)
+      end
+    end
   end
 
   describe '.report_metrics' do
diff --git a/spec/lib/gitlab/database/query_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzer_spec.rb
index 34c72893c53..3b4cbc79de2 100644
--- a/spec/lib/gitlab/database/query_analyzer_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzer_spec.rb
@@ -4,6 +4,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   let(:analyzer) { double(:query_analyzer) }
+  let(:user_analyzer) { double(:query_analyzer) }
   let(:disabled_analyzer) { double(:disabled_query_analyzer) }
 
   before do
@@ -53,6 +54,10 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
 
         expect { |b| described_class.instance.within(&b) }.to yield_control
       end
+
+      it 'raises exception when trying to re-define analyzers' do
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to raise_error /Query analyzers are already defined, cannot re-define them/
+      end
     end
 
     context 'when initializer is enabled' do
@@ -75,6 +80,18 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
         expect { |b| described_class.instance.within(&b) }.to yield_control
       end
     end
+
+    context 'when user analyzers are used' do
+      it 'calls begin! and end!' do
+        expect(analyzer).not_to receive(:begin!)
+        allow(user_analyzer).to receive(:enabled?).and_return(true)
+        allow(user_analyzer).to receive(:suppressed?).and_return(false)
+        expect(user_analyzer).to receive(:begin!)
+        expect(user_analyzer).to receive(:end!)
+
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
+      end
+    end
   end
 
   describe '#process_sql' do
diff --git a/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb b/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb
new file mode 100644
index 00000000000..a2c7916fa01
--- /dev/null
+++ b/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb
@@ -0,0 +1,161 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas, query_analyzers: false do
+  let(:analyzer) { described_class }
+
+  context 'properly analyzes queries' do
+    using RSpec::Parameterized::TableSyntax
+
+    where do
+      examples = {
+        "for SELECT on projects" => {
+          sql: "SELECT 1 FROM projects",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :dml_not_allowed,
+            gitlab_main: :success,
+            gitlab_ci: :dml_access_denied # cross-schema access
+          }
+        },
+        "for INSERT" => {
+          sql: "INSERT INTO projects VALUES (1)",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :dml_not_allowed,
+            gitlab_main: :success,
+            gitlab_ci: :dml_access_denied # cross-schema access
+          }
+        },
+        "for CREATE INDEX" => {
+          sql: "CREATE INDEX index_projects_on_hidden ON projects (hidden)",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        },
+        "for CREATE SCHEMA" => {
+          sql: "CREATE SCHEMA __test_schema",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            # TODO: This is currently not properly detected
+            gitlab_main: :success,
+            gitlab_ci: :success
+          }
+        },
+        "for CREATE FUNCTION" => {
+          sql: "CREATE FUNCTION add(integer, integer) RETURNS integer AS 'select $1 + $2;' LANGUAGE SQL",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        },
+        "for CREATE TRIGGER" => {
+          sql: "CREATE TRIGGER check_projects BEFORE UPDATE ON projects FOR EACH ROW EXECUTE PROCEDURE check_projects_update()",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        }
+      }
+
+      # Expands all examples into individual tests
+      examples.flat_map do |name, configuration|
+        configuration[:expected_allowed_gitlab_schemas].map do |allowed_gitlab_schema, expectation|
+          [
+            "#{name} for allowed_gitlab_schema=#{allowed_gitlab_schema}",
+            {
+              sql: configuration[:sql],
+              allowed_gitlab_schema: allowed_gitlab_schema, # nil, gitlab_main
+              expectation: expectation # success, dml_access_denied, ...
+            }
+          ]
+        end
+      end.to_h
+    end
+
+    with_them do
+      subject do
+        process_sql(sql) do
+          analyzer.allowed_gitlab_schemas = [allowed_gitlab_schema] unless allowed_gitlab_schema == :no_schema
+        end
+      end
+
+      it do
+        case expectation
+        when :success
+          expect { subject }.not_to raise_error
+        when :ddl_not_allowed
+          expect { subject }.to raise_error(described_class::DDLNotAllowedError)
+        when :dml_not_allowed
+          expect { subject }.to raise_error(described_class::DMLNotAllowedError)
+        when :dml_access_denied
+          expect { subject }.to raise_error(described_class::DMLAccessDeniedError)
+        else
+          raise "invalid expectation: #{expectation}"
+        end
+      end
+    end
+  end
+
+  describe '.require_ddl_mode!' do
+    subject { described_class.require_ddl_mode! }
+
+    it "when not configured does not raise exception" do
+      expect { subject }.not_to raise_error
+    end
+
+    it "when no schemas are configured does not raise exception (DDL mode)" do
+      with_analyzer do
+        expect { subject }.not_to raise_error
+      end
+    end
+
+    it "with schemas configured does raise exception (DML mode)" do
+      with_analyzer do
+        analyzer.allowed_gitlab_schemas = %i[gitlab_main]
+
+        expect { subject }.to raise_error(described_class::DMLNotAllowedError)
+      end
+    end
+  end
+
+  describe '.require_dml_mode!' do
+    subject { described_class.require_dml_mode! }
+
+    it "when not configured does not raise exception" do
+      expect { subject }.not_to raise_error
+    end
+
+    it "when no schemas are configured does raise exception (DDL mode)" do
+      with_analyzer do
+        expect { subject }.to raise_error(described_class::DDLNotAllowedError)
+      end
+    end
+
+    it "with schemas configured does raise exception (DML mode)" do
+      with_analyzer do
+        analyzer.allowed_gitlab_schemas = %i[gitlab_main]
+
+        expect { subject }.not_to raise_error
+      end
+    end
+  end
+
+  def with_analyzer
+    Gitlab::Database::QueryAnalyzer.instance.within([analyzer]) do
+      yield
+    end
+  end
+
+  def process_sql(sql, model = ActiveRecord::Base)
+    with_analyzer do
+      yield if block_given?
+
+      # Skip load balancer and retrieve connection assigned to model
+      Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_projects_spec.rb b/spec/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_projects_spec.rb
index 3799fe3c316..50071e3e22b 100644
--- a/spec/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_projects_spec.rb
+++ b/spec/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_projects_spec.rb
@@ -49,7 +49,7 @@ RSpec.describe Gitlab::Database::RenameReservedPathsMigration::V1::RenameProject
 
     it 'invalidates the markdown cache of related projects' do
       expect(subject).to receive(:remove_cached_html_for_projects)
-                           .with(projects.map(&:id))
+                           .with(a_collection_containing_exactly(*projects.map(&:id)))
 
       subject.rename_projects
     end
diff --git a/spec/lib/gitlab/database/transaction/context_spec.rb b/spec/lib/gitlab/database/transaction/context_spec.rb
index 37cfc841d48..33a47150060 100644
--- a/spec/lib/gitlab/database/transaction/context_spec.rb
+++ b/spec/lib/gitlab/database/transaction/context_spec.rb
@@ -135,4 +135,24 @@ RSpec.describe Gitlab::Database::Transaction::Context do
 
     it_behaves_like 'logs transaction data'
   end
+
+  context 'when there are too many external HTTP requests' do
+    before do
+      allow(::Gitlab::Metrics::Subscribers::ExternalHttp)
+        .to receive(:request_count)
+        .and_return(100)
+    end
+
+    it_behaves_like 'logs transaction data'
+  end
+
+  context 'when there are too many too long external HTTP requests' do
+    before do
+      allow(::Gitlab::Metrics::Subscribers::ExternalHttp)
+        .to receive(:duration)
+        .and_return(5.5)
+    end
+
+    it_behaves_like 'logs transaction data'
+  end
 end
diff --git a/spec/lib/gitlab/database/transaction/observer_spec.rb b/spec/lib/gitlab/database/transaction/observer_spec.rb
index e5cc0106c9b..074c18d406e 100644
--- a/spec/lib/gitlab/database/transaction/observer_spec.rb
+++ b/spec/lib/gitlab/database/transaction/observer_spec.rb
@@ -25,7 +25,7 @@ RSpec.describe Gitlab::Database::Transaction::Observer do
           User.first
 
           expect(transaction_context).to be_a(::Gitlab::Database::Transaction::Context)
-          expect(context.keys).to match_array(%i(start_time depth savepoints queries backtraces))
+          expect(context.keys).to match_array(%i(start_time depth savepoints queries backtraces external_http_count_start external_http_duration_start))
           expect(context[:depth]).to eq(2)
           expect(context[:savepoints]).to eq(1)
           expect(context[:queries].length).to eq(1)
@@ -38,6 +38,71 @@ RSpec.describe Gitlab::Database::Transaction::Observer do
       expect(context[:backtraces].length).to eq(1)
     end
 
+    describe 'tracking external network requests', :request_store do
+      it 'tracks external requests' do
+        perform_stubbed_external_http_request(duration: 0.25)
+        perform_stubbed_external_http_request(duration: 1.25)
+
+        ActiveRecord::Base.transaction do
+          User.first
+
+          expect(context[:external_http_count_start]).to eq(2)
+          expect(context[:external_http_duration_start]).to eq(1.5)
+
+          perform_stubbed_external_http_request(duration: 1)
+          perform_stubbed_external_http_request(duration: 3)
+
+          expect(transaction_context.external_http_requests_count).to eq 2
+          expect(transaction_context.external_http_requests_duration).to eq 4
+        end
+      end
+
+      context 'when external HTTP requests duration has been exceeded' do
+        it 'logs transaction details including exceeding thresholds' do
+          expect(Gitlab::AppJsonLogger).to receive(:info).with(
+            hash_including(
+              external_http_requests_count: 2,
+              external_http_requests_duration: 12
+            )
+          )
+
+          ActiveRecord::Base.transaction do
+            User.first
+
+            perform_stubbed_external_http_request(duration: 2)
+            perform_stubbed_external_http_request(duration: 10)
+          end
+        end
+      end
+
+      context 'when external HTTP requests count has been exceeded' do
+        it 'logs transaction details including exceeding thresholds' do
+          expect(Gitlab::AppJsonLogger).to receive(:info).with(
+            hash_including(external_http_requests_count: 55)
+          )
+
+          ActiveRecord::Base.transaction do
+            User.first
+
+            55.times { perform_stubbed_external_http_request(duration: 0.01) }
+          end
+        end
+      end
+
+      def perform_stubbed_external_http_request(duration:)
+        ::Gitlab::Metrics::Subscribers::ExternalHttp.new.request(
+          instance_double(
+            'ActiveSupport::Notifications::Event',
+            payload: {
+              method: 'GET', code: '200', duration: duration,
+              scheme: 'http', host: 'example.gitlab.com', port: 80, path: '/'
+            },
+            time: Time.current
+          )
+        )
+      end
+    end
+
     describe '.extract_sql_command' do
       using RSpec::Parameterized::TableSyntax
 
diff --git a/spec/lib/gitlab/database/type/color_spec.rb b/spec/lib/gitlab/database/type/color_spec.rb
new file mode 100644
index 00000000000..84fd8d0bbce
--- /dev/null
+++ b/spec/lib/gitlab/database/type/color_spec.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::Gitlab::Database::Type::Color do
+  subject(:type) { described_class.new }
+
+  let(:color) { ::Gitlab::Color.of('red') }
+
+  it 'serializes by calling #to_s' do
+    expect(type.serialize(color)).to eq(color.to_s)
+  end
+
+  it 'serializes nil to nil' do
+    expect(type.serialize(nil)).to be_nil
+  end
+
+  it 'casts by calling Color::new' do
+    expect(type.cast('#fff')).to eq(::Gitlab::Color.new('#fff'))
+  end
+
+  it 'accepts colors as arguments to cast' do
+    expect(type.cast(color)).to eq(color)
+  end
+
+  it 'allows nil database values' do
+    expect(type.cast(nil)).to be_nil
+  end
+
+  it 'tells us what is serializable' do
+    [nil, 'foo', color].each do |value|
+      expect(type.serializable?(value)).to be true
+    end
+  end
+
+  it 'tells us what is not serializable' do
+    [0, 3.2, true, Time.current, { some: 'hash' }].each do |value|
+      expect(type.serializable?(value)).to be false
+    end
+  end
+end