Merge branch '48967-disable-statement-timeout' into 'master'

`disable_statement_timeout` will no longer leak to other migrations

Closes #48967

See merge request gitlab-org/gitlab-ce!20503

26 files changed, 400 insertions, 275 deletions

diff --git a/changelogs/unreleased/48967-disable-statement-timeout.yml b/changelogs/unreleased/48967-disable-statement-timeout.yml
new file mode 100644
index 00000000000..2da800ed41e
--- /dev/null
+++ b/changelogs/unreleased/48967-disable-statement-timeout.yml
@@ -0,0 +1,5 @@
+---
+title: disable_statement_timeout no longer leak to other migrations
+merge_request: 20503
+author:
+type: fixed
diff --git a/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb b/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
index 1199073ed3a..12352d98a62 100644
--- a/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
+++ b/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
@@ -106,11 +106,11 @@ class ProjectForeignKeysWithCascadingDeletes < ActiveRecord::Migration
             # Disables statement timeouts for the current connection. This is
             # necessary as removing of orphaned data might otherwise exceed the
             # statement timeout.
-            disable_statement_timeout
+            disable_statement_timeout do
+              remove_orphans(*queue.pop) until queue.empty?
 
-            remove_orphans(*queue.pop) until queue.empty?
-
-            steal_from_queues(queues - [queue])
+              steal_from_queues(queues - [queue])
+            end
           end
         end
       end
diff --git a/db/migrate/20170724214302_add_lower_path_index_to_redirect_routes.rb b/db/migrate/20170724214302_add_lower_path_index_to_redirect_routes.rb
index db60c2087b9..a770ff63b4e 100644
--- a/db/migrate/20170724214302_add_lower_path_index_to_redirect_routes.rb
+++ b/db/migrate/20170724214302_add_lower_path_index_to_redirect_routes.rb
@@ -25,8 +25,9 @@ class AddLowerPathIndexToRedirectRoutes < ActiveRecord::Migration
     # trivial to write a query that checks for an index. BUT there is a
     # convenient `IF EXISTS` parameter for `DROP INDEX`.
     if supports_drop_index_concurrently?
-      disable_statement_timeout
-      execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
+      disable_statement_timeout do
+        execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
+      end
     else
       execute "DROP INDEX IF EXISTS #{INDEX_NAME};"
     end
diff --git a/db/migrate/20171220191323_add_index_on_namespaces_lower_name.rb b/db/migrate/20171220191323_add_index_on_namespaces_lower_name.rb
index d1a039ed551..130b24fe6f0 100644
--- a/db/migrate/20171220191323_add_index_on_namespaces_lower_name.rb
+++ b/db/migrate/20171220191323_add_index_on_namespaces_lower_name.rb
@@ -8,25 +8,25 @@ class AddIndexOnNamespacesLowerName < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
-    if Gitlab::Database.version.to_f >= 9.5
-      # Allow us to hot-patch the index manually ahead of the migration
-      execute "CREATE INDEX CONCURRENTLY IF NOT EXISTS #{INDEX_NAME} ON namespaces (lower(name));"
-    else
-      execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON namespaces (lower(name));"
+    disable_statement_timeout do
+      if Gitlab::Database.version.to_f >= 9.5
+        # Allow us to hot-patch the index manually ahead of the migration
+        execute "CREATE INDEX CONCURRENTLY IF NOT EXISTS #{INDEX_NAME} ON namespaces (lower(name));"
+      else
+        execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON namespaces (lower(name));"
+      end
     end
   end
 
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
-    if Gitlab::Database.version.to_f >= 9.2
-      execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
-    else
-      execute "DROP INDEX IF EXISTS #{INDEX_NAME};"
+    disable_statement_timeout do
+      if Gitlab::Database.version.to_f >= 9.2
+        execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
+      else
+        execute "DROP INDEX IF EXISTS #{INDEX_NAME};"
+      end
     end
   end
 end
diff --git a/db/migrate/20180113220114_rework_redirect_routes_indexes.rb b/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
index ab9971be074..53f82a31203 100644
--- a/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
+++ b/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
@@ -18,51 +18,51 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
   OLD_INDEX_NAME_PATH_LOWER = "index_on_redirect_routes_lower_path"
 
   def up
