Move stages status migration to the background worker

db/post_migrate/20170630111158_migrate_stages_statuses.rb

@@ -5,73 +5,22 @@ class MigrateStagesStatuses < ActiveRecord::Migration
 
   disable_ddl_transaction!
 
-  STATUSES = { created: 0, pending: 1, running: 2, success: 3,
-               failed: 4, canceled: 5, skipped: 6, manual: 7 }
+  BATCH_SIZE = 10000
+  MIGRATION = 'MigrateStageStatus'.freeze
 
-  class Build < ActiveRecord::Base
-    self.table_name = 'ci_builds'
-
-    scope :latest, -> { where(retried: [false, nil]) }
-    scope :created, -> { where(status: 'created') }
-    scope :running, -> { where(status: 'running') }
-    scope :pending, -> { where(status: 'pending') }
-    scope :success, -> { where(status: 'success') }
-    scope :failed, -> { where(status: 'failed')  }
-    scope :canceled, -> { where(status: 'canceled')  }
-    scope :skipped, -> { where(status: 'skipped')  }
-    scope :manual, -> { where(status: 'manual')  }
-
-    scope :failed_but_allowed, -> do
-      where(allow_failure: true, status: [:failed, :canceled])
-    end
-
-    scope :exclude_ignored, -> do
-      where("allow_failure = ? OR status IN (?)",
-        false, %w[created pending running success skipped])
-    end
-
-    def self.status_sql
-      scope_relevant = latest.exclude_ignored
-      scope_warnings = latest.failed_but_allowed
-
-      builds = scope_relevant.select('count(*)').to_sql
-      created = scope_relevant.created.select('count(*)').to_sql
-      success = scope_relevant.success.select('count(*)').to_sql
-      manual = scope_relevant.manual.select('count(*)').to_sql
-      pending = scope_relevant.pending.select('count(*)').to_sql
-      running = scope_relevant.running.select('count(*)').to_sql
-      skipped = scope_relevant.skipped.select('count(*)').to_sql
-      canceled = scope_relevant.canceled.select('count(*)').to_sql
-      warnings = scope_warnings.select('count(*) > 0').to_sql
-
-      <<-SQL.strip_heredoc
-        (CASE
-          WHEN (#{builds}) = (#{skipped}) AND (#{warnings}) THEN #{STATUSES[:success]}
-          WHEN (#{builds}) = (#{skipped}) THEN #{STATUSES[:skipped]}
-          WHEN (#{builds}) = (#{success}) THEN #{STATUSES[:success]}
-          WHEN (#{builds}) = (#{created}) THEN #{STATUSES[:created]}
-          WHEN (#{builds}) = (#{success}) + (#{skipped}) THEN #{STATUSES[:success]}
-          WHEN (#{builds}) = (#{success}) + (#{skipped}) + (#{canceled}) THEN #{STATUSES[:canceled]}
-          WHEN (#{builds}) = (#{created}) + (#{skipped}) + (#{pending}) THEN #{STATUSES[:pending]}
-          WHEN (#{running}) + (#{pending}) > 0 THEN #{STATUSES[:running]}
-          WHEN (#{manual}) > 0 THEN #{STATUSES[:manual]}
-          WHEN (#{created}) > 0 THEN #{STATUSES[:running]}
-          ELSE #{STATUSES[:failed]}
-        END)
-      SQL
-    end
+  class Stage < ActiveRecord::Base
+    self.table_name = 'ci_stages'
   end
 
   def up
-    disable_statement_timeout
+    index = 1
 
-    status_sql = Build
-      .where('ci_builds.commit_id = ci_stages.pipeline_id')
-      .where('ci_builds.stage = ci_stages.name')
-      .status_sql
+    Stage.where(status: nil).in_batches(of: BATCH_SIZE) do |relation|
+      jobs = relation.pluck(:id).map { |id| [MIGRATION, [id]] }
+      schedule = index * 5.minutes
+      index += 1
 
-    update_column_in_batches(:ci_stages, :status, Arel.sql("(#{status_sql})")) do |table, query|
-      query.where(table[:status].eq(nil))
+      BackgroundMigrationWorker.perform_bulk_in(schedule, jobs)
     end
   end
 

lib/gitlab/background_migration/migrate_stage_status.rb

+module Gitlab
+  module BackgroundMigration
+    class MigrateStageStatus
+      STATUSES = { created: 0, pending: 1, running: 2, success: 3,
+                   failed: 4, canceled: 5, skipped: 6, manual: 7 }
+
+      class Build < ActiveRecord::Base
+        self.table_name = 'ci_builds'
+
+        scope :latest, -> { where(retried: [false, nil]) }
+        scope :created, -> { where(status: 'created') }
+        scope :running, -> { where(status: 'running') }
+        scope :pending, -> { where(status: 'pending') }
+        scope :success, -> { where(status: 'success') }
+        scope :failed, -> { where(status: 'failed')  }
+        scope :canceled, -> { where(status: 'canceled')  }
+        scope :skipped, -> { where(status: 'skipped')  }
+        scope :manual, -> { where(status: 'manual')  }
+
+        scope :failed_but_allowed, -> do
+          where(allow_failure: true, status: [:failed, :canceled])
+        end
+
+        scope :exclude_ignored, -> do
+          where("allow_failure = ? OR status IN (?)",
+            false, %w[created pending running success skipped])
+        end
+
+        def self.status_sql
+          scope_relevant = latest.exclude_ignored
+          scope_warnings = latest.failed_but_allowed
+
+          builds = scope_relevant.select('count(*)').to_sql
+          created = scope_relevant.created.select('count(*)').to_sql
+          success = scope_relevant.success.select('count(*)').to_sql
+          manual = scope_relevant.manual.select('count(*)').to_sql
+          pending = scope_relevant.pending.select('count(*)').to_sql
+          running = scope_relevant.running.select('count(*)').to_sql
+          skipped = scope_relevant.skipped.select('count(*)').to_sql
+          canceled = scope_relevant.canceled.select('count(*)').to_sql
+          warnings = scope_warnings.select('count(*) > 0').to_sql
+
+          <<-SQL.strip_heredoc
+            (CASE
+              WHEN (#{builds}) = (#{skipped}) AND (#{warnings}) THEN #{STATUSES[:success]}
+              WHEN (#{builds}) = (#{skipped}) THEN #{STATUSES[:skipped]}
+              WHEN (#{builds}) = (#{success}) THEN #{STATUSES[:success]}
+              WHEN (#{builds}) = (#{created}) THEN #{STATUSES[:created]}
+              WHEN (#{builds}) = (#{success}) + (#{skipped}) THEN #{STATUSES[:success]}
+              WHEN (#{builds}) = (#{success}) + (#{skipped}) + (#{canceled}) THEN #{STATUSES[:canceled]}
+              WHEN (#{builds}) = (#{created}) + (#{skipped}) + (#{pending}) THEN #{STATUSES[:pending]}
+              WHEN (#{running}) + (#{pending}) > 0 THEN #{STATUSES[:running]}
+              WHEN (#{manual}) > 0 THEN #{STATUSES[:manual]}
+              WHEN (#{created}) > 0 THEN #{STATUSES[:running]}
+              ELSE #{STATUSES[:failed]}
+            END)
+          SQL
+        end
+      end
+
+      def perform(id)
+        status_sql = Build
+          .where('ci_builds.commit_id = ci_stages.pipeline_id')
+          .where('ci_builds.stage = ci_stages.name')
+          .status_sql
+
+        sql = <<-SQL
+          UPDATE ci_stages SET status = (#{status_sql})
+            WHERE ci_stages.id = #{id.to_i}
+        SQL
+
+        ActiveRecord::Base.connection.execute(sql)
+      end
+    end
+  end
+end

spec/migrations/migrate_stage_id_reference_in_background_spec.rb

@@ -2,19 +2,6 @@ require 'spec_helper'
 require Rails.root.join('db', 'post_migrate', '20170628080858_migrate_stage_id_reference_in_background')
 
 describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
-  matcher :be_scheduled_migration do |delay, *expected|
-    match do |migration|
-      BackgroundMigrationWorker.jobs.any? do |job|
-        job['args'] == [migration, expected] &&
-          job['at'].to_i == (delay.to_i + Time.now.to_i)
-      end
-    end
-
-    failure_message do |migration|
-      "Migration `#{migration}` with args `#{expected.inspect}` not scheduled!"
-    end
-  end
-
   let(:jobs) { table(:ci_builds) }
   let(:stages) { table(:ci_stages) }
   let(:pipelines) { table(:ci_pipelines) }

spec/migrations/migrate_stages_statuses_spec.rb

@@ -11,6 +11,8 @@ describe MigrateStagesStatuses, :migration do
                failed: 4, canceled: 5, skipped: 6, manual: 7 }
 
   before do
+    stub_const("#{described_class.name}::BATCH_SIZE", 2)
+
     projects.create!(id: 1, name: 'gitlab1', path: 'gitlab1')
     projects.create!(id: 2, name: 'gitlab2', path: 'gitlab2')
 
@@ -31,15 +33,31 @@ describe MigrateStagesStatuses, :migration do
   end
 
   it 'correctly migrates stages statuses' do
-    expect(stages.where(status: nil).count).to eq 3
+    Sidekiq::Testing.inline! do
+      expect(stages.where(status: nil).count).to eq 3
 
-    migrate!
+      migrate!
 
-    expect(stages.where(status: nil)).to be_empty
-    expect(stages.all.order('id ASC').pluck(:status))
-      .to eq [STATUSES[:running], STATUSES[:failed], STATUSES[:success]]
+      expect(stages.where(status: nil)).to be_empty
+      expect(stages.all.order('id ASC').pluck(:status))
+        .to eq [STATUSES[:running], STATUSES[:failed], STATUSES[:success]]
+    end
   end
 
+  it 'correctly schedules background migrations' do
+    Sidekiq::Testing.fake! do
+      Timecop.freeze do
+        migrate!
+
+        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 1)
+        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 2)
+        expect(described_class::MIGRATION).to be_scheduled_migration(10.minutes, 3)
+        expect(BackgroundMigrationWorker.jobs.size).to eq 3
+      end
+    end
+  end
+
+
   def create_job(project:, pipeline:, stage:, status:, **opts)
     stages  = { test: 1, build: 2, deploy: 3}
 

spec/support/background_migrations_matchers.rb

+RSpec::Matchers.define :be_scheduled_migration do |delay, *expected|
+  match do |migration|
+    BackgroundMigrationWorker.jobs.any? do |job|
+      job['args'] == [migration, expected] &&
+        job['at'].to_i == (delay.to_i + Time.now.to_i)
+    end
+  end
+
+  failure_message do |migration|
+    "Migration `#{migration}` with args `#{expected.inspect}` " \
+      'not scheduled in expected time!'
+  end
+end