Do not override original AR5 batching interface

config/initializers/ar5_batching.rb

@@ -15,7 +15,7 @@ module ActiveRecord
       relation = relation.where(arel_table[primary_key].lteq(finish)) if finish
       batch_relation = relation
 
-      1.step do |index|
+      loop do
         if load
           records = batch_relation.records
           ids = records.map(&:id)
@@ -31,7 +31,7 @@ module ActiveRecord
         primary_key_offset = ids.last
         raise ArgumentError.new("Primary key not included in the custom select clause") unless primary_key_offset
 
-        yield yielded_relation, index
+        yield yielded_relation
 
         break if ids.length < of
         batch_relation = relation.where(arel_table[primary_key].gt(primary_key_offset))

db/post_migrate/20170628080858_migrate_stage_id_reference_in_background.rb

@@ -12,9 +12,12 @@ class MigrateStageIdReferenceInBackground < ActiveRecord::Migration
   end
 
   def up
-    Build.where(stage_id: nil).in_batches(of: BATCH_SIZE) do |relation, index|
-      schedule = index * 5.minutes
+    index = 1
+
+    Build.where(stage_id: nil).in_batches(of: BATCH_SIZE) do |relation|
       jobs = relation.pluck(:id).map { |id| [MIGRATION, [id]] }
+      schedule = index * 5.minutes
+      index += 1
 
       BackgroundMigrationWorker.perform_bulk_in(schedule, jobs)
     end