-    disable_statement_timeout
-
-    # this is a plain btree on a single boolean column. It'll never be
-    # selective enough to be valuable. This class is called by
-    # setup_postgresql.rake so it needs to be able to handle this
-    # index not existing.
-    if index_exists?(:redirect_routes, :permanent)
-      remove_concurrent_index(:redirect_routes, :permanent)
-    end
+    disable_statement_timeout do
+      # this is a plain btree on a single boolean column. It'll never be
+      # selective enough to be valuable. This class is called by
+      # setup_postgresql.rake so it needs to be able to handle this
+      # index not existing.
+      if index_exists?(:redirect_routes, :permanent)
+        remove_concurrent_index(:redirect_routes, :permanent)
+      end
 
-    # If we're on MySQL then the existing index on path is ok. But on
-    # Postgres we need to clean things up:
-    return unless Gitlab::Database.postgresql?
+      # If we're on MySQL then the existing index on path is ok. But on
+      # Postgres we need to clean things up:
+      break unless Gitlab::Database.postgresql?
 
-    if_not_exists = Gitlab::Database.version.to_f >= 9.5 ? "IF NOT EXISTS" : ""
+      if_not_exists = Gitlab::Database.version.to_f >= 9.5 ? "IF NOT EXISTS" : ""
 
-    # Unique index on lower(path) across both types of redirect_routes:
-    execute("CREATE UNIQUE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_UNIQUE} ON redirect_routes (lower(path) varchar_pattern_ops);")
+      # Unique index on lower(path) across both types of redirect_routes:
+      execute("CREATE UNIQUE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_UNIQUE} ON redirect_routes (lower(path) varchar_pattern_ops);")
 
-    # Make two indexes on path -- one for permanent and one for temporary routes:
-    execute("CREATE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_PERM} ON redirect_routes (lower(path) varchar_pattern_ops) where (permanent);")
-    execute("CREATE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_TEMP} ON redirect_routes (lower(path) varchar_pattern_ops) where (not permanent or permanent is null) ;")
+      # Make two indexes on path -- one for permanent and one for temporary routes:
+      execute("CREATE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_PERM} ON redirect_routes (lower(path) varchar_pattern_ops) where (permanent);")
+      execute("CREATE INDEX CONCURRENTLY #{if_not_exists} #{INDEX_NAME_TEMP} ON redirect_routes (lower(path) varchar_pattern_ops) where (not permanent or permanent is null) ;")
 
-    # Remove the old indexes:
+      # Remove the old indexes:
 
-    # This one needed to be on lower(path) but wasn't so it's replaced with the two above
-    execute "DROP INDEX CONCURRENTLY IF EXISTS #{OLD_INDEX_NAME_PATH_TPOPS};"
+      # This one needed to be on lower(path) but wasn't so it's replaced with the two above
+      execute "DROP INDEX CONCURRENTLY IF EXISTS #{OLD_INDEX_NAME_PATH_TPOPS};"
 
-    # This one isn't needed because we only ever do = and LIKE on this
-    # column so the varchar_pattern_ops index is sufficient
-    execute "DROP INDEX CONCURRENTLY IF EXISTS #{OLD_INDEX_NAME_PATH_LOWER};"
+      # This one isn't needed because we only ever do = and LIKE on this
+      # column so the varchar_pattern_ops index is sufficient
+      execute "DROP INDEX CONCURRENTLY IF EXISTS #{OLD_INDEX_NAME_PATH_LOWER};"
+    end
   end
 
   def down
-    disable_statement_timeout
+    disable_statement_timeout do
+      add_concurrent_index(:redirect_routes, :permanent)
 
-    add_concurrent_index(:redirect_routes, :permanent)
+      break unless Gitlab::Database.postgresql?
 
-    return unless Gitlab::Database.postgresql?
+      execute("CREATE INDEX CONCURRENTLY #{OLD_INDEX_NAME_PATH_TPOPS} ON redirect_routes (path varchar_pattern_ops);")
+      execute("CREATE INDEX CONCURRENTLY #{OLD_INDEX_NAME_PATH_LOWER} ON redirect_routes (LOWER(path));")
 
-    execute("CREATE INDEX CONCURRENTLY #{OLD_INDEX_NAME_PATH_TPOPS} ON redirect_routes (path varchar_pattern_ops);")
-    execute("CREATE INDEX CONCURRENTLY #{OLD_INDEX_NAME_PATH_LOWER} ON redirect_routes (LOWER(path));")
-
-    execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_UNIQUE};")
-    execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};")
-    execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};")
+      execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_UNIQUE};")
+      execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};")
+      execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};")
+    end
   end
 end
