Schedule stage_id background migration in ranges

2 files changed, 9 insertions, 4 deletions

diff --git a/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb b/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
index d1d0a968588..0705b26056d 100644
--- a/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
+++ b/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
@@ -1,15 +1,20 @@
 module Gitlab
   module BackgroundMigration
     class MigrateBuildStageIdReference
-      def perform(id)
+      def perform(start_id, stop_id)
+        scope = if stop_id.nonzero?
+                  "ci_builds.id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}"
+                else
+                  "ci_builds.id >= #{start_id.to_i}"
+                end
+
         sql = <<-SQL.strip_heredoc
           UPDATE "ci_builds"
             SET "stage_id" =
               (SELECT id FROM ci_stages
                 WHERE ci_stages.pipeline_id = ci_builds.commit_id
                 AND ci_stages.name = ci_builds.stage)
-          WHERE "ci_builds"."id" = #{id.to_i}
-            AND "ci_builds"."stage_id" IS NULL
+          WHERE #{scope} AND "ci_builds"."stage_id" IS NULL
         SQL
 
         ActiveRecord::Base.connection.execute(sql)
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 5acc96331c0..b6883704da0 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -296,7 +296,7 @@ module Gitlab
           MSG
         end
 
-       # raise 'This method requires a block!' unless block_given?
+        raise ArgumentError, 'This method requires a block!' unless block_given?
 
         table_arel = Arel::Table.new(table)