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.

1 files changed, 93 insertions, 13 deletions

diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index eb7148ff108..c993fa12803 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/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