MigrationHelper `disable_statement_timeout` accepts `transaction: false`

By default statement_timeout will only be enabled during transaction
lifetime, therefore not leaking outside of it.

With `transaction: false` it will set for entire session, but requires
a block to passed. It yields control and cleans up session after block
finishes, also preventing leaking outside of it.

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(transaction: false) 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(transaction: false) 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(transaction: false) 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(transaction: false) do
+          execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
+        end
       end
 
       def foreign_key_exists?(source, target = nil, column: nil)
@@ -224,8 +227,49 @@ 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)
-      def disable_statement_timeout
-        execute('SET statement_timeout TO 0') if Database.postgresql?
+      #
+      # 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 `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)
+        # 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 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
+
+          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.'
+          end
+
+          execute('SET statement_timeout TO 0')
+
+          yield
+
+          execute('RESET ALL')
+        end
       end
 
       def true_value
@@ -367,30 +411,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(transaction: false) 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
 

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -51,7 +51,7 @@ describe Gitlab::Database::MigrationHelpers do
       context 'using 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 })
@@ -228,17 +228,21 @@ describe Gitlab::Database::MigrationHelpers do
         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,22 +295,98 @@ describe Gitlab::Database::MigrationHelpers do
 
   describe '#disable_statement_timeout' do
     context 'using PostgreSQL' do
-      it 'disables statement timeouts' do
-        expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
+      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')
+
+          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
 
-        expect(model).to receive(:execute).with('SET statement_timeout TO 0')
+          it 'defines statement to 0 only for current transaction' do
+            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
 
-        model.disable_statement_timeout
+            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
+      end
+
+      context 'with transaction: false' do
+        it 'disables statement timeouts on session level' 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)
+        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 block' do
+            expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
+
+            model.disable_statement_timeout(transaction: false) 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
-      it 'does nothing' do
-        expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
+      context 'with transaction: true' do
+        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
+      end
 
-        model.disable_statement_timeout
+      context 'with transaction: false' do
+        it 'executes yielded 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
+        end
       end
     end
   end