disable_statement_timeout doesn't require any argument anymore

it will decide the method for disable statement_timeout upon
per transaction or per session, based on how it's called.

When calling with a block, block will be executed and it will use
session based statement_timeout, otherwise will default to existing
behavior.

25 files changed, 101 insertions, 114 deletions

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 35e95e6b80a..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,7 +106,7 @@ 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(transaction: false) do
+            disable_statement_timeout do
               remove_orphans(*queue.pop) until queue.empty?
 
               steal_from_queues(queues - [queue])
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 28f36e8bf18..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,7 +25,7 @@ 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(transaction: false) do
+      disable_statement_timeout do
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
       end
     else
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 1c14f274b8b..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,7 +8,7 @@ class AddIndexOnNamespacesLowerName < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    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));"
@@ -21,7 +21,7 @@ class AddIndexOnNamespacesLowerName < ActiveRecord::Migration
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       if Gitlab::Database.version.to_f >= 9.2
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
       else
diff --git a/db/migrate/20180113220114_rework_redirect_routes_indexes.rb b/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
index 7748104f436..53f82a31203 100644
--- a/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
+++ b/db/migrate/20180113220114_rework_redirect_routes_indexes.rb
@@ -18,7 +18,7 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
   OLD_INDEX_NAME_PATH_LOWER = "index_on_redirect_routes_lower_path"
 
   def up
-    disable_statement_timeout(transaction: false) do
+    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
@@ -29,7 +29,7 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
 
       # 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?
+      break unless Gitlab::Database.postgresql?
 
       if_not_exists = Gitlab::Database.version.to_f >= 9.5 ? "IF NOT EXISTS" : ""
 
@@ -52,10 +52,10 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       add_concurrent_index(:redirect_routes, :permanent)
 
-      return unless Gitlab::Database.postgresql?
+      break 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));")
diff --git a/db/migrate/20180403035759_create_project_ci_cd_settings.rb b/db/migrate/20180403035759_create_project_ci_cd_settings.rb
index 61b43ef2038..173e662cffc 100644
--- a/db/migrate/20180403035759_create_project_ci_cd_settings.rb
+++ b/db/migrate/20180403035759_create_project_ci_cd_settings.rb
@@ -13,7 +13,7 @@ class CreateProjectCiCdSettings < ActiveRecord::Migration
       end
     end
 
-    disable_statement_timeout(transaction: false) do
+    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'
 
diff --git a/db/migrate/20180420010616_cleanup_build_stage_migration.rb b/db/migrate/20180420010616_cleanup_build_stage_migration.rb
index a535bc03194..5e9fe756efd 100644
--- a/db/migrate/20180420010616_cleanup_build_stage_migration.rb
+++ b/db/migrate/20180420010616_cleanup_build_stage_migration.rb
@@ -14,7 +14,7 @@ class CleanupBuildStageMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       ##
       # We steal from the background migrations queue to catch up with the
       # scheduled migrations set.
@@ -55,7 +55,7 @@ class CleanupBuildStageMigration < ActiveRecord::Migration
 
   def down
     if index_exists_by_name?(:ci_builds, TMP_INDEX)
-      disable_statement_timeout(transaction: false) do
+      disable_statement_timeout do
         remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
       end
     end
diff --git a/db/migrate/20180504195842_project_name_lower_index.rb b/db/migrate/20180504195842_project_name_lower_index.rb
index 856ccf0620e..74f3673bb03 100644
--- a/db/migrate/20180504195842_project_name_lower_index.rb
+++ b/db/migrate/20180504195842_project_name_lower_index.rb
@@ -13,7 +13,7 @@ class ProjectNameLowerIndex < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON projects (LOWER(name))"
     end
   end
