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

9 files changed, 190 insertions, 251 deletions

diff --git a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
index e42a6c970ea..99c7d70724c 100644
--- a/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migrations/background_migration_helpers_spec.rb
@@ -7,78 +7,6 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     ActiveRecord::Migration.new.extend(described_class)
   end
 
-  describe '#bulk_queue_background_migration_jobs_by_range' do
-    context 'when the model has an ID column' do
-      let!(:id1) { create(:user).id }
-      let!(:id2) { create(:user).id }
-      let!(:id3) { create(:user).id }
-
-      before do
-        User.class_eval do
-          include EachBatch
-        end
-      end
-
-      context 'with enough rows to bulk queue jobs more than once' do
-        before do
-          stub_const('Gitlab::Database::Migrations::BackgroundMigrationHelpers::JOB_BUFFER_SIZE', 1)
-        end
-
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          end
-        end
-
-        it 'queues jobs in groups of buffer size 1' do
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id1, id2]]])
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id3, id3]]])
-
-          model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-        end
-      end
-
-      context 'with not enough rows to bulk queue jobs more than once' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id2]])
-            expect(BackgroundMigrationWorker.jobs[1]['args']).to eq(['FooJob', [id3, id3]])
-          end
-        end
-
-        it 'queues jobs in bulk all at once (big buffer size)' do
-          expect(BackgroundMigrationWorker).to receive(:bulk_perform_async).with([['FooJob', [id1, id2]],
-                                                                                  ['FooJob', [id3, id3]]])
-
-          model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob', batch_size: 2)
-        end
-      end
-
-      context 'without specifying batch_size' do
-        it 'queues jobs correctly' do
-          Sidekiq::Testing.fake! do
-            model.bulk_queue_background_migration_jobs_by_range(User, 'FooJob')
-
-            expect(BackgroundMigrationWorker.jobs[0]['args']).to eq(['FooJob', [id1, id3]])
-          end
-        end
-      end
-    end
-
-    context "when the model doesn't have an ID column" do
-      it 'raises error (for now)' do
-        expect do
-          model.bulk_queue_background_migration_jobs_by_range(ProjectAuthorization, 'FooJob')
-        end.to raise_error(StandardError, /does not have an ID/)
-      end
-    end
-  end
-
   describe '#queue_background_migration_jobs_by_range_at_intervals' do
     context 'when the model has an ID column' do
       let!(:id1) { create(:user).id }
@@ -354,161 +282,6 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
     end
   end
 
