From fb89ba24825853ca29b804a4a08f7c83210f09db Mon Sep 17 00:00:00 2001
From: Grzegorz Bizon <grzesiek.bizon@gmail.com>
Date: Fri, 7 Jul 2017 09:35:28 +0200
Subject: [PATCH] Schedule stage_id background migration in ranges

---
 ..._migrate_stage_id_reference_in_background.rb | 17 ++++++++---------
 .../migrate_build_stage_id_reference.rb         | 11 ++++++++---
 lib/gitlab/database/migration_helpers.rb        |  2 +-
 ...ate_stage_id_reference_in_background_spec.rb | 11 +++++------
 4 files changed, 22 insertions(+), 19 deletions(-)

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 ebec4cb6bb7..0d108d18501 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,19 +7,18 @@ class MigrateStageIdReferenceInBackground < ActiveRecord::Migration
 
   disable_ddl_transaction!
 
-  class Build < ActiveRecord::Base
-    self.table_name = 'ci_builds'
-  end
-
+  ##
+  # It will take around 3 days to process 20M ci_builds.
+  #
   def up
-    index = 1
+    opts = { scope: ->(table, query) { query.where(table[:stage_id].eq(nil)) },
+             of: BATCH_SIZE }
 
-    Build.where(stage_id: nil).in_batches(of: BATCH_SIZE) do |relation|
-      jobs = relation.pluck(:id).map { |id| [MIGRATION, [id]] }
+    walk_table_in_batches(:ci_builds, **opts) do |index, start_id, stop_id|
       schedule = index * 2.minutes
-      index += 1
 
-      BackgroundMigrationWorker.perform_bulk_in(schedule, jobs)
+      BackgroundMigrationWorker
+        .perform_in(schedule, MIGRATION, [start_id, stop_id])
     end
   end
 
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)
 
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 a32a7fceb68..f3dde8b59c0 100644
--- a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
+++ b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
@@ -21,7 +21,7 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
   let(:projects) { table(:projects) }
 
   before do
-    stub_const('MigrateStageIdReferenceInBackground::BATCH_SIZE', 2)
+    stub_const("#{described_class.name}::BATCH_SIZE", 2)
 
     projects.create!(id: 123, name: 'gitlab1', path: 'gitlab1')
     projects.create!(id: 345, name: 'gitlab2', path: 'gitlab2')
@@ -47,11 +47,10 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
       Timecop.freeze do
         migrate!
 
-        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 1)
-        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 2)
-        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 3)
-        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 4)
-        expect(BackgroundMigrationWorker.jobs.size).to eq 5
+        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(BackgroundMigrationWorker.jobs.size).to eq 3
       end
     end
   end