@@ -21,7 +21,7 @@ class ProjectNameLowerIndex < ActiveRecord::Migration
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       if supports_drop_index_concurrently?
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME}"
       else
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 7fbc3d4461d..b330da13d43 100644
--- a/db/optional_migrations/composite_primary_keys.rb
+++ b/db/optional_migrations/composite_primary_keys.rb
@@ -29,7 +29,7 @@ class CompositePrimaryKeysMigration < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       TABLES.each do |index|
         add_primary_key(index)
       end
@@ -39,7 +39,7 @@ class CompositePrimaryKeysMigration < ActiveRecord::Migration
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       TABLES.each do |index|
         remove_primary_key(index)
       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 8b05a6bea75..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,7 +8,7 @@ class EnableAutoCancelPendingPipelinesForAll < ActiveRecord::Migration
   DOWNTIME = false
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       update_column_in_batches(:projects, :auto_cancel_pending_pipelines, 1)
     end
   end
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 70de1fe2757..079f0e7511f 100644
--- a/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
+++ b/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
@@ -10,7 +10,7 @@ class UpdateRetriedForCiBuild < ActiveRecord::Migration
     if Gitlab::Database.mysql?
       up_mysql
     else
-      disable_statement_timeout(transaction: false) do
+      disable_statement_timeout do
         up_postgres
       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 188817c3324..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
@@ -10,7 +10,7 @@ class AddHeadPipelineForEachMergeRequest < ActiveRecord::Migration
     pipelines = Arel::Table.new(:ci_pipelines)
     merge_requests = Arel::Table.new(:merge_requests)
 
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       head_id = pipelines
         .project(Arel::Nodes::NamedFunction.new('max', [pipelines[:id]]))
         .from(pipelines)
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 7edbab67c29..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,7 +87,7 @@ class RenameAllReservedPathsAgain < ActiveRecord::Migration
   ].freeze
 
   def up
-    disable_statement_timeout(transaction: false) do
+    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) }
@@ -95,7 +95,7 @@ class RenameAllReservedPathsAgain < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       revert_renames
     end
   end
diff --git a/db/post_migrate/20170526185842_migrate_pipeline_stages.rb b/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
index 5df4ad8b666..736aff77f02 100644
--- a/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
+++ b/db/post_migrate/20170526185842_migrate_pipeline_stages.rb
@@ -6,7 +6,7 @@ class MigratePipelineStages < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout(transaction: false) do
+    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
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 4d5ac81382e..a7bfba0ab2b 100644
--- a/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
+++ b/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
@@ -13,7 +13,7 @@ class MigrateBuildStageReferenceAgain < ActiveRecord::Migration
            AND ci_stages.name = ci_builds.stage)
     SQL
 
-    disable_statement_timeout(transaction: false) do
+    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
@@ -21,7 +21,7 @@ class MigrateBuildStageReferenceAgain < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       update_column_in_batches(:ci_builds, :stage_id, nil)
     end
   end
diff --git a/db/post_migrate/20170711145558_migrate_stages_statuses.rb b/db/post_migrate/20170711145558_migrate_stages_statuses.rb
index 73ada165d0b..265f7317b9b 100644
--- a/db/post_migrate/20170711145558_migrate_stages_statuses.rb
+++ b/db/post_migrate/20170711145558_migrate_stages_statuses.rb
@@ -26,8 +26,7 @@ class MigrateStagesStatuses < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout(transaction: false) do
-
+    disable_statement_timeout do
       # rubocop:disable Migration/UpdateLargeTable
       update_column_in_batches(:ci_stages, :status, nil)
     end
diff --git a/db/post_migrate/20171207150343_remove_soft_removed_objects.rb b/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
index ded2a372b3c..3109b6dbf8e 100644
--- a/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
+++ b/db/post_migrate/20171207150343_remove_soft_removed_objects.rb
@@ -78,7 +78,7 @@ class RemoveSoftRemovedObjects < ActiveRecord::Migration
   MODELS = [Issue, MergeRequest, CiPipelineSchedule, CiTrigger].freeze
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       remove_personal_routes
       remove_personal_namespaces
       remove_group_namespaces
