1 files changed, 19 insertions, 17 deletions
diff --git a/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb b/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
index 73c23dffca0..5418f442e79 100644
--- a/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
+++ b/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
@@ -12,32 +12,34 @@ class CleanupStagesPositionMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
+    disable_statement_timeout do
+      Gitlab::BackgroundMigration.steal('MigrateStageIndex')
 
-    Gitlab::BackgroundMigration.steal('MigrateStageIndex')
-
-    unless index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
-      add_concurrent_index(:ci_stages, :id, where: 'position IS NULL', name: TMP_INDEX_NAME)
-    end
+      unless index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
+        add_concurrent_index(:ci_stages, :id, where: 'position IS NULL', name: TMP_INDEX_NAME)
+      end
 
-    migratable = <<~SQL
-      position IS NULL AND EXISTS (
-        SELECT 1 FROM ci_builds WHERE stage_id = ci_stages.id AND stage_idx IS NOT NULL
-      )
-    SQL
+      migratable = <<~SQL
+        position IS NULL AND EXISTS (
+          SELECT 1 FROM ci_builds WHERE stage_id = ci_stages.id AND stage_idx IS NOT NULL
+        )
+      SQL
 
-    Stages.where(migratable).each_batch(of: 1000) do |batch|
-      batch.pluck(:id).each do |stage|
-        Gitlab::BackgroundMigration::MigrateStageIndex.new.perform(stage, stage)
+      Stages.where(migratable).each_batch(of: 1000) do |batch|
+        batch.pluck(:id).each do |stage|
+          Gitlab::BackgroundMigration::MigrateStageIndex.new.perform(stage, stage)
+        end
       end
-    end
 
-    remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
+      remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
+    end
   end
 
   def down
     if index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
-      remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
+      disable_statement_timeout do
+        remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
+      end
     end
   end
 end