Support multiple dbs in MigrationHelpers

- Remove hardcoded ActiveRecord::Base.connection for WithLockRetries
- Pass connection as argument for WithLockRetries
- Start to use migration connection for triggers

Related to https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67753

Changelog: added

lib/gitlab/database/migration_helpers.rb

@@ -603,17 +603,17 @@ module Gitlab
       # new_column - The name of the new column.
       # trigger_name - The name of the trigger to use (optional).
       def install_rename_triggers(table, old, new, trigger_name: nil)
-        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).create(old, new, trigger_name: trigger_name)
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table, connection: connection).create(old, new, trigger_name: trigger_name)
       end
 
       # Removes the triggers used for renaming a column concurrently.
       def remove_rename_triggers(table, trigger)
-        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).drop(trigger)
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table, connection: connection).drop(trigger)
       end
 
       # Returns the (base) name to use for triggers when renaming columns.
       def rename_trigger_name(table, old, new)
-        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).name(old, new)
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table, connection: connection).name(old, new)
       end
 
       # Changes the type of a column concurrently.

lib/gitlab/database/with_lock_retries.rb

@@ -61,7 +61,7 @@ module Gitlab
         [10.seconds, 10.minutes]
       ].freeze
 
-      def initialize(logger: NULL_LOGGER, allow_savepoints: true, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV)
+      def initialize(logger: NULL_LOGGER, allow_savepoints: true, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV, connection: ActiveRecord::Base.connection)
         @logger = logger
         @klass = klass
         @allow_savepoints = allow_savepoints
@@ -69,6 +69,7 @@ module Gitlab
         @env = env
         @current_iteration = 1
         @log_params = { method: 'with_lock_retries', class: klass.to_s }
+        @connection = connection
       end
 
       # Executes a block of code, retrying it whenever a database lock can't be acquired in time
@@ -96,7 +97,7 @@ module Gitlab
           run_block_with_lock_timeout
         rescue ActiveRecord::LockWaitTimeout
           if retry_with_lock_timeout?
-            disable_idle_in_transaction_timeout if ActiveRecord::Base.connection.transaction_open?
+            disable_idle_in_transaction_timeout if connection.transaction_open?
             wait_until_next_retry
             reset_db_settings
 
@@ -116,16 +117,16 @@ module Gitlab
 
       private
 
-      attr_reader :logger, :env, :block, :current_iteration, :log_params, :timing_configuration
+      attr_reader :logger, :env, :block, :current_iteration, :log_params, :timing_configuration, :connection
 
       def run_block
         block.call
       end
 
       def run_block_with_lock_timeout
-        raise "WithLockRetries should not run inside already open transaction" if ActiveRecord::Base.connection.transaction_open? && @allow_savepoints.blank?
+        raise "WithLockRetries should not run inside already open transaction" if connection.transaction_open? && @allow_savepoints.blank?
 
-        ActiveRecord::Base.transaction(requires_new: true) do # rubocop:disable Performance/ActiveRecordSubtransactions
+        connection.transaction(requires_new: true) do # rubocop:disable Performance/ActiveRecordSubtransactions
           execute("SET LOCAL lock_timeout TO '#{current_lock_timeout_in_ms}ms'")
 
           log(message: 'Lock timeout is set', current_iteration: current_iteration, lock_timeout_in_ms: current_lock_timeout_in_ms)
@@ -152,7 +153,7 @@ module Gitlab
         log(message: "Couldn't acquire lock to perform the migration", current_iteration: current_iteration)
         log(message: "Executing the migration without lock timeout", current_iteration: current_iteration)
 
-        disable_lock_timeout if ActiveRecord::Base.connection.transaction_open?
+        disable_lock_timeout if connection.transaction_open?
 
         run_block
 
@@ -168,7 +169,7 @@ module Gitlab
       end
 
       def execute(statement)
-        ActiveRecord::Base.connection.execute(statement)
+        connection.execute(statement)
       end
 
       def retry_count

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -798,13 +798,13 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     # This spec runs without an enclosing transaction (:delete truncation method for db_cleaner)
     context 'when the statement_timeout is already disabled', :delete do
       before do
-        ActiveRecord::Base.connection.execute('SET statement_timeout TO 0')
+        ActiveRecord::Migration.connection.execute('SET statement_timeout TO 0')
       end
 
       after do
-        # Use ActiveRecord::Base.connection instead of model.execute
+        # Use ActiveRecord::Migration.connection instead of model.execute
         # so that this call is not counted below
-        ActiveRecord::Base.connection.execute('RESET statement_timeout')
+        ActiveRecord::Migration.connection.execute('RESET statement_timeout')
       end
 
       it 'yields control without disabling the timeout or resetting' do
@@ -954,10 +954,11 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           let(:trigger_name) { model.rename_trigger_name(:users, :id, :new) }
           let(:user) { create(:user) }
           let(:copy_trigger) { double('copy trigger') }
+          let(:connection) { ActiveRecord::Migration.connection }
 
           before do
             expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
-              .with(:users).and_return(copy_trigger)
+              .with(:users, connection: connection).and_return(copy_trigger)
           end
 
           it 'copies the value to the new column using the type_cast_function', :aggregate_failures do
@@ -1300,11 +1301,13 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   end
 
   describe '#install_rename_triggers' do
+    let(:connection) { ActiveRecord::Migration.connection }
+
     it 'installs the triggers' do
       copy_trigger = double('copy trigger')
 
       expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
-        .with(:users).and_return(copy_trigger)
+        .with(:users, connection: connection).and_return(copy_trigger)
 
       expect(copy_trigger).to receive(:create).with(:old, :new, trigger_name: 'foo')
 
@@ -1313,11 +1316,13 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   end
 
   describe '#remove_rename_triggers' do
+    let(:connection) { ActiveRecord::Migration.connection }
+
     it 'removes the function and trigger' do
       copy_trigger = double('copy trigger')
 
       expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
-        .with('bar').and_return(copy_trigger)
+        .with('bar', connection: connection).and_return(copy_trigger)
 
       expect(copy_trigger).to receive(:drop).with('foo')
 
@@ -2194,7 +2199,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
   describe '#index_exists_by_name?' do
     it 'returns true if an index exists' do