-  describe '#queue_batched_background_migration' do
-    let(:pgclass_info) { instance_double('Gitlab::Database::PgClass', cardinality_estimate: 42) }
-
-    before do
-      allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
-    end
-
-    context 'when such migration already exists' do
-      it 'does not create duplicate migration' do
-        create(
-          :batched_background_migration,
-          job_class_name: 'MyJobClass',
-          table_name: :projects,
-          column_name: :id,
-          interval: 10.minutes,
-          min_value: 5,
-          max_value: 1005,
-          batch_class_name: 'MyBatchClass',
-          batch_size: 200,
-          sub_batch_size: 20,
-          job_arguments: [[:id], [:id_convert_to_bigint]]
-        )
-
-        expect do
-          model.queue_batched_background_migration(
-            'MyJobClass',
-            :projects,
-            :id,
-            [:id], [:id_convert_to_bigint],
-            job_interval: 5.minutes,
-            batch_min_value: 5,
-            batch_max_value: 1000,
-            batch_class_name: 'MyBatchClass',
-            batch_size: 100,
-            sub_batch_size: 10)
-        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
-      end
-    end
-
-    it 'creates the database record for the migration' do
-      expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info)
-
-      expect do
-        model.queue_batched_background_migration(
-          'MyJobClass',
-          :projects,
-          :id,
-          job_interval: 5.minutes,
-          batch_min_value: 5,
-          batch_max_value: 1000,
-          batch_class_name: 'MyBatchClass',
-          batch_size: 100,
-          sub_batch_size: 10)
-      end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-      expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
-        job_class_name: 'MyJobClass',
-        table_name: 'projects',
-        column_name: 'id',
-        interval: 300,
-        min_value: 5,
-        max_value: 1000,
-        batch_class_name: 'MyBatchClass',
-        batch_size: 100,
-        sub_batch_size: 10,
-        job_arguments: %w[],
-        status: 'active',
-        total_tuple_count: pgclass_info.cardinality_estimate)
-    end
-
-    context 'when the job interval is lower than the minimum' do
-      let(:minimum_delay) { described_class::BATCH_MIN_DELAY }
-
-      it 'sets the job interval to the minimum value' do
-        expect do
-          model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: minimum_delay - 1.minute)
-        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-        created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-        expect(created_migration.interval).to eq(minimum_delay)
-      end
-    end
-
-    context 'when additional arguments are passed to the method' do
-      it 'saves the arguments on the database record' do
-        expect do
-          model.queue_batched_background_migration(
-            'MyJobClass',
-            :projects,
-            :id,
-            'my',
-            'arguments',
-            job_interval: 5.minutes,
-            batch_max_value: 1000)
-        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-        expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
-          job_class_name: 'MyJobClass',
-          table_name: 'projects',
-          column_name: 'id',
-          interval: 300,
-          min_value: 1,
-          max_value: 1000,
-          job_arguments: %w[my arguments])
-      end
-    end
-
-    context 'when the max_value is not given' do
-      context 'when records exist in the database' do
-        let!(:event1) { create(:event) }
-        let!(:event2) { create(:event) }
-        let!(:event3) { create(:event) }
-
-        it 'creates the record with the current max value' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-          expect(created_migration.max_value).to eq(event3.id)
-        end
-
-        it 'creates the record with an active status' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_active
-        end
-      end
-
-      context 'when the database is empty' do
-        it 'sets the max value to the min value' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
-
-          expect(created_migration.max_value).to eq(created_migration.min_value)
-        end
-
-        it 'creates the record with a finished status' do
-          expect do
-            model.queue_batched_background_migration('MyJobClass', :projects, :id, job_interval: 5.minutes)
-          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
-
-          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_finished
-        end
-      end
-    end
-  end
-
   describe '#migrate_async' do
     it 'calls BackgroundMigrationWorker.perform_async' do
       expect(BackgroundMigrationWorker).to receive(:perform_async).with("Class", "hello", "world")
@@ -583,7 +356,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
   end
 
   describe '#finalized_background_migration' do
-    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(:main, BackgroundMigrationWorker) }
+    let(:job_coordinator) { Gitlab::BackgroundMigration::JobCoordinator.new(BackgroundMigrationWorker) }
 
     let!(:job_class_name) { 'TestJob' }
     let!(:job_class) { Class.new }
@@ -605,7 +378,7 @@ RSpec.describe Gitlab::Database::Migrations::BackgroundMigrationHelpers do
       job_class.define_method(:perform, job_perform_method)
 
       allow(Gitlab::BackgroundMigration).to receive(:coordinator_for_database)
-        .with(:main).and_return(job_coordinator)
+        .with('main').and_return(job_coordinator)
 
       expect(job_coordinator).to receive(:migration_class_for)
         .with(job_class_name).at_least(:once) { job_class }
