Support multiple databases for batched background migrations

Update CopyColumnUsingBackgroundMigrationJob to use the right connection

Changelog: added

Relates to gitlab-org/gitlab#351585

lib/gitlab/background_migration/copy_column_using_background_migration_job.rb

@@ -13,7 +13,7 @@ module BackgroundMigration
     # - We skip the NULL checks as they may result in not using an index scan
     # - The table that is migrated does _not_ need `id` as the primary key
     #   We use the provided primary_key column to perform the update.
-    class CopyColumnUsingBackgroundMigrationJob
+    class CopyColumnUsingBackgroundMigrationJob < BaseJob
       include Gitlab::Database::DynamicModelHelpers
 
       # start_id - The start ID of the range of rows to update.
@@ -52,10 +52,6 @@ def batch_metrics
 
       private
 
-      def connection
-        ActiveRecord::Base.connection
-      end
-
       def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
         define_batchable_model(source_table, connection: connection).where(source_key_column => start_id..stop_id)
       end

lib/gitlab/database/background_migration/batched_migration_wrapper.rb

@@ -6,6 +6,10 @@ module BackgroundMigration
       class BatchedMigrationWrapper
         extend Gitlab::Utils::StrongMemoize
 
+        def initialize(connection: ApplicationRecord.connection)
+          @connection = connection
+        end
+
         # Wraps the execution of a batched_background_migration.
         #
         # Updates the job's tracking records with the status of the migration
@@ -29,12 +33,14 @@ def perform(batch_tracking_record)
 
         private
 
+        attr_reader :connection
+
         def start_tracking_execution(tracking_record)
           tracking_record.run!
         end
 
         def execute_batch(tracking_record)
-          job_instance = tracking_record.migration_job_class.new
+          job_instance = migration_instance_for(tracking_record.migration_job_class)
 
           job_instance.perform(
             tracking_record.min_value,
@@ -50,6 +56,14 @@ def execute_batch(tracking_record)
           end
         end
 
+        def migration_instance_for(job_class)
+          if job_class < Gitlab::BackgroundMigration::BaseJob
+            job_class.new(connection: connection)
+          else
+            job_class.new
+          end
+        end
+
         def track_prometheus_metrics(tracking_record)
           migration = tracking_record.batched_migration
           base_labels = migration.prometheus_labels

spec/lib/gitlab/background_migration/copy_column_using_background_migration_job_spec.rb

@@ -7,13 +7,14 @@
   let(:test_table) { table(table_name) }
   let(:sub_batch_size) { 1000 }
   let(:pause_ms) { 0 }
+  let(:connection) { ApplicationRecord.connection }
 
   let(:helpers) do
     ActiveRecord::Migration.new.extend(Gitlab::Database::MigrationHelpers)
   end
 
   before do
-    ActiveRecord::Base.connection.execute(<<~SQL)
+    connection.execute(<<~SQL)
       CREATE TABLE #{table_name}
       (
        id integer NOT NULL,
@@ -34,12 +35,14 @@
 
   after do
     # Make sure that the temp table we created is dropped (it is not removed by the database_cleaner)
-    ActiveRecord::Base.connection.execute(<<~SQL)
+    connection.execute(<<~SQL)
       DROP TABLE IF EXISTS #{table_name};
     SQL
   end
 
-  subject(:copy_columns) { described_class.new }
+  subject(:copy_columns) { described_class.new(connection: connection) }
+
+  it { expect(described_class).to be < Gitlab::BackgroundMigration::BaseJob }
 
   describe '#perform' do
     let(:migration_class) { described_class.name }

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

@@ -194,4 +194,44 @@
     it_behaves_like 'an error is raised', RuntimeError.new('Something broke!')
     it_behaves_like 'an error is raised', SignalException.new('SIGTERM')
   end
+
+  context 'when the batched background migration does not inherit from BaseJob' do
+    let(:migration_class) { Class.new }
+
+    before do
+      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) }
+
+    it 'does not pass any argument' do
+      expect(Gitlab::BackgroundMigration::Foo).to receive(:new).with(no_args).and_return(job_instance)
+
+      expect(job_instance).to receive(:perform)
+
+      described_class.new(connection: connection).perform(job_record)
+    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) }
+
+    let(:migration_class) { Class.new(::Gitlab::BackgroundMigration::BaseJob) }
+
+    before do
+      stub_const('Gitlab::BackgroundMigration::Foo', migration_class)
+    end
+
+    it 'passes the correct connection' do
+      expect(Gitlab::BackgroundMigration::Foo).to receive(:new).with(connection: connection).and_return(job_instance)
+
+      expect(job_instance).to receive(:perform)
+
+      described_class.new(connection: connection).perform(job_record)
+    end
+  end
 end