Add new migration helpers for data migration in batches

db/post_migrate/20210525075724_clean_up_pending_builds_table.rb

 # frozen_string_literal: true
 
 class CleanUpPendingBuildsTable < ActiveRecord::Migration[6.0]
+  include ::Gitlab::Database::MigrationHelpers
+
   BATCH_SIZE = 1000
 
   disable_ddl_transaction!
@@ -8,7 +10,7 @@ class CleanUpPendingBuildsTable < ActiveRecord::Migration[6.0]
   def up
     return unless Gitlab.dev_or_test_env? || Gitlab.com?
 
-    each_batch('ci_pending_builds', of: BATCH_SIZE) do |min, max|
+    each_batch_range('ci_pending_builds', of: BATCH_SIZE) do |min, max|
       execute <<~SQL
         DELETE FROM ci_pending_builds
           USING ci_builds
@@ -23,19 +25,4 @@ class CleanUpPendingBuildsTable < ActiveRecord::Migration[6.0]
   def down
     # noop
   end
-
-  private
-
-  def each_batch(table_name, scope: ->(table) { table.all }, of: 1000)
-    table = Class.new(ActiveRecord::Base) do
-      include EachBatch
-
-      self.table_name = table_name
-      self.inheritance_column = :_type_disabled
-    end
-
-    scope.call(table).each_batch(of: of) do |batch|
-      yield batch.pluck('MIN(id), MAX(id)').first
-    end
-  end
 end

lib/gitlab/database/migration_helpers.rb

@@ -1591,6 +1591,31 @@ into similar problems in the future (e.g. when new tables are created).
         raise
       end
 
+      def each_batch(table_name, scope: ->(table) { table.all }, of: 1000)
+        if transaction_open?
+          raise <<~MSG.squish
+            each_batch should not run inside a transaction, you can disable
+            transactions by calling disable_ddl_transaction! in the body of
+            your migration class
+          MSG
+        end
+
+        table = Class.new(ActiveRecord::Base) do
+          include EachBatch
+
+          self.table_name = table_name
+          self.inheritance_column = :_type_disabled
+        end
+
+        scope.call(table).each_batch(of: of) { |batch| yield batch }
+      end
+
+      def each_batch_range(table_name, scope: ->(table) { table.all }, of: 1000)
+        each_batch(table_name, scope: scope, of: of) do |batch|
+          yield batch.pluck('MIN(id), MAX(id)').first
+        end
+      end
+
       private
 
       def validate_check_constraint_name!(constraint_name)