diff --git a/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb b/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
new file mode 100644
index 00000000000..c45149d67bf
--- /dev/null
+++ b/spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb
@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers do
+  let(:migration) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  describe '#queue_batched_background_migration' do
+    let(:pgclass_info) { instance_double('Gitlab::Database::PgClass', cardinality_estimate: 42) }
+
+    before do
+      allow(Gitlab::Database::PgClass).to receive(:for_table).and_call_original
+    end
+
+    context 'when such migration already exists' do
+      it 'does not create duplicate migration' do
+        create(
+          :batched_background_migration,
+          job_class_name: 'MyJobClass',
+          table_name: :projects,
+          column_name: :id,
+          interval: 10.minutes,
+          min_value: 5,
+          max_value: 1005,
+          batch_class_name: 'MyBatchClass',
+          batch_size: 200,
+          sub_batch_size: 20,
+          job_arguments: [[:id], [:id_convert_to_bigint]]
+        )
+
+        expect do
+          migration.queue_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            [:id], [:id_convert_to_bigint],
+            job_interval: 5.minutes,
+            batch_min_value: 5,
+            batch_max_value: 1000,
+            batch_class_name: 'MyBatchClass',
+            batch_size: 100,
+            sub_batch_size: 10)
+        end.not_to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }
+      end
+    end
+
+    it 'creates the database record for the migration' do
+      expect(Gitlab::Database::PgClass).to receive(:for_table).with(:projects).and_return(pgclass_info)
+
+      expect do
+        migration.queue_batched_background_migration(
+          'MyJobClass',
+          :projects,
+          :id,
+          job_interval: 5.minutes,
+          batch_min_value: 5,
+          batch_max_value: 1000,
+          batch_class_name: 'MyBatchClass',
+          batch_size: 100,
+          sub_batch_size: 10)
+      end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+      expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
+        job_class_name: 'MyJobClass',
+        table_name: 'projects',
+        column_name: 'id',
+        interval: 300,
+        min_value: 5,
+        max_value: 1000,
+        batch_class_name: 'MyBatchClass',
+        batch_size: 100,
+        sub_batch_size: 10,
+        job_arguments: %w[],
+        status: 'active',
+        total_tuple_count: pgclass_info.cardinality_estimate)
+    end
+
+    context 'when the job interval is lower than the minimum' do
+      let(:minimum_delay) { described_class::BATCH_MIN_DELAY }
+
+      it 'sets the job interval to the minimum value' do
+        expect do
+          migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: minimum_delay - 1.minute)
+        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+        created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+        expect(created_migration.interval).to eq(minimum_delay)
+      end
+    end
+
+    context 'when additional arguments are passed to the method' do
+      it 'saves the arguments on the database record' do
+        expect do
+          migration.queue_batched_background_migration(
+            'MyJobClass',
+            :projects,
+            :id,
+            'my',
+            'arguments',
+            job_interval: 5.minutes,
+            batch_max_value: 1000)
+        end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to have_attributes(
+          job_class_name: 'MyJobClass',
+          table_name: 'projects',
+          column_name: 'id',
+          interval: 300,
+          min_value: 1,
+          max_value: 1000,
+          job_arguments: %w[my arguments])
+      end
+    end
+
+    context 'when the max_value is not given' do
+      context 'when records exist in the database' do
+        let!(:event1) { create(:event) }
+        let!(:event2) { create(:event) }
+        let!(:event3) { create(:event) }
+
+        it 'creates the record with the current max value' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+          expect(created_migration.max_value).to eq(event3.id)
+        end
+
+        it 'creates the record with an active status' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_active
+        end
+      end
+
+      context 'when the database is empty' do
+        it 'sets the max value to the min value' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :events, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          created_migration = Gitlab::Database::BackgroundMigration::BatchedMigration.last
+
+          expect(created_migration.max_value).to eq(created_migration.min_value)
+        end
+
+        it 'creates the record with a finished status' do
+          expect do
+            migration.queue_batched_background_migration('MyJobClass', :projects, :id, job_interval: 5.minutes)
+          end.to change { Gitlab::Database::BackgroundMigration::BatchedMigration.count }.by(1)
+
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration.last).to be_finished
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
index 841d2a98a16..902d8e13a63 100644
--- a/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
+++ b/spec/lib/gitlab/database/migrations/instrumentation_spec.rb
@@ -3,6 +3,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Instrumentation do
   let(:result_dir) { Dir.mktmpdir }
+  let(:connection) { ActiveRecord::Migration.connection }
 
   after do
     FileUtils.rm_rf(result_dir)
@@ -14,11 +15,11 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     let(:migration_version) { '12345' }
 
     it 'executes the given block' do
-      expect { |b| subject.observe(version: migration_version, name: migration_name, &b) }.to yield_control
+      expect { |b| subject.observe(version: migration_version, name: migration_name, connection: connection, &b) }.to yield_control
     end
 
     context 'behavior with observers' do
-      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
       let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
 
@@ -29,7 +30,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       it 'instantiates observer with observation' do
         expect(Gitlab::Database::Migrations::Observers::MigrationObserver)
           .to receive(:new)
-          .with(instance_of(Gitlab::Database::Migrations::Observation), anything) { |observation| expect(observation.version).to eq(migration_version) }
+          .with(instance_of(Gitlab::Database::Migrations::Observation), anything, connection) { |observation| expect(observation.version).to eq(migration_version) }
           .and_return(observer)
 
         subject
@@ -63,7 +64,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'on successful execution' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
       it 'records walltime' do
         expect(subject.walltime).not_to be_nil
