Add method to finalize batched background migration

https://gitlab.com/gitlab-org/gitlab/-/issues/292874

lib/gitlab/database/background_migration/batched_migration.rb

@@ -29,11 +29,16 @@ module Gitlab
           paused: 0,
           active: 1,
           finished: 3,
-          failed: 4
+          failed: 4,
+          finalizing: 5
         }
 
         attribute :pause_ms, :integer, default: 100
 
+        def self.find_for_configuration(job_class_name, table_name, column_name, job_arguments)
+          for_configuration(job_class_name, table_name, column_name, job_arguments).first
+        end
+
         def self.active_migration
           active.queue_order.first
         end

lib/gitlab/database/background_migration/batched_migration_runner.rb

@@ -4,6 +4,12 @@ module Gitlab
   module Database
     module BackgroundMigration
       class BatchedMigrationRunner
+        FailedToFinalize = Class.new(RuntimeError)
+
+        def self.finalize(job_class_name, table_name, column_name, job_arguments)
+          new.finalize(job_class_name, table_name, column_name, job_arguments)
+        end
+
         def initialize(migration_wrapper = BatchedMigrationWrapper.new)
           @migration_wrapper = migration_wrapper
         end
@@ -37,10 +43,34 @@ module Gitlab
             raise 'this method is not intended for use in real environments'
           end
 
-          while migration.active?
-            run_migration_job(migration)
+          run_migration_while(migration, :active)
+        end
 
-            migration.reload_last_job
+        # Finalize migration for given configuration.
+        #
+        # If the migration is already finished, do nothing. Otherwise change its status to `finalizing`
+        # in order to prevent it being picked up by the background worker. Perform all pending jobs,
+        # then keep running until migration is finished.
+        def finalize(job_class_name, table_name, column_name, job_arguments)
+          migration = BatchedMigration.find_for_configuration(job_class_name, table_name, column_name, job_arguments)
+
+          if migration.nil?
+            configuration = {
+              job_class_name: job_class_name,
+              table_name: table_name,
+              column_name: column_name,
+              job_arguments: job_arguments
+            }
+            Gitlab::AppLogger.warn "Could not find batched background migration for the given configuration: #{configuration}"
+          else
+            return if migration.finished?
+
+            migration.finalizing!
+            migration.batched_jobs.pending.each { |job| migration_wrapper.perform(job) }
+
+            run_migration_while(migration, :finalizing)
+
+            raise FailedToFinalize unless migration.finished?
           end
         end
 
@@ -90,6 +120,14 @@ module Gitlab
             active_migration.finished!
           end
         end
+
+        def run_migration_while(migration, status)
+          while migration.status == status.to_s
+            run_migration_job(migration)
+
+            migration.reload_last_job
+          end
+        end
       end
     end
   end

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

@@ -281,4 +281,142 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       end
     end
   end
+
+  describe '#finalize' do
+    let(:migration_wrapper) { Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new }
+
+    let(:migration_helpers) { ActiveRecord::Migration.new }
+    let(:table_name) { :_batched_migrations_test_table }
+    let(:column_name) { :some_id }
+    let(:job_arguments) { [:some_id, :some_id_convert_to_bigint] }
+
+    let(:migration_status) { :active }
+
+    let!(:batched_migration) do
+      create(
+        :batched_background_migration,
+        status: migration_status,
+        max_value: 8,
+        batch_size: 2,
+        sub_batch_size: 1,
+        interval: 0,
+        table_name: table_name,
+        column_name: column_name,
+        job_arguments: job_arguments,
+        pause_ms: 0
+      )
+    end
+
+    before do
+      migration_helpers.drop_table table_name, if_exists: true
+      migration_helpers.create_table table_name, id: false do |t|
+        t.integer :some_id, primary_key: true
+        t.integer :some_id_convert_to_bigint
+      end
+
+      migration_helpers.execute("INSERT INTO #{table_name} VALUES (1, 1), (2, 2), (3, NULL), (4, NULL), (5, NULL), (6, NULL), (7, NULL), (8, NULL)")
+    end
+
+    after do
+      migration_helpers.drop_table table_name, if_exists: true
+    end
+
+    context 'when the migration is not yet completed' do
+      before do
+        common_attributes = {
+          batched_migration: batched_migration,
+          batch_size: 2,
+          sub_batch_size: 1,
+          pause_ms: 0
+        }
+
+        create(:batched_background_migration_job, common_attributes.merge(status: :succeeded, min_value: 1, max_value: 2))
+        create(:batched_background_migration_job, common_attributes.merge(status: :pending, min_value: 3, max_value: 4))
+        create(:batched_background_migration_job, common_attributes.merge(status: :failed, min_value: 5, max_value: 6, attempts: 1))
+      end
+
+      it 'completes the migration' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .and_return(batched_migration)
+
+        expect(batched_migration).to receive(:finalizing!).and_call_original
+
+        expect do
+          runner.finalize(
+            batched_migration.job_class_name,
+            table_name,
+            column_name,
+            job_arguments
+          )
+        end.to change { batched_migration.reload.status }.from('active').to('finished')
+
+        expect(batched_migration.batched_jobs).to all(be_succeeded)
+
+        not_converted = migration_helpers.execute("SELECT * FROM #{table_name} WHERE some_id_convert_to_bigint IS NULL")
+        expect(not_converted.to_a).to be_empty
+      end
+
+      context 'when migration fails to complete' do
+        it 'raises an error' do
+          batched_migration.batched_jobs.failed.update_all(attempts: Gitlab::Database::BackgroundMigration::BatchedJob::MAX_ATTEMPTS)
+
+          expect do
+            runner.finalize(
+              batched_migration.job_class_name,
+              table_name,
+              column_name,
+              job_arguments
+            )
+          end.to raise_error described_class::FailedToFinalize
+        end
+      end
+    end
+
+    context 'when the migration is already finished' do
+      let(:migration_status) { :finished }
+
+      it 'is a no-op' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .and_return(batched_migration)
+
+        expect(batched_migration).not_to receive(:finalizing!)
+
+        runner.finalize(
+          batched_migration.job_class_name,
+          table_name,
+          column_name,
+          job_arguments
+        )
+      end
+    end
+
+    context 'when the migration does not exist' do
+      it 'is a no-op' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, [:some, :other, :arguments])
+          .and_return(nil)
+
+        configuration = {
+          job_class_name: batched_migration.job_class_name,
+          table_name: table_name.to_sym,
+          column_name: column_name.to_sym,
+          job_arguments: [:some, :other, :arguments]
+        }
+
+        expect(Gitlab::AppLogger).to receive(:warn)
+          .with("Could not find batched background migration for the given configuration: #{configuration}")
+
+        expect(batched_migration).not_to receive(:finalizing!)
+
+        runner.finalize(
+          batched_migration.job_class_name,
+          table_name,
+          column_name,
+          [:some, :other, :arguments]
+        )
+      end
+    end
+  end
 end

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

@@ -387,4 +387,22 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
       expect(actual).to contain_exactly(migration)
     end
   end
+
+  describe '.find_for_configuration' do
+    it 'returns nill if such migration does not exists' do
+      expect(described_class.find_for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to be_nil
+    end
+
+    it 'returns the migration when it exists' do
+      migration = create(
+        :batched_background_migration,
+        job_class_name: 'MyJobClass',
+        table_name: :projects,
+        column_name: :id,
+        job_arguments: [[:id], [:id_convert_to_bigint]]
+      )
+
+      expect(described_class.find_for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])).to eq(migration)
+    end
+  end
 end