diff --git a/db/migrate/20180403035759_create_project_ci_cd_settings.rb b/db/migrate/20180403035759_create_project_ci_cd_settings.rb
index 06856af6204..173e662cffc 100644
--- a/db/migrate/20180403035759_create_project_ci_cd_settings.rb
+++ b/db/migrate/20180403035759_create_project_ci_cd_settings.rb
@@ -13,16 +13,16 @@ class CreateProjectCiCdSettings < ActiveRecord::Migration
       end
     end
 
-    disable_statement_timeout
+    disable_statement_timeout do
+      # This particular INSERT will take between 10 and 20 seconds.
+      execute 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects'
 
-    # This particular INSERT will take between 10 and 20 seconds.
-    execute 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects'
+      # We add the index and foreign key separately so the above INSERT statement
+      # takes as little time as possible.
+      add_concurrent_index(:project_ci_cd_settings, :project_id, unique: true)
 
-    # We add the index and foreign key separately so the above INSERT statement
-    # takes as little time as possible.
-    add_concurrent_index(:project_ci_cd_settings, :project_id, unique: true)
-
-    add_foreign_key_with_retry
+      add_foreign_key_with_retry
+    end
   end
 
   def down
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
diff --git a/db/migrate/20180504195842_project_name_lower_index.rb b/db/migrate/20180504195842_project_name_lower_index.rb
index d6f25d3d4ab..74f3673bb03 100644
--- a/db/migrate/20180504195842_project_name_lower_index.rb
+++ b/db/migrate/20180504195842_project_name_lower_index.rb
@@ -13,20 +13,20 @@ class ProjectNameLowerIndex < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
-    execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON projects (LOWER(name))"
+    disable_statement_timeout do
+      execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON projects (LOWER(name))"
+    end
   end
 
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
-    if supports_drop_index_concurrently?
-      execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME}"
-    else
-      execute "DROP INDEX IF EXISTS #{INDEX_NAME}"
+    disable_statement_timeout do
+      if supports_drop_index_concurrently?
+        execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME}"
+      else
+        execute "DROP INDEX IF EXISTS #{INDEX_NAME}"
+      end
     end
   end
 end
diff --git a/db/migrate/20180702124358_remove_orphaned_routes.rb b/db/migrate/20180702124358_remove_orphaned_routes.rb
index 6f6e289ba87..4068e479b6c 100644
--- a/db/migrate/20180702124358_remove_orphaned_routes.rb
+++ b/db/migrate/20180702124358_remove_orphaned_routes.rb
@@ -28,16 +28,16 @@ class RemoveOrphanedRoutes < ActiveRecord::Migration
     # which is pretty close to our 15 second statement timeout. To ensure a
     # smooth deployment procedure we disable the statement timeouts for this
     # migration, just in case.
-    disable_statement_timeout
-
-    # On GitLab.com there are around 4000 orphaned project routes, and around
-    # 150 orphaned namespace routes.
-    [
-      Route.orphaned_project_routes,
-      Route.orphaned_namespace_routes
-    ].each do |relation|
-      relation.each_batch(of: 1_000) do |batch|
-        batch.delete_all
+    disable_statement_timeout do
+      # On GitLab.com there are around 4000 orphaned project routes, and around
+      # 150 orphaned namespace routes.
+      [
+        Route.orphaned_project_routes,
+        Route.orphaned_namespace_routes
+      ].each do |relation|
+        relation.each_batch(of: 1_000) do |batch|
+          batch.delete_all
+        end
       end
     end
   end
diff --git a/db/optional_migrations/composite_primary_keys.rb b/db/optional_migrations/composite_primary_keys.rb
index d45705021b0..b330da13d43 100644
--- a/db/optional_migrations/composite_primary_keys.rb
+++ b/db/optional_migrations/composite_primary_keys.rb
@@ -29,18 +29,20 @@ class CompositePrimaryKeysMigration < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-    TABLES.each do |index|
-      add_primary_key(index)
+    disable_statement_timeout do
+      TABLES.each do |index|
+        add_primary_key(index)
+      end
     end
   end
 
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-    TABLES.each do |index|
-      remove_primary_key(index)
+    disable_statement_timeout do
+      TABLES.each do |index|
+        remove_primary_key(index)
+      end
     end
   end
 
diff --git a/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb b/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
index bba37e32c01..845c6f0557f 100644
--- a/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
+++ b/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
@@ -8,9 +8,9 @@ class EnableAutoCancelPendingPipelinesForAll < ActiveRecord::Migration
   DOWNTIME = false
 
   def up
-    disable_statement_timeout
-
-    update_column_in_batches(:projects, :auto_cancel_pending_pipelines, 1)
+    disable_statement_timeout do
+      update_column_in_batches(:projects, :auto_cancel_pending_pipelines, 1)
+    end
   end
 
   def down
diff --git a/db/post_migrate/20170503004427_update_retried_for_ci_build.rb b/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
index b0b58ab3011..079f0e7511f 100644
--- a/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
+++ b/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
@@ -7,12 +7,12 @@ class UpdateRetriedForCiBuild < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
     if Gitlab::Database.mysql?
       up_mysql
     else
-      up_postgres
+      disable_statement_timeout do
+        up_postgres
+      end
     end
   end
 
diff --git a/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb b/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
index 81e9d050668..5df3ab71648 100644
--- a/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
+++ b/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
@@ -7,20 +7,20 @@ class AddHeadPipelineForEachMergeRequest < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
     pipelines = Arel::Table.new(:ci_pipelines)
     merge_requests = Arel::Table.new(:merge_requests)
 
-    head_id = pipelines
-      .project(Arel::Nodes::NamedFunction.new('max', [pipelines[:id]]))
-      .from(pipelines)
-      .where(pipelines[:ref].eq(merge_requests[:source_branch]))
-      .where(pipelines[:project_id].eq(merge_requests[:source_project_id]))
+    disable_statement_timeout do
+      head_id = pipelines
+        .project(Arel::Nodes::NamedFunction.new('max', [pipelines[:id]]))
+        .from(pipelines)
+        .where(pipelines[:ref].eq(merge_requests[:source_branch]))
+        .where(pipelines[:project_id].eq(merge_requests[:source_project_id]))
 
-    sub_query = Arel::Nodes::SqlLiteral.new(Arel::Nodes::Grouping.new(head_id).to_sql)
+      sub_query = Arel::Nodes::SqlLiteral.new(Arel::Nodes::Grouping.new(head_id).to_sql)
 
-    update_column_in_batches(:merge_requests, :head_pipeline_id, sub_query)
+      update_column_in_batches(:merge_requests, :head_pipeline_id, sub_query)
+    end
   end
 
   def down
diff --git a/db/post_migrate/20170525140254_rename_all_reserved_paths_again.rb b/db/post_migrate/20170525140254_rename_all_reserved_paths_again.rb
index 2125cc046e5..c996ddbec84 100644
--- a/db/post_migrate/20170525140254_rename_all_reserved_paths_again.rb
+++ b/db/post_migrate/20170525140254_rename_all_reserved_paths_again.rb
@@ -87,16 +87,16 @@ class RenameAllReservedPathsAgain < ActiveRecord::Migration
   ].freeze
 
   def up
-    disable_statement_timeout
-
-    TOP_LEVEL_ROUTES.each { |route| rename_root_paths(route) }
-    PROJECT_WILDCARD_ROUTES.each { |route| rename_wildcard_paths(route) }
-    GROUP_ROUTES.each { |route| rename_child_paths(route) }
+    disable_statement_timeout do
+      TOP_LEVEL_ROUTES.each { |route| rename_root_paths(route) }
+      PROJECT_WILDCARD_ROUTES.each { |route| rename_wildcard_paths(route) }
+      GROUP_ROUTES.each { |route| rename_child_paths(route) }
+    end
   end
 
   def down
-    disable_statement_timeout
-
-    revert_renames
+    disable_statement_timeout do
+      revert_renames
+    end
   end
 end
diff --git a/db/post_migrate/20170526185842_migrate_pipeline_stages.rb b/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
index afd4db183c2..736aff77f02 100644
--- a/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
+++ b/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
@@ -6,17 +6,17 @@ class MigratePipelineStages < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
-    execute <<-SQL.strip_heredoc
-      INSERT INTO ci_stages (project_id, pipeline_id, name)
-        SELECT project_id, commit_id, stage FROM ci_builds
-          WHERE stage IS NOT NULL
-          AND stage_id IS NULL
-          AND EXISTS (SELECT 1 FROM projects WHERE projects.id = ci_builds.project_id)
-          AND EXISTS (SELECT 1 FROM ci_pipelines WHERE ci_pipelines.id = ci_builds.commit_id)
-          GROUP BY project_id, commit_id, stage
-          ORDER BY MAX(stage_idx)
-    SQL
+    disable_statement_timeout do
+      execute <<-SQL.strip_heredoc
+        INSERT INTO ci_stages (project_id, pipeline_id, name)
+          SELECT project_id, commit_id, stage FROM ci_builds
+            WHERE stage IS NOT NULL
+            AND stage_id IS NULL
+            AND EXISTS (SELECT 1 FROM projects WHERE projects.id = ci_builds.project_id)
+            AND EXISTS (SELECT 1 FROM ci_pipelines WHERE ci_pipelines.id = ci_builds.commit_id)
+            GROUP BY project_id, commit_id, stage
+            ORDER BY MAX(stage_idx)
+      SQL
+    end
   end
 end
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..a7bfba0ab2b 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 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 do
+      update_column_in_batches(:ci_builds, :stage_id, nil)
+    end
   end
 end
