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

changelogs/unreleased/48967-disable-statement-timeout.yml

+---
+title: disable_statement_timeout no longer leak to other migrations
+merge_request: 20503
+author:
+type: fixed

db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb

@@ -106,14 +106,14 @@ 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?
 
               steal_from_queues(queues - [queue])
             end
           end
         end
+      end
 
       threads.each(&:join)
     end

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
+      disable_statement_timeout do
         execute "DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME};"
+      end
     else
       execute "DROP INDEX IF EXISTS #{INDEX_NAME};"
     end

db/migrate/20171220191323_add_index_on_namespaces_lower_name.rb

@@ -8,8 +8,7 @@ class AddIndexOnNamespacesLowerName < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
+    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));"
@@ -17,16 +16,17 @@ class AddIndexOnNamespacesLowerName < ActiveRecord::Migration
         execute "CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON namespaces (lower(name));"
       end
     end
+  end
 
   def down
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
+    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

db/migrate/20180113220114_rework_redirect_routes_indexes.rb

@@ -18,8 +18,7 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
   OLD_INDEX_NAME_PATH_LOWER = "index_on_redirect_routes_lower_path"
 
   def up
-    disable_statement_timeout
-
+    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
@@ -30,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" : ""
 
@@ -50,13 +49,13 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
       # 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)
 
-    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));")
@@ -65,4 +64,5 @@ class ReworkRedirectRoutesIndexes < ActiveRecord::Migration
       execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_PERM};")
       execute("DROP INDEX CONCURRENTLY IF EXISTS #{INDEX_NAME_TEMP};")
     end
+  end
 end

db/migrate/20180403035759_create_project_ci_cd_settings.rb

@@ -13,8 +13,7 @@ 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'
 
@@ -24,6 +23,7 @@ class CreateProjectCiCdSettings < ActiveRecord::Migration
 
       add_foreign_key_with_retry
     end
+  end
 
   def down
     drop_table :project_ci_cd_settings

db/migrate/20180420010616_cleanup_build_stage_migration.rb

@@ -14,8 +14,7 @@ class CleanupBuildStageMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       ##
       # We steal from the background migrations queue to catch up with the
       # scheduled migrations set.
@@ -52,10 +51,13 @@ class CleanupBuildStageMigration < ActiveRecord::Migration
       #
       remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
     end
+  end
 
   def down
     if index_exists_by_name?(:ci_builds, TMP_INDEX)
+      disable_statement_timeout do
         remove_concurrent_index_by_name(:ci_builds, TMP_INDEX)
       end
     end
+  end
 end

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
-
+    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
-
+    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

db/migrate/20180702124358_remove_orphaned_routes.rb

@@ -28,8 +28,7 @@ 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
-
+    disable_statement_timeout do
       # On GitLab.com there are around 4000 orphaned project routes, and around
       # 150 orphaned namespace routes.
       [
@@ -41,6 +40,7 @@ class RemoveOrphanedRoutes < ActiveRecord::Migration
         end
       end
     end
+  end
 
   def down
     # There is no way to restore orphaned routes, and this doesn't make any

db/optional_migrations/composite_primary_keys.rb

@@ -29,20 +29,22 @@ class CompositePrimaryKeysMigration < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
+    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
+    disable_statement_timeout do
       TABLES.each do |index|
         remove_primary_key(index)
       end
     end
+  end
 
   private
 

db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb

@@ -8,10 +8,10 @@ class EnableAutoCancelPendingPipelinesForAll < ActiveRecord::Migration
   DOWNTIME = false
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       update_column_in_batches(:projects, :auto_cancel_pending_pipelines, 1)
     end
+  end
 
   def down
     # Nothing we can do!

db/post_migrate/20170503004427_update_retried_for_ci_build.rb

@@ -7,14 +7,14 @@ class UpdateRetriedForCiBuild < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
     if Gitlab::Database.mysql?
       up_mysql
     else
+      disable_statement_timeout do
         up_postgres
       end
     end
+  end
 
   def down
   end

db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb

@@ -7,11 +7,10 @@ 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)
 
+    disable_statement_timeout do
       head_id = pipelines
         .project(Arel::Nodes::NamedFunction.new('max', [pipelines[:id]]))
         .from(pipelines)
@@ -22,6 +21,7 @@ class AddHeadPipelineForEachMergeRequest < ActiveRecord::Migration
 
       update_column_in_batches(:merge_requests, :head_pipeline_id, sub_query)
     end
+  end
 
   def down
   end

db/post_migrate/20170525140254_rename_all_reserved_paths_again.rb

@@ -87,16 +87,16 @@ class RenameAllReservedPathsAgain < ActiveRecord::Migration
   ].freeze
 
   def up
-    disable_statement_timeout
-
+    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
-
+    disable_statement_timeout do
       revert_renames
     end
+  end
 end

db/post_migrate/20170526185842_migrate_pipeline_stages.rb

@@ -6,8 +6,7 @@ class MigratePipelineStages < ActiveRecord::Migration
   disable_ddl_transaction!
 
   def up
-    disable_statement_timeout
-
+    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
@@ -19,4 +18,5 @@ class MigratePipelineStages < ActiveRecord::Migration
             ORDER BY MAX(stage_idx)
       SQL
     end
+  end
 end

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
 
