1 files changed, 39 insertions, 37 deletions
diff --git a/db/migrate/20180420010616_cleanup_build_stage_migration.rb b/db/migrate/20180420010616_cleanup_build_stage_migration.rb
index 24777294101..5e9fe756efd 100644
--- a/db/migrate/20180420010616_cleanup_build_stage_migration.rb
+++ b/db/migrate/20180420010616_cleanup_build_stage_migration.rb
@@ -14,48 +14,50 @@ class CleanupBuildStageMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
-    ##
-    # We steal from the background migrations queue to catch up with the
-    # scheduled migrations set.
-    #
-    Gitlab::BackgroundMigration.steal('MigrateBuildStage')
-
-    ##
-    # We add temporary index, to make iteration over batches more performant.
-    # Conditional here is to avoid the need of doing that in a separate
-    # migration file to make this operation idempotent.
-    #
-    unless index_exists_by_name?(:ci_builds, TMP_INDEX)
-      add_concurrent_index(:ci_builds, :id, where: 'stage_id IS NULL', name: TMP_INDEX)
-    end
-
-    ##
-    # We check if there are remaining rows that should be migrated (for example
-    # if Sidekiq / Redis fails / is restarted, what could result in not all
-    # background migrations being executed correctly.
-    #
-    # We migrate remaining rows synchronously in a blocking way, to make sure
-    # that when this migration is done we are confident that all rows are
-    # already migrated.
-    #
-    Build.where('stage_id IS NULL').each_batch(of: 50) do |batch|
-      range = batch.pluck('MIN(id)', 'MAX(id)').first
-
-      Gitlab::BackgroundMigration::MigrateBuildStage.new.perform(*range)
+    disable_statement_timeout do
+      ##
+      # We steal from the background migrations queue to catch up with the
+      # scheduled migrations set.
+      #
+      Gitlab::BackgroundMigration.steal('MigrateBuildStage')
+
+      ##
+      # We add temporary index, to make iteration over batches more performant.
+      # Conditional here is to avoid the need of doing that in a separate
+      # migration file to make this operation idempotent.
+      #
+      unless index_exists_by_name?(:ci_builds, TMP_INDEX)
+        add_concurrent_index(:ci_builds, :id, where: 'stage_id IS NULL', name: TMP_INDEX)
+      end
+
+      ##
+      # We check if there are remaining rows that should be migrated (for example
+      # if Sidekiq / Redis fails / is restarted, what could result in not all
+      # background migrations being executed correctly.
+      #
+      # We migrate remaining rows synchronously in a blocking way, to make sure
+      # that when this migration is done we are confident that all rows are
+      # already migrated.
+      #
+      Build.where('stage_id IS NULL').each_batch(of: 50) do |batch|
+        range = batch.pluck('MIN(id)', 'MAX(id)').first
+
+        Gitlab::BackgroundMigration::MigrateBuildStage.new.perform(*range)
+      end
+
+      ##
+      # We remove temporary index, because it is not required during standard
+      # operations and runtime.
+      #
+      remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
     end
-
-    ##
-    # We remove temporary index, because it is not required during standard
-    # operations and runtime.
-    #
-    remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
   end
 
   def down
     if index_exists_by_name?(:ci_builds, TMP_INDEX)
-      remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
+      disable_statement_timeout do
+        remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
+      end
     end
   end
 end