diff --git a/db/post_migrate/20170711145558_migrate_stages_statuses.rb b/db/post_migrate/20170711145558_migrate_stages_statuses.rb
index 65755c0e824..265f7317b9b 100644
--- a/db/post_migrate/20170711145558_migrate_stages_statuses.rb
+++ b/db/post_migrate/20170711145558_migrate_stages_statuses.rb
@@ -26,9 +26,9 @@ class MigrateStagesStatuses < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout
-
-    # rubocop:disable Migration/UpdateLargeTable
-    update_column_in_batches(:ci_stages, :status, nil)
+    disable_statement_timeout do
+      # rubocop:disable Migration/UpdateLargeTable
+      update_column_in_batches(:ci_stages, :status, nil)
+    end
   end
 end
diff --git a/db/post_migrate/20171207150343_remove_soft_removed_objects.rb b/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
index 3e2dedfdd6a..3109b6dbf8e 100644
--- a/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
+++ b/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
@@ -78,12 +78,12 @@ class RemoveSoftRemovedObjects < ActiveRecord::Migration
   MODELS = [Issue, MergeRequest, CiPipelineSchedule, CiTrigger].freeze
 
   def up
-    disable_statement_timeout
-
-    remove_personal_routes
-    remove_personal_namespaces
-    remove_group_namespaces
-    remove_simple_soft_removed_rows
+    disable_statement_timeout do
+      remove_personal_routes
+      remove_personal_namespaces
+      remove_group_namespaces
+      remove_simple_soft_removed_rows
+    end
   end
 
   def down
diff --git a/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb b/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
index 61ea85eb2a7..269f1287f91 100644
--- a/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
+++ b/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
@@ -38,29 +38,29 @@ class RemoveRedundantPipelineStages < ActiveRecord::Migration
   end
 
   def remove_redundant_pipeline_stages!
-    disable_statement_timeout
-
-    redundant_stages_ids = <<~SQL
-      SELECT id FROM ci_stages WHERE (pipeline_id, name) IN (
-        SELECT pipeline_id, name FROM ci_stages
-          GROUP BY pipeline_id, name HAVING COUNT(*) > 1
-      )
-    SQL
-
-    execute <<~SQL
-      UPDATE ci_builds SET stage_id = NULL WHERE stage_id IN (#{redundant_stages_ids})
-    SQL
-
-    if Gitlab::Database.postgresql?
-      execute <<~SQL
-        DELETE FROM ci_stages WHERE id IN (#{redundant_stages_ids})
+    disable_statement_timeout do
+      redundant_stages_ids = <<~SQL
+        SELECT id FROM ci_stages WHERE (pipeline_id, name) IN (
+          SELECT pipeline_id, name FROM ci_stages
+            GROUP BY pipeline_id, name HAVING COUNT(*) > 1
+        )
       SQL
-    else # We can't modify a table we are selecting from on MySQL
+
       execute <<~SQL
-        DELETE a FROM ci_stages AS a, ci_stages AS b
-          WHERE a.pipeline_id = b.pipeline_id AND a.name = b.name
-            AND a.id <> b.id
+        UPDATE ci_builds SET stage_id = NULL WHERE stage_id IN (#{redundant_stages_ids})
       SQL
+
+      if Gitlab::Database.postgresql?
+        execute <<~SQL
+          DELETE FROM ci_stages WHERE id IN (#{redundant_stages_ids})
+        SQL
+      else # We can't modify a table we are selecting from on MySQL
+        execute <<~SQL
+          DELETE a FROM ci_stages AS a, ci_stages AS b
+            WHERE a.pipeline_id = b.pipeline_id AND a.name = b.name
+              AND a.id <> b.id
+        SQL
+      end
     end
   end
 end
diff --git a/db/post_migrate/20180305100050_remove_permanent_from_redirect_routes.rb b/db/post_migrate/20180305100050_remove_permanent_from_redirect_routes.rb
index db5165dbe70..aa19732ca1c 100644
--- a/db/post_migrate/20180305100050_remove_permanent_from_redirect_routes.rb
+++ b/db/post_migrate/20180305100050_remove_permanent_from_redirect_routes.rb
@@ -15,10 +15,10 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     # ReworkRedirectRoutesIndexes:
     # https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/16211
     if Gitlab::Database.postgresql?
-      disable_statement_timeout
-
-      execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};"
-      execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};"
+      disable_statement_timeout do
+        execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};"
+        execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};"
+      end
     end
 
     remove_column(:redirect_routes, :permanent)
