Commit cd71cf54 authored by dfrazao-gitlab's avatar dfrazao-gitlab

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
parent 809fe53f
......@@ -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.
......
......@@ -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
......
......@@ -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
......
......@@ -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
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment