Fix connection search path being reset

RESET ALL is reseting too many parameters, we should reset only the
parameter we just set, namely statement_timeout

https://gitlab.com/gitlab-org/gitlab/-/issues/337813

Changelog: fixed

lib/gitlab/database/migration_helpers.rb

@@ -341,9 +341,9 @@ module Gitlab
       # - 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
+      # a block so we can automatically execute `RESET statement_timeout` after block finishes
       # otherwise the statement will still be disabled until connection is dropped
-      # or `RESET ALL` is executed
+      # or `RESET statement_timeout` is executed
       def disable_statement_timeout
         if block_given?
           if statement_timeout_disabled?
@@ -357,7 +357,7 @@ module Gitlab
 
               yield
             ensure
-              execute('RESET ALL')
+              execute('RESET statement_timeout')
             end
           end
         else

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -414,9 +414,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           expect(model).to receive(:execute).with(/REFERENCES users \(id\)/)
 
@@ -428,9 +428,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           expect(model).to receive(:execute).with(/REFERENCES users \(id_convert_to_bigint\)/)
 
@@ -446,9 +446,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
             expect(model).to receive(:with_lock_retries).and_call_original
             expect(model).to receive(:disable_statement_timeout).and_call_original
             expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-            expect(model).to receive(:execute).with(/statement_timeout/)
+            expect(model).to receive(:execute).with(/SET statement_timeout TO/)
             expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-            expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+            expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
             expect(model).to receive(:execute).with(/ON DELETE SET NULL/)
 
@@ -463,9 +463,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
             expect(model).to receive(:with_lock_retries).and_call_original
             expect(model).to receive(:disable_statement_timeout).and_call_original
             expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-            expect(model).to receive(:execute).with(/statement_timeout/)
+            expect(model).to receive(:execute).with(/SET statement_timeout TO/)
             expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-            expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+            expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
             expect(model).to receive(:execute).with(/ON DELETE CASCADE/)
 
@@ -480,9 +480,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
             expect(model).to receive(:with_lock_retries).and_call_original
             expect(model).to receive(:disable_statement_timeout).and_call_original
             expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-            expect(model).to receive(:execute).with(/statement_timeout/)
+            expect(model).to receive(:execute).with(/SET statement_timeout TO/)
             expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-            expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+            expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
             expect(model).not_to receive(:execute).with(/ON DELETE/)
 
@@ -498,10 +498,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/NOT VALID/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
         end
@@ -527,10 +527,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
             expect(model).to receive(:with_lock_retries).and_call_original
             expect(model).to receive(:disable_statement_timeout).and_call_original
             expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-            expect(model).to receive(:execute).with(/statement_timeout/)
+            expect(model).to receive(:execute).with(/SET statement_timeout TO/)
             expect(model).to receive(:execute).ordered.with(/NOT VALID/)
             expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT.+foo/)
-            expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+            expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
             model.add_concurrent_foreign_key(:projects, :users, column: :user_id, name: :foo)
           end
@@ -557,10 +557,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
               expect(model).to receive(:with_lock_retries).and_call_original
               expect(model).to receive(:disable_statement_timeout).and_call_original
               expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-              expect(model).to receive(:execute).with(/statement_timeout/)
+              expect(model).to receive(:execute).with(/SET statement_timeout TO/)
               expect(model).to receive(:execute).ordered.with(/NOT VALID/)
               expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT.+bar/)
-              expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+              expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
               model.add_concurrent_foreign_key(:projects, :users, column: :user_id, name: :bar)
             end
@@ -592,9 +592,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           expect(model).to receive(:execute).with('LOCK TABLE users, projects IN SHARE ROW EXCLUSIVE MODE')
           expect(model).to receive(:execute).with(/REFERENCES users \(id\)/)
@@ -614,9 +614,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).not_to receive(:concurrent_foreign_key_name)
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/ALTER TABLE projects VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
         end
 
         model.validate_foreign_key(:projects, :user_id, name: :foo)
@@ -631,9 +631,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:concurrent_foreign_key_name)
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:execute).ordered.with(/ALTER TABLE projects VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
         end
 
         model.validate_foreign_key(:projects, :user_id)
@@ -748,7 +748,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       end
 
       after do
-        model.execute('RESET ALL')
+        model.execute('RESET statement_timeout')
       end
 
       it 'defines statement to 0 only for current transaction' do
@@ -765,7 +765,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       context 'when passing a blocks' do
         it 'disables statement timeouts on session level and executes the block' do
           expect(model).to receive(:execute).with('SET statement_timeout TO 0')
-          expect(model).to receive(:execute).with('RESET ALL').at_least(:once)
+          expect(model).to receive(:execute).with('RESET statement_timeout').at_least(:once)
 
           expect { |block| model.disable_statement_timeout(&block) }.to yield_control
         end
@@ -777,7 +777,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           end
 
           after do
-            model.execute('RESET ALL')
+            model.execute('RESET statement_timeout')
           end
 
           it 'defines statement to 0 for any code run inside the block' do
@@ -804,12 +804,12 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       after do
         # Use ActiveRecord::Base.connection instead of model.execute
         # so that this call is not counted below
-        ActiveRecord::Base.connection.execute('RESET ALL')
+        ActiveRecord::Base.connection.execute('RESET statement_timeout')
       end
 
       it 'yields control without disabling the timeout or resetting' do
         expect(model).not_to receive(:execute).with('SET statement_timeout TO 0')
-        expect(model).not_to receive(:execute).with('RESET ALL')
+        expect(model).not_to receive(:execute).with('RESET statement_timeout')
 
         expect { |block| model.disable_statement_timeout(&block) }.to yield_control
       end
@@ -2532,7 +2532,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
 
@@ -2542,7 +2542,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
                        .and_return(true).exactly(1)
 
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           model.add_check_constraint(
             :test_table,
@@ -2576,7 +2576,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
           expect(model).to receive(:disable_statement_timeout).and_call_original
           expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).with(/SET statement_timeout TO/)
           expect(model).to receive(:with_lock_retries).and_call_original
           expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
 
@@ -2585,7 +2585,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
                        .and_return(true).exactly(1)
 
           expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+          expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
           model.add_check_constraint(
             :test_table,
@@ -2618,9 +2618,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         expect(model).to receive(:check_constraint_exists?).and_return(true)
         expect(model).to receive(:disable_statement_timeout).and_call_original
         expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-        expect(model).to receive(:execute).with(/statement_timeout/)
+        expect(model).to receive(:execute).with(/SET statement_timeout TO/)
         expect(model).to receive(:execute).ordered.with(validate_sql)
-        expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+        expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
         model.validate_check_constraint(:test_table, 'check_name')
       end

spec/support/shared_examples/lib/gitlab/migration_helpers_shared_examples.rb

@@ -14,10 +14,10 @@ RSpec.shared_examples 'performs validation' do |validation_option|
   it 'performs validation' do
     expect(model).to receive(:disable_statement_timeout).and_call_original
     expect(model).to receive(:statement_timeout_disabled?).and_return(false)
-    expect(model).to receive(:execute).with(/statement_timeout/)
+    expect(model).to receive(:execute).with(/SET statement_timeout TO/)
     expect(model).to receive(:execute).ordered.with(/NOT VALID/)
     expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
-    expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+    expect(model).to receive(:execute).ordered.with(/RESET statement_timeout/)
 
     model.add_concurrent_foreign_key(*args, **options.merge(validation_option))
   end