@@ -28,10 +28,10 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     add_column(:redirect_routes, :permanent, :boolean)
 
     if Gitlab::Database.postgresql?
-      disable_statement_timeout
-
-      execute("CREATE INDEX CONCURRENTLY #{INDEX_NAME_PERM} ON redirect_routes (lower(path) varchar_pattern_ops) where (permanent);")
-      execute("CREATE INDEX CONCURRENTLY #{INDEX_NAME_TEMP} ON redirect_routes (lower(path) varchar_pattern_ops) where (not permanent or permanent is null) ;")
+      disable_statement_timeout do
+        execute("CREATE INDEX CONCURRENTLY #{INDEX_NAME_PERM} ON redirect_routes (lower(path) varchar_pattern_ops) where (permanent);")
+        execute("CREATE INDEX CONCURRENTLY #{INDEX_NAME_TEMP} ON redirect_routes (lower(path) varchar_pattern_ops) where (not permanent or permanent is null) ;")
+      end
     end
   end
 end
diff --git a/db/post_migrate/20180306164012_add_path_index_to_redirect_routes.rb b/db/post_migrate/20180306164012_add_path_index_to_redirect_routes.rb
index d6fb4f06695..ca9212fae27 100644
--- a/db/post_migrate/20180306164012_add_path_index_to_redirect_routes.rb
+++ b/db/post_migrate/20180306164012_add_path_index_to_redirect_routes.rb
@@ -20,10 +20,10 @@ class AddPathIndexToRedirectRoutes < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
-    unless index_exists_by_name?(:redirect_routes, INDEX_NAME)
-      execute("CREATE UNIQUE INDEX CONCURRENTLY #{INDEX_NAME} ON redirect_routes (lower(path) varchar_pattern_ops);")
+    disable_statement_timeout do
+      unless index_exists_by_name?(:redirect_routes, INDEX_NAME)
+        execute("CREATE UNIQUE INDEX CONCURRENTLY #{INDEX_NAME} ON redirect_routes (lower(path) varchar_pattern_ops);")
+      end
     end
   end
 
diff --git a/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb b/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
index e19387bce1e..c32123454f9 100644
--- a/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
+++ b/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
@@ -17,13 +17,13 @@ class RescheduleBuildsStagesMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
-    Build.where('stage_id IS NULL').tap do |relation|
-      queue_background_migration_jobs_by_range_at_intervals(relation,
-                                                            MIGRATION,
-                                                            5.minutes,
-                                                            batch_size: BATCH_SIZE)
+    disable_statement_timeout do
+      Build.where('stage_id IS NULL').tap do |relation|
+        queue_background_migration_jobs_by_range_at_intervals(relation,
+                                                              MIGRATION,
+                                                              5.minutes,
+                                                              batch_size: BATCH_SIZE)
+      end
     end
   end
 
diff --git a/db/post_migrate/20180420080616_schedule_stages_index_migration.rb b/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
index 1d0daad002f..eb82f098639 100644
--- a/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
+++ b/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
@@ -13,13 +13,13 @@ class ScheduleStagesIndexMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
-    Stage.all.tap do |relation|
-      queue_background_migration_jobs_by_range_at_intervals(relation,
-                                                            MIGRATION,
-                                                            5.minutes,
-                                                            batch_size: BATCH_SIZE)
+    disable_statement_timeout do
+      Stage.all.tap do |relation|
+        queue_background_migration_jobs_by_range_at_intervals(relation,
+                                                              MIGRATION,
+                                                              5.minutes,
+                                                              batch_size: BATCH_SIZE)
+      end
     end
   end
 
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
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index f39b3b6eb5b..7f012312819 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -58,7 +58,6 @@ module Gitlab
 
         if Database.postgresql?
           options = options.merge({ algorithm: :concurrently })
-          disable_statement_timeout
         end
 
         if index_exists?(table_name, column_name, options)
@@ -66,7 +65,9 @@ module Gitlab
           return
         end
 
-        add_index(table_name, column_name, options)
+        disable_statement_timeout do
+          add_index(table_name, column_name, options)
+        end
       end
 
       # Removes an existed index, concurrently when supported
