1 files changed, 9 insertions, 6 deletions
diff --git a/lib/gitlab/database/dynamic_model_helpers.rb b/lib/gitlab/database/dynamic_model_helpers.rb
index 220062f1bc6..ad7dea8f0d9 100644
--- a/lib/gitlab/database/dynamic_model_helpers.rb
+++ b/lib/gitlab/database/dynamic_model_helpers.rb
@@ -5,16 +5,19 @@ module Gitlab
     module DynamicModelHelpers
       BATCH_SIZE = 1_000
 
-      def define_batchable_model(table_name)
-        Class.new(ActiveRecord::Base) do
+      def define_batchable_model(table_name, connection:)
+        klass = Class.new(ActiveRecord::Base) do
           include EachBatch
 
           self.table_name = table_name
           self.inheritance_column = :_type_disabled
         end
+
+        klass.connection = connection
+        klass
       end
 
-      def each_batch(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
+      def each_batch(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
         if transaction_open?
           raise <<~MSG.squish
             each_batch should not run inside a transaction, you can disable
@@ -23,12 +26,12 @@ module Gitlab
           MSG
         end
 
-        scope.call(define_batchable_model(table_name))
+        scope.call(define_batchable_model(table_name, connection: connection))
           .each_batch(of: of) { |batch| yield batch }
       end
 
-      def each_batch_range(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
-        each_batch(table_name, scope: scope, of: of) do |batch|
+      def each_batch_range(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
+        each_batch(table_name, connection: connection, scope: scope, of: of) do |batch|
           yield batch.pluck('MIN(id), MAX(id)').first
         end
       end