1 files changed, 55 insertions, 47 deletions

diff --git a/lib/gitlab/database/migrations/test_batched_background_runner.rb b/lib/gitlab/database/migrations/test_batched_background_runner.rb
index c27ae6a2c5d..46855ca1921 100644
--- a/lib/gitlab/database/migrations/test_batched_background_runner.rb
+++ b/lib/gitlab/database/migrations/test_batched_background_runner.rb
@@ -5,65 +5,73 @@ module Gitlab
     module Migrations
       class TestBatchedBackgroundRunner < BaseBackgroundRunner
         include Gitlab::Database::DynamicModelHelpers
-        attr_reader :connection
 
         def initialize(result_dir:, connection:)
-          super(result_dir: result_dir)
+          super(result_dir: result_dir, connection: connection)
           @connection = connection
         end
 
         def jobs_by_migration_name
-          Gitlab::Database::BackgroundMigration::BatchedMigration
-            .executable
-            .created_after(3.hours.ago) # Simple way to exclude migrations already running before migration testing
-            .to_h do |migration|
-            batching_strategy = migration.batch_class.new(connection: connection)
-
-            smallest_batch_start = migration.next_min_value
-
-            table_max_value = define_batchable_model(migration.table_name, connection: connection)
-                                .maximum(migration.column_name)
-
-            largest_batch_start = table_max_value - migration.batch_size
-
-            # variance is the portion of the batch range that we shrink between variance * 0 and variance * 1
-            # to pick actual batches to sample.
-            variance = largest_batch_start - smallest_batch_start
-
-            batch_starts = uniform_fractions
-                             .lazy # frac varies from 0 to 1, values in smallest_batch_start..largest_batch_start
-                             .map { |frac| (variance * frac).to_i + smallest_batch_start }
-
-            # Track previously run batches so that we stop sampling if a new batch would intersect an older one
-            completed_batches = []
-
-            jobs_to_sample = batch_starts
-                               # Stop sampling if a batch would intersect a previous batch
-                               .take_while { |start| completed_batches.none? { |batch| batch.cover?(start) } }
-                               .map do |batch_start|
-              next_bounds = batching_strategy.next_batch(
-                migration.table_name,
-                migration.column_name,
-                batch_min_value: batch_start,
-                batch_size: migration.batch_size,
-                job_arguments: migration.job_arguments
-              )
-
-              batch_min, batch_max = next_bounds
-
-              job = migration.create_batched_job!(batch_min, batch_max)
-
-              completed_batches << (batch_min..batch_max)
+          Gitlab::Database::SharedModel.using_connection(connection) do
+            Gitlab::Database::BackgroundMigration::BatchedMigration
+              .executable
+              .created_after(3.hours.ago) # Simple way to exclude migrations already running before migration testing
+              .to_h do |migration|
+              batching_strategy = migration.batch_class.new(connection: connection)
+
+              smallest_batch_start = migration.next_min_value
+
+              table_max_value = define_batchable_model(migration.table_name, connection: connection)
+                                  .maximum(migration.column_name)
+
+              largest_batch_start = table_max_value - migration.batch_size
+
+              # variance is the portion of the batch range that we shrink between variance * 0 and variance * 1
+              # to pick actual batches to sample.
+              variance = largest_batch_start - smallest_batch_start
+
+              batch_starts = uniform_fractions
+                               .lazy # frac varies from 0 to 1, values in smallest_batch_start..largest_batch_start
+                               .map { |frac| (variance * frac).to_i + smallest_batch_start }
+
+              # Track previously run batches so that we stop sampling if a new batch would intersect an older one
+              completed_batches = []
+
+              jobs_to_sample = batch_starts
+                                 # Stop sampling if a batch would intersect a previous batch
+                                 .take_while { |start| completed_batches.none? { |batch| batch.cover?(start) } }
+                                 .map do |batch_start|
+                # The current block is lazily evaluated as part of the jobs_to_sample enumerable
+                # so it executes after the enclosing using_connection block has already executed
+                # Therefore we need to re-associate with the explicit connection again
+                Gitlab::Database::SharedModel.using_connection(connection) do
+                  next_bounds = batching_strategy.next_batch(
+                    migration.table_name,
+                    migration.column_name,
+                    batch_min_value: batch_start,
+                    batch_size: migration.batch_size,
+                    job_arguments: migration.job_arguments
+                  )
+
+                  batch_min, batch_max = next_bounds
+
+                  job = migration.create_batched_job!(batch_min, batch_max)
+
+                  completed_batches << (batch_min..batch_max)
+
+                  job
+                end
+              end
 
-              job
+              [migration.job_class_name, jobs_to_sample]
             end
-
-            [migration.job_class_name, jobs_to_sample]
           end
         end
 
         def run_job(job)
-          Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection).perform(job)
+          Gitlab::Database::SharedModel.using_connection(connection) do
+            Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection).perform(job)
+          end
         end
 
         def uniform_fractions