@@ -87,7 +88,6 @@ module Gitlab
 
         if supports_drop_index_concurrently?
           options = options.merge({ algorithm: :concurrently })
-          disable_statement_timeout
         end
 
         unless index_exists?(table_name, column_name, options)
@@ -95,7 +95,9 @@ module Gitlab
           return
         end
 
-        remove_index(table_name, options.merge({ column: column_name }))
+        disable_statement_timeout do
+          remove_index(table_name, options.merge({ column: column_name }))
+        end
       end
 
       # Removes an existing index, concurrently when supported
@@ -116,7 +118,6 @@ module Gitlab
 
         if supports_drop_index_concurrently?
           options = options.merge({ algorithm: :concurrently })
-          disable_statement_timeout
         end
 
         unless index_exists_by_name?(table_name, index_name)
@@ -124,7 +125,9 @@ module Gitlab
           return
         end
 
-        remove_index(table_name, options.merge({ name: index_name }))
+        disable_statement_timeout do
+          remove_index(table_name, options.merge({ name: index_name }))
+        end
       end
 
       # Only available on Postgresql >= 9.2
@@ -171,8 +174,6 @@ module Gitlab
           on_delete = 'SET NULL' if on_delete == :nullify
         end
 
-        disable_statement_timeout
-
         key_name = concurrent_foreign_key_name(source, column)
 
         unless foreign_key_exists?(source, target, column: column)
@@ -199,7 +200,9 @@ module Gitlab
         # while running.
         #
         # Note this is a no-op in case the constraint is VALID already
-        execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
+        disable_statement_timeout do
+          execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
+        end
       end
 
       def foreign_key_exists?(source, target = nil, column: nil)
@@ -224,8 +227,48 @@ module Gitlab
       # Long-running migrations may take more than the timeout allowed by
       # the database. Disable the session's statement timeout to ensure
       # migrations don't get killed prematurely. (PostgreSQL only)
+      #
+      # There are two possible ways to disable the statement timeout:
+      #
+      # - Per transaction (this is the preferred and default mode)
+      # - Per connection (requires a cleanup after the execution)
+      #
+      # When using a per connection disable statement, code must be inside
+      # a block so we can automatically execute `RESET ALL` after block finishes
+      # otherwise the statement will still be disabled until connection is dropped
+      # or `RESET ALL` is executed
       def disable_statement_timeout
