Recover from unique constraint violation in stages migration

3 files changed, 21 insertions, 10 deletions

diff --git a/db/post_migrate/20180212101928_schedule_build_stage_migration.rb b/db/post_migrate/20180212101928_schedule_build_stage_migration.rb
index f03a3f8d355..ad65a6fb16e 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 = 500
+  BATCH_SIZE = 800
 
   disable_ddl_transaction!
 
diff --git a/lib/gitlab/background_migration/migrate_build_stage.rb b/lib/gitlab/background_migration/migrate_build_stage.rb
index 3e297a90b19..adca16751af 100644
--- a/lib/gitlab/background_migration/migrate_build_stage.rb
+++ b/lib/gitlab/background_migration/migrate_build_stage.rb
@@ -13,19 +13,20 @@ module Gitlab
         class Build < ActiveRecord::Base
           self.table_name = 'ci_builds'
 
-          def ensure_stage!
-            find || create!
+          def ensure_stage!(attempts: 2)
+            find_stage || create_stage!
           rescue ActiveRecord::RecordNotUnique
-            # TODO
+            retry if (attempts -= 1) > 0
+            raise
           end
 
-          def find
-            Stage.find_by(name: self.stage,
+          def find_stage
+            Stage.find_by(name: self.stage || 'test',
                           pipeline_id: self.commit_id,
                           project_id: self.project_id)
           end
 
-          def create!
+          def create_stage!
             Stage.create!(name: self.stage || 'test',
                           pipeline_id: self.commit_id,
                           project_id: self.project_id)
@@ -34,11 +35,10 @@ module Gitlab
       end
 
       def perform(start_id, stop_id)
-        # TODO, should we disable_statement_timeout?
-        # TODO, use plain SQL query?
+        # TODO, statement timeout?
 
         stages = Migratable::Build.where('stage_id IS NULL')
-          .where("id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}")
+          .where('id BETWEEN ? AND ?', start_id, stop_id)
           .map { |build| build.ensure_stage! }
           .compact.map(&:id)
 
diff --git a/spec/lib/gitlab/background_migration/migrate_build_stage_spec.rb b/spec/lib/gitlab/background_migration/migrate_build_stage_spec.rb
index 27cb63848de..e112e9e9e3d 100644
--- a/spec/lib/gitlab/background_migration/migrate_build_stage_spec.rb
+++ b/spec/lib/gitlab/background_migration/migrate_build_stage_spec.rb
@@ -40,4 +40,15 @@ describe Gitlab::BackgroundMigration::MigrateBuildStage, :migration, schema: 201
                                                       STATUSES[:failed],
                                                       STATUSES[:pending]]
   end
+
+  it 'recovers from unique constraint violation only twice' do
+    allow(described_class::Migratable::Stage)
+      .to receive(:find_by).and_return(nil)
+
+    expect(described_class::Migratable::Stage)
+      .to receive(:find_by).exactly(3).times
+
+    expect { described_class.new.perform(1, 6) }
+      .to raise_error ActiveRecord::RecordNotUnique
+  end
 end