Add steal method for BatchedBackgroundMigrationHelpers

- Add support for inline execution for batched background migrations
- Run inline execution before ensuring that the migration is completed

Changelog: added

Relates to https://gitlab.com/gitlab-org/gitlab/-/merge_requests/86551

lib/gitlab/database/migration_helpers.rb

@@ -943,7 +943,7 @@ module Gitlab
         execute("DELETE FROM batched_background_migrations WHERE #{conditions}")
       end
-      def ensure_batched_background_migration_is_finished(job_class_name:, table_name:, column_name:, job_arguments:)
+      def ensure_batched_background_migration_is_finished(job_class_name:, table_name:, column_name:, job_arguments:, finalize: true)
         migration = Gitlab::Database::BackgroundMigration::BatchedMigration
           .for_configuration(job_class_name, table_name, column_name, job_arguments).first
@@ -954,9 +954,13 @@ module Gitlab
           job_arguments: job_arguments
         }
-        if migration.nil?
-          Gitlab::AppLogger.warn "Could not find batched background migration for the given configuration: #{configuration}"
-        elsif !migration.finished?
+        return Gitlab::AppLogger.warn "Could not find batched background migration for the given configuration: #{configuration}" if migration.nil?
+
+        return if migration.finished?
+
+        Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.finalize(job_class_name, table_name, column_name, job_arguments, connection: connection) if finalize
+
+        unless migration.reload.finished? # rubocop:disable Cop/ActiveRecordAssociationReload
           raise "Expected batched background migration for the given configuration to be marked as 'finished', " \
             "but it is '#{migration.status_name}':" \
             "\t#{configuration}" \

lib/gitlab/database/migrations/batched_background_migration_helpers.rb

@@ -122,6 +122,14 @@ module Gitlab
           migration.save!
           migration
         end
+
+        def finalize_batched_background_migration(job_class_name:, table_name:, column_name:, job_arguments:)
+          migration = Gitlab::Database::BackgroundMigration::BatchedMigration.find_for_configuration(job_class_name, table_name, column_name, job_arguments)
+
+          raise 'Could not find batched background migration' if migration.nil?
+
+          Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.finalize(job_class_name, table_name, column_name, job_arguments, connection: connection)
+        end
       end
     end
   end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2228,12 +2228,17 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   end
   describe '#ensure_batched_background_migration_is_finished' do
+    let(:job_class_name) { 'CopyColumnUsingBackgroundMigrationJob' }
+    let(:table) { :events }
+    let(:column_name) { :id }
+    let(:job_arguments) { [["id"], ["id_convert_to_bigint"], nil] }
+
     let(:configuration) do
       {
-        job_class_name: 'CopyColumnUsingBackgroundMigrationJob',
-        table_name: :events,
-        column_name: :id,
-        job_arguments: [["id"], ["id_convert_to_bigint"], nil]
+        job_class_name: job_class_name,
+        table_name: table,
+        column_name: column_name,
+        job_arguments: job_arguments
       }
     end
@@ -2242,6 +2247,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     it 'raises an error when migration exists and is not marked as finished' do
       create(:batched_background_migration, :active, configuration)
+      allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+        allow(runner).to receive(:finalize).with(job_class_name, table, column_name, job_arguments).and_return(false)
+      end
+
       expect { ensure_batched_background_migration_is_finished }
         .to raise_error "Expected batched background migration for the given configuration to be marked as 'finished', but it is 'active':" \
             "\t#{configuration}" \
@@ -2269,6 +2278,28 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       expect { ensure_batched_background_migration_is_finished }
         .not_to raise_error
     end
+
+    it 'finalizes the migration' do
+      migration = create(:batched_background_migration, :active, configuration)
+
+      allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+        expect(runner).to receive(:finalize).with(job_class_name, table, column_name, job_arguments).and_return(migration.finish!)
+      end
+
+      ensure_batched_background_migration_is_finished
+    end
+
+    context 'when the flag finalize is false' do
+      it 'does not finalize the migration' do
+        create(:batched_background_migration, :active, configuration)
+
+        allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+          expect(runner).not_to receive(:finalize).with(job_class_name, table, column_name, job_arguments)
+        end
+
+        expect { model.ensure_batched_background_migration_is_finished(**configuration.merge(finalize: false)) }.to raise_error(RuntimeError)
+      end
+    end
   end
   describe '#index_exists_by_name?' do

spec/lib/gitlab/database/migrations/batched_background_migration_helpers_spec.rb

@@ -163,4 +163,24 @@ RSpec.describe Gitlab::Database::Migrations::BatchedBackgroundMigrationHelpers d
       end
     end
   end
+
+  describe '#finalize_batched_background_migration' do
+    let!(:batched_migration) { create(:batched_background_migration, job_class_name: 'MyClass', table_name: :projects, column_name: :id, job_arguments: []) }
+
+    it 'finalizes the migration' do
+      allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+        expect(runner).to receive(:finalize).with('MyClass', :projects, :id, [])
+      end
+
+      migration.finalize_batched_background_migration(job_class_name: 'MyClass', table_name: :projects, column_name: :id, job_arguments: [])
+    end
+
+    context 'when the migration does not exist' do
+      it 'raises an exception' do
+        expect do
+          migration.finalize_batched_background_migration(job_class_name: 'MyJobClass', table_name: :projects, column_name: :id, job_arguments: [])
+        end.to raise_error(RuntimeError, 'Could not find batched background migration')
+      end
+    end
+  end
 end

spec/migrations/finalize_project_namespaces_backfill_spec.rb

@@ -9,9 +9,11 @@ RSpec.describe FinalizeProjectNamespacesBackfill, :migration do
   let_it_be(:migration) { described_class::MIGRATION }
   describe '#up' do
-    shared_examples 'raises migration not finished exception' do
-      it 'raises exception' do
-        expect { migrate! }.to raise_error(/Expected batched background migration for the given configuration to be marked as 'finished'/)
+    shared_examples 'finalizes the migration' do
+      it 'finalizes the migration' do
+        allow_next_instance_of(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner) do |runner|
+          expect(runner).to receive(:finalize).with('"ProjectNamespaces::BackfillProjectNamespaces"', :projects, :id, [nil, "up"])
+        end
       end
     end
@@ -61,7 +63,7 @@ RSpec.describe FinalizeProjectNamespacesBackfill, :migration do
             project_namespace_backfill.update!(status: status)
           end
-          it_behaves_like 'raises migration not finished exception'
+          it_behaves_like 'finalizes the migration'
         end
       end
     end