@@ -83,7 +84,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'upon failure' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) { raise 'something went wrong' } }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' } }
 
       it 'raises the exception' do
         expect { subject }.to raise_error(/something went wrong/)
@@ -93,7 +94,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
         subject { instance.observations.first }
 
         before do
-          instance.observe(version: migration_version, name: migration_name) { raise 'something went wrong' }
+          instance.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' }
         rescue StandardError
           # ignore
         end
@@ -125,8 +126,8 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       let(:migration2) { double('migration2', call: nil) }
 
       it 'records observations for all migrations' do
-        subject.observe(version: migration_version, name: migration_name) {}
-        subject.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } rescue nil
+        subject.observe(version: migration_version, name: migration_name, connection: connection) {}
+        subject.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' } rescue nil
 
         expect(subject.observations.size).to eq(2)
       end
diff --git a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
index 191ac29e3b3..5a19ae6581d 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_details_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
-  subject { described_class.new(observation, directory_path) }
+  subject { described_class.new(observation, directory_path, connection) }
 
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
-  let(:connection) { ActiveRecord::Base.connection }
   let(:query) { "select date_trunc('day', $1::timestamptz) + $2 * (interval '1 hour')" }
   let(:query_binds) { [Time.current, 3] }
   let(:directory_path) { Dir.mktmpdir }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
index 2e70a85fd5b..7b01e39f5f1 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_log_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
-  subject { described_class.new(observation, directory_path) }
+  subject { described_class.new(observation, directory_path, connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select 1' }
   let(:directory_path) { Dir.mktmpdir }
   let(:migration_version) { 20210422152437 }
diff --git a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
index 9727a215d71..2515f0d4a06 100644
--- a/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
-  subject { described_class.new(observation, double("unused path")) }
+  subject { described_class.new(observation, double("unused path"), connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
 
   def mock_pgss(enabled: true)
     if enabled
diff --git a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
index e689759c574..4b08838d6bb 100644
--- a/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb
@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
-  subject { described_class.new(observation, double('unused path')) }
+  subject { described_class.new(observation, double('unused path'), connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select pg_database_size(current_database())' }
 
   it 'records the size change' do
diff --git a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
index e65f89747c4..b26bb8fbe41 100644
--- a/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
+++ b/spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb
@@ -2,8 +2,9 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
-  subject(:transaction_duration_observer) { described_class.new(observation, directory_path) }
+  subject(:transaction_duration_observer) { described_class.new(observation, directory_path, connection) }
 
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
   let(:directory_path) { Dir.mktmpdir }
   let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
@@ -78,17 +79,17 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
   end
 
   def run_real_transactions
-    ActiveRecord::Base.transaction do
+    ApplicationRecord.transaction do
     end
   end
 
   def run_sub_transactions
-    ActiveRecord::Base.transaction(requires_new: true) do
+    ApplicationRecord.transaction(requires_new: true) do
     end
   end
 
   def run_transaction
-    ActiveRecord::Base.connection_pool.with_connection do |connection|
+    ApplicationRecord.connection_pool.with_connection do |connection|
       Gitlab::Database::SharedModel.using_connection(connection) do
         Gitlab::Database::SharedModel.transaction do
           Gitlab::Database::SharedModel.transaction(requires_new: true) do
diff --git a/spec/lib/gitlab/database/migrations/runner_spec.rb b/spec/lib/gitlab/database/migrations/runner_spec.rb
index 52fb5ec2ba8..4616bd6941e 100644
--- a/spec/lib/gitlab/database/migrations/runner_spec.rb
+++ b/spec/lib/gitlab/database/migrations/runner_spec.rb
@@ -76,7 +76,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       it 'runs the unapplied migrations in version order', :aggregate_failures do
         up.run
 
-        expect(migration_runs.map(&:dir)).to eq([:up, :up])
+        expect(migration_runs.map(&:dir)).to match_array([:up, :up])
         expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
       end
     end
@@ -101,7 +101,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       it 'runs the applied migrations for the current branch in reverse order', :aggregate_failures do
         down.run
 
-        expect(migration_runs.map(&:dir)).to eq([:down, :down])
+        expect(migration_runs.map(&:dir)).to match_array([:down, :down])
         expect(migration_runs.map(&:version_to_migrate)).to eq(applied_migrations_this_branch.reverse.map(&:version))
       end
     end