-      ActiveRecord::Base.connection.execute(
+      ActiveRecord::Migration.connection.execute(
         'CREATE INDEX test_index_for_index_exists ON projects (path);'
       )
 
@@ -2209,7 +2214,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
     context 'when an index with a function exists' do
       before do
-        ActiveRecord::Base.connection.execute(
+        ActiveRecord::Migration.connection.execute(
           'CREATE INDEX test_index ON projects (LOWER(path));'
         )
       end
@@ -2222,15 +2227,15 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
     context 'when an index exists for a table with the same name in another schema' do
       before do
-        ActiveRecord::Base.connection.execute(
+        ActiveRecord::Migration.connection.execute(
           'CREATE SCHEMA new_test_schema'
         )
 
-        ActiveRecord::Base.connection.execute(
+        ActiveRecord::Migration.connection.execute(
           'CREATE TABLE new_test_schema.projects (id integer, name character varying)'
         )
 
-        ActiveRecord::Base.connection.execute(
+        ActiveRecord::Migration.connection.execute(
           'CREATE INDEX test_index_on_name ON new_test_schema.projects (LOWER(name));'
         )
       end
@@ -2463,19 +2468,19 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
   describe '#check_constraint_exists?' do
     before do
-      ActiveRecord::Base.connection.execute(
+      ActiveRecord::Migration.connection.execute(
         'ALTER TABLE projects ADD CONSTRAINT check_1 CHECK (char_length(path) <= 5) NOT VALID'
       )
 
-      ActiveRecord::Base.connection.execute(
+      ActiveRecord::Migration.connection.execute(
         'CREATE SCHEMA new_test_schema'
       )
 
-      ActiveRecord::Base.connection.execute(
+      ActiveRecord::Migration.connection.execute(
         'CREATE TABLE new_test_schema.projects (id integer, name character varying)'
       )
 
-      ActiveRecord::Base.connection.execute(
+      ActiveRecord::Migration.connection.execute(
         'ALTER TABLE new_test_schema.projects ADD CONSTRAINT check_2 CHECK (char_length(name) <= 5)'
       )
     end

spec/lib/gitlab/database/with_lock_retries_spec.rb

@@ -7,6 +7,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
   let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
   let(:allow_savepoints) { true }
+  let(:connection) { ActiveRecord::Base.connection }
 
   let(:timing_configuration) do
     [
@@ -67,7 +68,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
             WHERE t.relkind = 'r' AND l.mode = 'ExclusiveLock' AND t.relname = '#{Project.table_name}'
           """
 
-          expect(ActiveRecord::Base.connection.execute(check_exclusive_lock_query).to_a).to be_present
+          expect(connection.execute(check_exclusive_lock_query).to_a).to be_present
         end
       end
 
@@ -96,8 +97,8 @@ RSpec.describe Gitlab::Database::WithLockRetries do
               lock_fiber.resume
             end
 
-            ActiveRecord::Base.transaction do
-              ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+            connection.transaction do
+              connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
               lock_acquired = true
             end
           end
@@ -115,7 +116,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
         context 'setting the idle transaction timeout' do
           context 'when there is no outer transaction: disable_ddl_transaction! is set in the migration' do
             it 'does not disable the idle transaction timeout' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+              allow(connection).to receive(:transaction_open?).and_return(false)
               allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
               allow(subject).to receive(:run_block_with_lock_timeout).once
 
@@ -127,7 +128,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
 
           context 'when there is outer transaction: disable_ddl_transaction! is not set in the migration' do
             it 'disables the idle transaction timeout so the code can sleep and retry' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+              allow(connection).to receive(:transaction_open?).and_return(true)
 
               n = 0
               allow(subject).to receive(:run_block_with_lock_timeout).twice do
@@ -152,7 +153,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
 
         context 'when there is no outer transaction: disable_ddl_transaction! is set in the migration' do
           it 'does not disable the lock_timeout' do
-            allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+            allow(connection).to receive(:transaction_open?).and_return(false)
             allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
 
             expect(subject).not_to receive(:disable_lock_timeout)
@@ -163,7 +164,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
 
         context 'when there is outer transaction: disable_ddl_transaction! is not set in the migration' do
           it 'disables the lock_timeout' do
-            allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+            allow(connection).to receive(:transaction_open?).and_return(true)
             allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
 
             expect(subject).to receive(:disable_lock_timeout)
@@ -198,8 +199,8 @@ RSpec.describe Gitlab::Database::WithLockRetries do
             subject.run(raise_on_exhaustion: true) do
               lock_attempts += 1
 
-              ActiveRecord::Base.transaction do
-                ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+              connection.transaction do
+                connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
                 lock_acquired = true
               end
             end
@@ -213,11 +214,11 @@ RSpec.describe Gitlab::Database::WithLockRetries do
       context 'when statement timeout is reached' do
         it 'raises QueryCanceled error' do
           lock_acquired = false
-          ActiveRecord::Base.connection.execute("SET LOCAL statement_timeout='100ms'")
+          connection.execute("SET LOCAL statement_timeout='100ms'")
 
           expect do
             subject.run do
-              ActiveRecord::Base.connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
+              connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
               lock_acquired = true
             end
           end.to raise_error(ActiveRecord::QueryCanceled)
@@ -230,11 +231,11 @@ RSpec.describe Gitlab::Database::WithLockRetries do
 
   context 'restore local database variables' do
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW lock_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW lock_timeout").to_a }
     end
 
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
     end
   end
 
@@ -242,10 +243,10 @@ RSpec.describe Gitlab::Database::WithLockRetries do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET LOCAL lock_timeout` using the configured timeout value in milliseconds' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("RESET idle_in_transaction_session_timeout; RESET lock_timeout").and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1", "TRANSACTION").and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("SET LOCAL lock_timeout TO '15ms'").and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1", "TRANSACTION").and_call_original
+      expect(connection).to receive(:execute).with("RESET idle_in_transaction_session_timeout; RESET lock_timeout").and_call_original
+      expect(connection).to receive(:execute).with("SAVEPOINT active_record_1", "TRANSACTION").and_call_original
+      expect(connection).to receive(:execute).with("SET LOCAL lock_timeout TO '15ms'").and_call_original
+      expect(connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1", "TRANSACTION").and_call_original
 
       subject.run { }
     end
@@ -262,13 +263,13 @@ RSpec.describe Gitlab::Database::WithLockRetries do
     let(:allow_savepoints) { false }
 
     it 'prevents running inside already open transaction' do
-      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+      allow(connection).to receive(:transaction_open?).and_return(true)
 
       expect { subject.run { } }.to raise_error(/should not run inside already open transaction/)
     end
 
     it 'does not raise the error if not inside open transaction' do
-      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+      allow(connection).to receive(:transaction_open?).and_return(false)
 
       expect { subject.run { } }.not_to raise_error
     end