diff --git a/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb b/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
index 8352a0f8191..269f1287f91 100644
--- a/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
+++ b/db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb
@@ -38,7 +38,7 @@ class RemoveRedundantPipelineStages < ActiveRecord::Migration
   end
 
   def remove_redundant_pipeline_stages!
-    disable_statement_timeout(transaction: false) do
+    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
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 41075f29ce3..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,7 +15,7 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     # ReworkRedirectRoutesIndexes:
     # https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/16211
     if Gitlab::Database.postgresql?
-      disable_statement_timeout(transaction: false) do
+      disable_statement_timeout do
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};"
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};"
       end
@@ -28,7 +28,7 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     add_column(:redirect_routes, :permanent, :boolean)
 
     if Gitlab::Database.postgresql?
-      disable_statement_timeout(transaction: false) do
+      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
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 894ed2cf21e..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,7 +20,7 @@ class AddPathIndexToRedirectRoutes < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout(transaction: false) do
+    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
diff --git a/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb b/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
index 8f7b2e8d852..c32123454f9 100644
--- a/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
+++ b/db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb
@@ -17,7 +17,7 @@ class RescheduleBuildsStagesMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       Build.where('stage_id IS NULL').tap do |relation|
         queue_background_migration_jobs_by_range_at_intervals(relation,
                                                               MIGRATION,
diff --git a/db/post_migrate/20180420080616_schedule_stages_index_migration.rb b/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
index c8bdf124810..eb82f098639 100644
--- a/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
+++ b/db/post_migrate/20180420080616_schedule_stages_index_migration.rb
@@ -13,7 +13,7 @@ class ScheduleStagesIndexMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       Stage.all.tap do |relation|
         queue_background_migration_jobs_by_range_at_intervals(relation,
                                                               MIGRATION,
diff --git a/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb b/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
index 561934d4353..5418f442e79 100644
--- a/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
+++ b/db/post_migrate/20180604123514_cleanup_stages_position_migration.rb
@@ -12,7 +12,7 @@ class CleanupStagesPositionMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout(transaction: false) do
+    disable_statement_timeout do
       Gitlab::BackgroundMigration.steal('MigrateStageIndex')
 
       unless index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
@@ -37,7 +37,7 @@ class CleanupStagesPositionMigration < ActiveRecord::Migration
 
   def down
     if index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
-      disable_statement_timeout(transaction: false) do
+      disable_statement_timeout do
         remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
       end
     end
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 1e6a620dd66..7f012312819 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -65,7 +65,7 @@ module Gitlab
           return
         end
 
-        disable_statement_timeout(transaction: false) do
+        disable_statement_timeout do
           add_index(table_name, column_name, options)
         end
       end
@@ -95,7 +95,7 @@ module Gitlab
           return
         end
 
-        disable_statement_timeout(transaction: false) do
+        disable_statement_timeout do
           remove_index(table_name, options.merge({ column: column_name }))
         end
       end
@@ -125,7 +125,7 @@ module Gitlab
           return
         end
 
-        disable_statement_timeout(transaction: false) do
+        disable_statement_timeout do
           remove_index(table_name, options.merge({ name: index_name }))
         end
       end
@@ -200,7 +200,7 @@ module Gitlab
         # while running.
         #
         # Note this is a no-op in case the constraint is VALID already
-        disable_statement_timeout(transaction: false) do
+        disable_statement_timeout do
           execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
         end
       end
@@ -233,13 +233,11 @@ module Gitlab
       # - 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 `RESET ALL` after it has executed otherwise the statement
-      # will still be disabled until connection is dropped or `RESET ALL` is executed
-      #
-      # - +transaction:+ true to disable for current transaction only *(default)*
-      # - +transaction:+ false to disable for current session (requires block)
-      def disable_statement_timeout(transaction: true)
+      # 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
         # bypass disabled_statement logic when not using postgres, but still execute block when one is given
         unless Database.postgresql?
           if block_given?
@@ -249,26 +247,27 @@ module Gitlab
           return
         end
 
-        if transaction
-          unless transaction_open?
-            raise 'disable_statement_timeout() cannot be run without a transaction, ' \
-              'use it inside a transaction block. Alternatively you can use: ' \
-              'disable_statement_timeout(transaction: false) { #code here } to make sure ' \
-              'statement_timeout is reset after the block execution is finished.'
-          end
+        if block_given?
+          begin
+            execute('SET statement_timeout TO 0')
 
-          execute('SET LOCAL statement_timeout TO 0')
-        else
-          unless block_given?
-            raise ArgumentError, 'disable_statement_timeout(transaction: false) requires a block encapsulating' \
-              'code that will be executed with the statement_timeout disabled.'
+            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:
 
-          execute('SET statement_timeout TO 0')
+              disable_statement_timeout { # code that requires disabled statement here }
 
-          yield
+              This will make sure statement_timeout is disabled before and reset after the block execution is finished.
+            ERROR
+          end
 
-          execute('RESET ALL')
+          execute('SET LOCAL statement_timeout TO 0')
         end
       end
 
@@ -411,7 +410,7 @@ module Gitlab
             'in the body of your migration class'
         end
 
-        disable_statement_timeout(transaction: false) do
+        disable_statement_timeout do
           transaction do
             if limit
               add_column(table, column, type, default: nil, limit: limit)
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 28a03e6f6f6..4e83b27e4a5 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -296,52 +296,43 @@ describe Gitlab::Database::MigrationHelpers do
 
   describe '#disable_statement_timeout' do
     context 'using PostgreSQL' do
-      context 'with transaction: true' 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 LOCAL statement_timeout TO 0')
+      it 'disables statement timeouts to current transaction only' do
+        expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
 
-          model.disable_statement_timeout
-        end
+        expect(model).to receive(:execute).with('SET LOCAL statement_timeout TO 0')
 
-        # 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
+        model.disable_statement_timeout
+      end
 
-          after do
-            model.execute('RESET ALL')
-          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
 
-          it 'defines statement to 0 only for current transaction' do
-            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+        after do
+          model.execute('RESET ALL')
+        end
 
-            model.connection.transaction do
-              model.disable_statement_timeout
-              expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
-            end
+        it 'defines statement to 0 only for current transaction' do
+          expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
 
-            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 'with transaction: false' do
-        it 'disables statement timeouts on session level' do
+      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')
 
-          model.disable_statement_timeout(transaction: false) do
-            # no-op
-          end
-        end
-
-        it 'raises error when no block is given' do
-          expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
-          expect { model.disable_statement_timeout(transaction: false) }.to raise_error(ArgumentError)
+          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
         end
 
         # this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
@@ -354,10 +345,10 @@ describe Gitlab::Database::MigrationHelpers do
             model.execute('RESET ALL')
           end
 
-          it 'defines statement to 0 for any code run inside block' do
+          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(transaction: false) do
+            model.disable_statement_timeout do
               model.connection.transaction do
                 expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
               end
@@ -370,23 +361,21 @@ describe Gitlab::Database::MigrationHelpers do
     end
 
     context 'using MySQL' do
-      context 'with transaction: true' do
-        it 'does nothing' do
-          expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
+      it 'does nothing' do
+        expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
 
-          expect(model).not_to receive(:execute)
+        expect(model).not_to receive(:execute)
 
-          model.disable_statement_timeout
-        end
+        model.disable_statement_timeout
       end
 
-      context 'with transaction: false' do
-        it 'executes yielded code' do
+      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(transaction: false, &block) }.to yield_control
+          expect { |block| model.disable_statement_timeout(&block) }.to yield_control
         end
       end
     end