Remove defaulted connections for BBM

Remove the connection defaults of ApplicationRecord.connection
throughout the implementation of backed background migrations. At this
point, we should only be using explicit connections passed by either the
calling worker or rake tasks.

lib/gitlab/database/background_migration/batched_migration_runner.rb

@@ -6,13 +6,13 @@ module Gitlab
       class BatchedMigrationRunner
         FailedToFinalize = Class.new(RuntimeError)
 
-        def self.finalize(job_class_name, table_name, column_name, job_arguments, connection: ApplicationRecord.connection)
+        def self.finalize(job_class_name, table_name, column_name, job_arguments, connection:)
           new(connection: connection).finalize(job_class_name, table_name, column_name, job_arguments)
         end
 
-        def initialize(migration_wrapper = BatchedMigrationWrapper.new, connection: ApplicationRecord.connection)
-          @migration_wrapper = migration_wrapper
+        def initialize(connection:, migration_wrapper: BatchedMigrationWrapper.new(connection: connection))
           @connection = connection
+          @migration_wrapper = migration_wrapper
         end
 
         # Runs the next batched_job for a batched_background_migration.
@@ -79,7 +79,7 @@ module Gitlab
 
         private
 
-        attr_reader :migration_wrapper, :connection
+        attr_reader :connection, :migration_wrapper
 
         def find_or_create_next_batched_job(active_migration)
           if next_batch_range = find_next_batch_range(active_migration)

lib/gitlab/database/background_migration/batched_migration_wrapper.rb

@@ -4,7 +4,7 @@ module Gitlab
   module Database
     module BackgroundMigration
       class BatchedMigrationWrapper
-        def initialize(connection: ApplicationRecord.connection, metrics: PrometheusMetrics.new)
+        def initialize(connection:, metrics: PrometheusMetrics.new)
           @connection = connection
           @metrics = metrics
         end

lib/tasks/gitlab/db.rake

@@ -340,9 +340,11 @@ namespace :gitlab do
 
     desc 'Run all pending batched migrations'
     task execute_batched_migrations: :environment do
-      Gitlab::Database::BackgroundMigration::BatchedMigration.with_status(:active).queue_order.each do |migration|
-        Gitlab::AppLogger.info("Executing batched migration #{migration.id} inline")
-        Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.new.run_entire_migration(migration)
+      Gitlab::Database::EachDatabase.each_database_connection do |connection, name|
+        Gitlab::Database::BackgroundMigration::BatchedMigration.with_status(:active).queue_order.each do |migration|
+          Gitlab::AppLogger.info("Executing batched migration #{migration.id} on database #{name} inline")
+          Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.new(connection: connection).run_entire_migration(migration)
+        end
       end
     end
 

spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb

@@ -3,8 +3,16 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
+  let(:connection) { Gitlab::Database.database_base_models[:main].connection }
   let(:migration_wrapper) { double('test wrapper') }
-  let(:runner) { described_class.new(migration_wrapper) }
+
+  let(:runner) { described_class.new(connection: connection, migration_wrapper: migration_wrapper) }
+
+  around do |example|
+    Gitlab::Database::SharedModel.using_connection(connection) do
+      example.run
+    end
+  end
 
   describe '#run_migration_job' do
     shared_examples_for 'it has completed the migration' do
@@ -295,7 +303,9 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
   end
 
   describe '#finalize' do
-    let(:migration_wrapper) { Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new }
+    let(:migration_wrapper) do
+      Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection)
+    end
 
     let(:migration_helpers) { ActiveRecord::Migration.new }
     let(:table_name) { :_test_batched_migrations_test_table }
@@ -443,8 +453,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
 
   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] }

spec/lib/gitlab/database/background_migration/batched_migration_wrapper_spec.rb

@@ -3,8 +3,9 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '#perform' do
-  subject { described_class.new(metrics: metrics_tracker).perform(job_record) }
+  subject { described_class.new(connection: connection, metrics: metrics_tracker).perform(job_record) }
 
+  let(:connection) { Gitlab::Database.database_base_models[:main].connection }
   let(:metrics_tracker) { instance_double('::Gitlab::Database::BackgroundMigration::PrometheusMetrics', track: nil) }
   let(:job_class) { Gitlab::BackgroundMigration::CopyColumnUsingBackgroundMigrationJob }
 
@@ -20,6 +21,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
 
   let(:job_instance) { double('job instance', batch_metrics: {}) }
 
+  around do |example|
+    Gitlab::Database::SharedModel.using_connection(connection) do
+      example.run
+    end
+  end
+
   before do
     allow(job_class).to receive(:new).and_return(job_instance)
   end
@@ -138,7 +145,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
       stub_const('Gitlab::BackgroundMigration::Foo', migration_class)
     end
 
-    let(:connection) { double(:connection) }
     let(:active_migration) { create(:batched_background_migration, :active, job_class_name: 'Foo') }
     let!(:job_record) { create(:batched_background_migration_job, batched_migration: active_migration) }
 
@@ -147,12 +153,11 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
 
       expect(job_instance).to receive(:perform)
 
-      described_class.new(connection: connection).perform(job_record)
+      subject
     end
   end
 
   context 'when the batched background migration inherits from BaseJob' do
-    let(:connection) { double(:connection) }
     let(:active_migration) { create(:batched_background_migration, :active, job_class_name: 'Foo') }
     let!(:job_record) { create(:batched_background_migration_job, batched_migration: active_migration) }
 
@@ -167,7 +172,7 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper, '
 
       expect(job_instance).to receive(:perform)
 
-      described_class.new(connection: connection).perform(job_record)
+      subject
     end
   end
 end

spec/tasks/gitlab/db_rake_spec.rb

@@ -708,22 +708,51 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe '#execute_batched_migrations' do
-    subject { run_rake_task('gitlab:db:execute_batched_migrations') }
+    subject(:execute_batched_migrations) { run_rake_task('gitlab:db:execute_batched_migrations') }
 
-    let(:migrations) { create_list(:batched_background_migration, 2, :active) }
-    let(:runner) { instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner') }
+    let(:connections) do
+      {
+        main: instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter),
+        ci: instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+      }
+    end
+
+    let(:runners) do
+      {
+        main: instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner'),
+        ci: instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner')
+      }
+    end
+
+    let(:migrations) do
+      {
+        main: build_list(:batched_background_migration, 1),
+        ci: build_list(:batched_background_migration, 1)
+      }
+    end
 
     before do
-      allow(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive_message_chain(:active, :queue_order).and_return(migrations)
-      allow(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:new).and_return(runner)
+      each_database = class_double('Gitlab::Database::EachDatabase').as_stubbed_const
+
+      allow(each_database).to receive(:each_database_connection)
+        .and_yield(connections[:main], 'main')
+        .and_yield(connections[:ci], 'ci')
+
+      keys = migrations.keys
+      allow(Gitlab::Database::BackgroundMigration::BatchedMigration)
+        .to receive_message_chain(:with_status, :queue_order) { migrations[keys.shift] }
     end
 
     it 'executes all migrations' do
-      migrations.each do |migration|
-        expect(runner).to receive(:run_entire_migration).with(migration)
+      [:main, :ci].each do |name|
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:new)
+          .with(connection: connections[name])
+          .and_return(runners[name])
+
+        expect(runners[name]).to receive(:run_entire_migration).with(migrations[name].first)
       end
 
-      subject
+      execute_batched_migrations
     end
   end