1 files changed, 7 insertions, 7 deletions
diff --git a/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb b/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
index 31a73bb3b27..4d5ac81382e 100644
--- a/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
+++ b/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
@@ -7,22 +7,22 @@ class MigrateBuildStageReferenceAgain < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
     stage_id = Arel.sql <<-SQL.strip_heredoc
       (SELECT id FROM ci_stages
          WHERE ci_stages.pipeline_id = ci_builds.commit_id
            AND ci_stages.name = ci_builds.stage)
     SQL
 
-    update_column_in_batches(:ci_builds, :stage_id, stage_id) do |table, query|
-      query.where(table[:stage_id].eq(nil))
+    disable_statement_timeout(transaction: false) do
+      update_column_in_batches(:ci_builds, :stage_id, stage_id) do |table, query|
+        query.where(table[:stage_id].eq(nil))
+      end
     end
   end
 
   def down
-    disable_statement_timeout
-
-    update_column_in_batches(:ci_builds, :stage_id, nil)
+    disable_statement_timeout(transaction: false) do
+      update_column_in_batches(:ci_builds, :stage_id, nil)
+    end
   end
 end