Merge branch 'pb-remove-default-bbm-connections' into 'master'

Remove defaulted connections for BBM

See merge request gitlab-org/gitlab!83645

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