+    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
-
+    disable_statement_timeout do
       update_column_in_batches(:ci_builds, :stage_id, nil)
     end
+  end
 end

db/post_migrate/20170711145558_migrate_stages_statuses.rb

@@ -26,9 +26,9 @@ class MigrateStagesStatuses < ActiveRecord::Migration
   end
 
   def down
-    disable_statement_timeout
-
+    disable_statement_timeout do
       # rubocop:disable Migration/UpdateLargeTable
       update_column_in_batches(:ci_stages, :status, nil)
     end
+  end
 end

db/post_migrate/20171207150343_remove_soft_removed_objects.rb

@@ -78,13 +78,13 @@ class RemoveSoftRemovedObjects < ActiveRecord::Migration
   MODELS = [Issue, MergeRequest, CiPipelineSchedule, CiTrigger].freeze
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       remove_personal_routes
       remove_personal_namespaces
       remove_group_namespaces
       remove_simple_soft_removed_rows
     end
+  end
 
   def down
     # The data removed by this migration can't be restored in an automated way.

db/post_migrate/20180119121225_remove_redundant_pipeline_stages.rb

@@ -38,8 +38,7 @@ class RemoveRedundantPipelineStages < ActiveRecord::Migration
   end
 
   def remove_redundant_pipeline_stages!
-    disable_statement_timeout
-
+    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
@@ -63,4 +62,5 @@ class RemoveRedundantPipelineStages < ActiveRecord::Migration
         SQL
       end
     end
+  end
 end

db/post_migrate/20180305100050_remove_permanent_from_redirect_routes.rb

@@ -15,11 +15,11 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     # ReworkRedirectRoutesIndexes:
     # https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/16211
     if Gitlab::Database.postgresql?
-      disable_statement_timeout
-
+      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)
   end
@@ -28,10 +28,10 @@ class RemovePermanentFromRedirectRoutes < ActiveRecord::Migration
     add_column(:redirect_routes, :permanent, :boolean)
 
     if Gitlab::Database.postgresql?
-      disable_statement_timeout
-
+      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

db/post_migrate/20180306164012_add_path_index_to_redirect_routes.rb

@@ -20,12 +20,12 @@ class AddPathIndexToRedirectRoutes < ActiveRecord::Migration
   def up
     return unless Gitlab::Database.postgresql?
 
-    disable_statement_timeout
-
+    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
 
   def down
     # Do nothing in the DOWN. Since the index above is originally created in the

db/post_migrate/20180405101928_reschedule_builds_stages_migration.rb

@@ -17,8 +17,7 @@ class RescheduleBuildsStagesMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       Build.where('stage_id IS NULL').tap do |relation|
         queue_background_migration_jobs_by_range_at_intervals(relation,
                                                               MIGRATION,
@@ -26,6 +25,7 @@ class RescheduleBuildsStagesMigration < ActiveRecord::Migration
                                                               batch_size: BATCH_SIZE)
       end
     end
+  end
 
   def down
     # noop

db/post_migrate/20180420080616_schedule_stages_index_migration.rb

@@ -13,8 +13,7 @@ class ScheduleStagesIndexMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       Stage.all.tap do |relation|
         queue_background_migration_jobs_by_range_at_intervals(relation,
                                                               MIGRATION,
@@ -22,6 +21,7 @@ class ScheduleStagesIndexMigration < ActiveRecord::Migration
                                                               batch_size: BATCH_SIZE)
       end
     end
+  end
 
   def down
     # noop

db/post_migrate/20180604123514_cleanup_stages_position_migration.rb

@@ -12,8 +12,7 @@ class CleanupStagesPositionMigration < ActiveRecord::Migration
   end
 
   def up
-    disable_statement_timeout
-
+    disable_statement_timeout do
       Gitlab::BackgroundMigration.steal('MigrateStageIndex')
 
       unless index_exists_by_name?(:ci_stages, TMP_INDEX_NAME)
@@ -34,10 +33,13 @@ class CleanupStagesPositionMigration < ActiveRecord::Migration
 
       remove_concurrent_index_by_name(:ci_stages, TMP_INDEX_NAME)
     end
+  end
 
   def down
     if index_exists_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

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,8 +65,10 @@ module Gitlab
           return
         end
 
+        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,8 +95,10 @@ module Gitlab
           return
         end
 
+        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,8 +125,10 @@ module Gitlab
           return
         end
 
+        disable_statement_timeout do
           remove_index(table_name, options.merge({ name: index_name }))
         end
+      end
 
       # Only available on Postgresql >= 9.2
       def supports_drop_index_concurrently?
@@ -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,8 +200,10 @@ module Gitlab
         # while running.
         #
         # Note this is a no-op in case the constraint is VALID already
+        disable_statement_timeout do
           execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
         end
+      end
 
       def foreign_key_exists?(source, target = nil, column: nil)
         foreign_keys(source).any? do |key|
@@ -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,8 +410,7 @@ module Gitlab
             'in the body of your migration class'
         end
 
-        disable_statement_timeout
-
+        disable_statement_timeout do
           transaction do
             if limit
               add_column(table, column, type, default: nil, limit: limit)
@@ -393,6 +435,7 @@ module Gitlab
             raise error
           end
         end
+      end
 
       # Renames a column without requiring downtime.
       #

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