Use a helper to schedule pipeline stages migration

2 files changed, 6 insertions, 7 deletions

diff --git a/db/post_migrate/20180212101928_schedule_build_stage_migration.rb b/db/post_migrate/20180212101928_schedule_build_stage_migration.rb
index ad65a6fb16e..df15b2cd9d4 100644
--- a/db/post_migrate/20180212101928_schedule_build_stage_migration.rb
+++ b/db/post_migrate/20180212101928_schedule_build_stage_migration.rb
@@ -3,7 +3,7 @@ class ScheduleBuildStageMigration < ActiveRecord::Migration
 
   DOWNTIME = false
   MIGRATION = 'MigrateBuildStage'.freeze
-  BATCH_SIZE = 800
+  BATCH_SIZE = 500
 
   disable_ddl_transaction!
 
@@ -15,10 +15,11 @@ class ScheduleBuildStageMigration < ActiveRecord::Migration
   def up
     disable_statement_timeout
 
-    Build.where('stage_id IS NULL').each_batch(of: BATCH_SIZE) do |builds, index|
-      builds.pluck('MIN(id)', 'MAX(id)').first.tap do |range|
-        BackgroundMigrationWorker.perform_in(index * 5.minutes, MIGRATION, range)
-      end
+    Build.where('stage_id IS NULL').tap do |relation|
+      queue_background_migration_jobs_by_range_at_intervals(relation,
+                                                            MIGRATION,
+                                                            5.minutes,
+                                                            batch_size: BATCH_SIZE)
     end
   end
 
diff --git a/lib/gitlab/background_migration/migrate_build_stage.rb b/lib/gitlab/background_migration/migrate_build_stage.rb
index adca16751af..8fe4f1a2289 100644
--- a/lib/gitlab/background_migration/migrate_build_stage.rb
+++ b/lib/gitlab/background_migration/migrate_build_stage.rb
@@ -35,8 +35,6 @@ module Gitlab
       end
 
       def perform(start_id, stop_id)
-        # TODO, statement timeout?
-
         stages = Migratable::Build.where('stage_id IS NULL')
           .where('id BETWEEN ? AND ?', start_id, stop_id)
           .map { |build| build.ensure_stage! }