Use new `each_batch` helper instead of custom one

In stage_id backgrond migration.

3 files changed, 13 insertions, 11 deletions

diff --git a/app/models/concerns/each_batch.rb b/app/models/concerns/each_batch.rb
index 6610e7967d1..53979c4c683 100644
--- a/app/models/concerns/each_batch.rb
+++ b/app/models/concerns/each_batch.rb
@@ -35,7 +35,7 @@ module EachBatch
       start_id = start[primary_key]
       arel_table = self.arel_table
 
-      loop do
+      1.step do |index|
         stop = except(:select)
           .select(primary_key)
           .where(arel_table[primary_key].gteq(start_id))
@@ -54,7 +54,7 @@ module EachBatch
 
         # Any ORDER BYs are useless for this relation and can lead to less
         # efficient UPDATE queries, hence we get rid of it.
-        yield relation.except(:order)
+        yield relation.except(:order), index
 
         break unless stop
       end
diff --git a/db/post_migrate/20170628080858_migrate_stage_id_reference_in_background.rb b/db/post_migrate/20170628080858_migrate_stage_id_reference_in_background.rb
index 0d108d18501..107a5661329 100644
--- a/db/post_migrate/20170628080858_migrate_stage_id_reference_in_background.rb
+++ b/db/post_migrate/20170628080858_migrate_stage_id_reference_in_background.rb
@@ -7,18 +7,20 @@ class MigrateStageIdReferenceInBackground < ActiveRecord::Migration
 
   disable_ddl_transaction!
 
+  class Build < ActiveRecord::Base
+    self.table_name = 'ci_builds'
+    include ::EachBatch
+  end
+
   ##
   # It will take around 3 days to process 20M ci_builds.
   #
   def up
-    opts = { scope: ->(table, query) { query.where(table[:stage_id].eq(nil)) },
-             of: BATCH_SIZE }
-
-    walk_table_in_batches(:ci_builds, **opts) do |index, start_id, stop_id|
+    Build.all.each_batch(of: BATCH_SIZE) do |relation, index|
+      range = relation.pluck('MIN(id)', 'MAX(id)').first
       schedule = index * 2.minutes
 
-      BackgroundMigrationWorker
-        .perform_in(schedule, MIGRATION, [start_id, stop_id])
+      BackgroundMigrationWorker.perform_in(schedule, MIGRATION, range)
     end
   end
 
diff --git a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
index f3dde8b59c0..90ad5c39089 100644
--- a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
+++ b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
@@ -47,9 +47,9 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
       Timecop.freeze do
         migrate!
 
-        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 1, 3)
-        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 3, 5)
-        expect(described_class::MIGRATION).to be_scheduled_migration(6.minutes, 5, 0)
+        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 1, 2)
+        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 3, 4)
+        expect(described_class::MIGRATION).to be_scheduled_migration(6.minutes, 5, 6)
         expect(BackgroundMigrationWorker.jobs.size).to eq 3
       end
     end