-        execute('SET statement_timeout TO 0') if Database.postgresql?
+        # bypass disabled_statement logic when not using postgres, but still execute block when one is given
+        unless Database.postgresql?
+          if block_given?
+            yield
+          end
+
+          return
+        end
+
+        if block_given?
+          begin
+            execute('SET statement_timeout TO 0')
+
+            yield
+          ensure
+            execute('RESET ALL')
+          end
+        else
+          unless transaction_open?
+            raise <<~ERROR
+              Cannot call disable_statement_timeout() without a transaction open or outside of a transaction block.
+              If you don't want to use a transaction wrap your code in a block call:
+
+              disable_statement_timeout { # code that requires disabled statement here }
+
+              This will make sure statement_timeout is disabled before and reset after the block execution is finished.
+            ERROR
+          end
+
+          execute('SET LOCAL statement_timeout TO 0')
+        end
       end
 
       def true_value
@@ -367,30 +410,30 @@ module Gitlab
             'in the body of your migration class'
         end
 
-        disable_statement_timeout
-
-        transaction do
-          if limit
-            add_column(table, column, type, default: nil, limit: limit)
-          else
-            add_column(table, column, type, default: nil)
+        disable_statement_timeout do
+          transaction do
+            if limit
+              add_column(table, column, type, default: nil, limit: limit)
+            else
+              add_column(table, column, type, default: nil)
+            end
+
+            # Changing the default before the update ensures any newly inserted
+            # rows already use the proper default value.
+            change_column_default(table, column, default)
           end
 
-          # Changing the default before the update ensures any newly inserted
-          # rows already use the proper default value.
-          change_column_default(table, column, default)
-        end
-
-        begin
-          update_column_in_batches(table, column, default, &block)
+          begin
+            update_column_in_batches(table, column, default, &block)
 
-          change_column_null(table, column, false) unless allow_null
-        # We want to rescue _all_ exceptions here, even those that don't inherit
-        # from StandardError.
-        rescue Exception => error # rubocop: disable all
-          remove_column(table, column)
+            change_column_null(table, column, false) unless allow_null
+          # We want to rescue _all_ exceptions here, even those that don't inherit
+          # from StandardError.
+          rescue Exception => error # rubocop: disable all
+            remove_column(table, column)
 
-          raise error
+            raise error
+          end
         end
       end
 
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index eb7148ff108..4e83b27e4a5 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -48,10 +48,10 @@ describe Gitlab::Database::MigrationHelpers do
         allow(model).to receive(:transaction_open?).and_return(false)
       end
 
-      context 'using PostgreSQL' do
+      context 'using PostgreSQL', :postgresql do
         before do
           allow(Gitlab::Database).to receive(:postgresql?).and_return(true)
-          allow(model).to receive(:disable_statement_timeout)
+          allow(model).to receive(:disable_statement_timeout).and_call_original
         end
 
         it 'creates the index concurrently' do
@@ -114,12 +114,12 @@ describe Gitlab::Database::MigrationHelpers do
       before do
         allow(model).to receive(:transaction_open?).and_return(false)
         allow(model).to receive(:index_exists?).and_return(true)
+        allow(model).to receive(:disable_statement_timeout).and_call_original
       end
 
       context 'using PostgreSQL' do
         before do
           allow(model).to receive(:supports_drop_index_concurrently?).and_return(true)
-          allow(model).to receive(:disable_statement_timeout)
         end
 
         describe 'by column name' do
@@ -162,7 +162,7 @@ describe Gitlab::Database::MigrationHelpers do
 
       context 'using MySQL' do
         it 'removes an index' do
-          expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
+          expect(Gitlab::Database).to receive(:postgresql?).and_return(false).twice
 
           expect(model).to receive(:remove_index)
             .with(:users, { column: :foo })
@@ -224,21 +224,26 @@ describe Gitlab::Database::MigrationHelpers do
 
       context 'using PostgreSQL' do
         before do
+          allow(Gitlab::Database).to receive(:postgresql?).and_return(true)
           allow(Gitlab::Database).to receive(:mysql?).and_return(false)
         end
 
         it 'creates a concurrent foreign key and validates it' do
-          expect(model).to receive(:disable_statement_timeout)
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:execute).with(/statement_timeout/)
           expect(model).to receive(:execute).ordered.with(/NOT VALID/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).with(/RESET ALL/)
 
           model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
         end
 
         it 'appends a valid ON DELETE statement' do
-          expect(model).to receive(:disable_statement_timeout)
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:execute).with(/statement_timeout/)
           expect(model).to receive(:execute).with(/ON DELETE SET NULL/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).with(/RESET ALL/)
 
           model.add_concurrent_foreign_key(:projects, :users,
                                            column: :user_id,
@@ -291,13 +296,68 @@ describe Gitlab::Database::MigrationHelpers do
 
   describe '#disable_statement_timeout' do
     context 'using PostgreSQL' do
-      it 'disables statement timeouts' do
+      it 'disables statement timeouts to current transaction only' do
         expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
 
-        expect(model).to receive(:execute).with('SET statement_timeout TO 0')
+        expect(model).to receive(:execute).with('SET LOCAL statement_timeout TO 0')
 
         model.disable_statement_timeout
       end
+
+      # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
+      context 'with real environment', :postgresql, :delete do
+        before do
+          model.execute("SET statement_timeout TO '20000'")
+        end
+
+        after do
+          model.execute('RESET ALL')
+        end
+
+        it 'defines statement to 0 only for current transaction' do
+          expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+
+          model.connection.transaction do
+            model.disable_statement_timeout
+            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+          end
+
+          expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+        end
+      end
+
+      context 'when passing a blocks' do
+        it 'disables statement timeouts on session level and executes the block' do
+          expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
+          expect(model).to receive(:execute).with('SET statement_timeout TO 0')
+          expect(model).to receive(:execute).with('RESET ALL')
+
+          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
+        end
+
+        # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
+        context 'with real environment', :postgresql, :delete do
+          before do
+            model.execute("SET statement_timeout TO '20000'")
+          end
+
+          after do
+            model.execute('RESET ALL')
+          end
+
+          it 'defines statement to 0 for any code run inside the block' do
+            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+
+            model.disable_statement_timeout do
+              model.connection.transaction do
+                expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+              end
+
+              expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
+            end
+          end
+        end
+      end
     end
 
     context 'using MySQL' do
@@ -308,6 +368,16 @@ describe Gitlab::Database::MigrationHelpers do
 
         model.disable_statement_timeout
       end
+
+      context 'when passing a blocks' do
+        it 'executes the block of code' do
+          expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
+
+          expect(model).not_to receive(:execute)
+
+          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
+        end
+      end
     end
   end