1 files changed, 12 insertions, 0 deletions
diff --git a/lib/gitlab/database/migrations/runner.rb b/lib/gitlab/database/migrations/runner.rb
index 3b6f52b43a8..4404b5bf961 100644
--- a/lib/gitlab/database/migrations/runner.rb
+++ b/lib/gitlab/database/migrations/runner.rb
@@ -21,6 +21,18 @@ module Gitlab
             TestBackgroundRunner.new(result_dir: BASE_RESULT_DIR.join('background_migrations'))
           end
 
+          def batched_background_migrations(for_database:)
+            runner = nil
+
+            # Only one loop iteration since we pass `only:` here
+            Gitlab::Database::EachDatabase.each_database_connection(only: for_database) do |connection|
+              runner = Gitlab::Database::Migrations::TestBatchedBackgroundRunner
+                         .new(result_dir: BASE_RESULT_DIR.join('background_migrations'), connection: connection)
+            end
+
+            runner
+          end
+
           def migration_context
             @migration_context ||= ApplicationRecord.connection.migration_context
           end