diff --git a/lib/gitlab/database/lock_writes_manager.rb b/lib/gitlab/database/lock_writes_manager.rb
index cd483d616bb4489fe448c7404df585f6c02fdad9..fe75cd763b49828e02441a2b9778f7070d409dd2 100644
--- a/lib/gitlab/database/lock_writes_manager.rb
+++ b/lib/gitlab/database/lock_writes_manager.rb
@@ -5,42 +5,63 @@ module Database
class LockWritesManager
TRIGGER_FUNCTION_NAME = 'gitlab_schema_prevent_write'
- def initialize(table_name:, connection:, database_name:, logger: nil)
+ def initialize(table_name:, connection:, database_name:, logger: nil, dry_run: false)
@table_name = table_name
@connection = connection
@database_name = database_name
@logger = logger
+ @dry_run = dry_run
+ end
+
+ def table_locked_for_writes?(table_name)
+ query = <<~SQL
+ SELECT COUNT(*) from information_schema.triggers
+ WHERE event_object_table = '#{table_name}'
+ AND trigger_name = '#{write_trigger_name(table_name)}'
+ SQL
+
+ connection.select_value(query) == 3
end
def lock_writes
+ if table_locked_for_writes?(table_name)
+ logger&.info "Skipping lock_writes, because #{table_name} is already locked for writes"
+ return
+ end
+
logger&.info "Database: '#{database_name}', Table: '#{table_name}': Lock Writes".color(:yellow)
- sql = <<-SQL
- DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name};
+ sql_statement = <<~SQL
CREATE TRIGGER #{write_trigger_name(table_name)}
BEFORE INSERT OR UPDATE OR DELETE OR TRUNCATE
ON #{table_name}
FOR EACH STATEMENT EXECUTE FUNCTION #{TRIGGER_FUNCTION_NAME}();
SQL
- with_retries(connection) do
- connection.execute(sql)
- end
+ execute_sql_statement(sql_statement)
end
def unlock_writes
logger&.info "Database: '#{database_name}', Table: '#{table_name}': Allow Writes".color(:green)
- sql = <<-SQL
- DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name}
+ sql_statement = <<~SQL
+ DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name};
SQL
- with_retries(connection) do
- connection.execute(sql)
- end
+ execute_sql_statement(sql_statement)
end
private
- attr_reader :table_name, :connection, :database_name, :logger
+ attr_reader :table_name, :connection, :database_name, :logger, :dry_run
+
+ def execute_sql_statement(sql)
+ if dry_run
+ logger&.info sql
+ else
+ with_retries(connection) do
+ connection.execute(sql)
+ end
+ end
+ end
def with_retries(connection, &block)
with_statement_timeout_retries do
diff --git a/spec/lib/gitlab/database/lock_writes_manager_spec.rb b/spec/lib/gitlab/database/lock_writes_manager_spec.rb
index eb527d492cf00d8b3064f0e1f100a0b05cb37a59..b1cc8add55a750d0eb0a33cc86261839ee83cee9 100644
--- a/spec/lib/gitlab/database/lock_writes_manager_spec.rb
+++ b/spec/lib/gitlab/database/lock_writes_manager_spec.rb
@@ -6,13 +6,15 @@
let(:connection) { ApplicationRecord.connection }
let(:test_table) { '_test_table' }
let(:logger) { instance_double(Logger) }
+ let(:dry_run) { false }
subject(:lock_writes_manager) do
described_class.new(
table_name: test_table,
connection: connection,
database_name: 'main',
- logger: logger
+ logger: logger,
+ dry_run: dry_run
)
end
@@ -27,6 +29,16 @@
SQL
end
+ describe "#table_locked_for_writes?" do
+ it 'returns false for a table that is not locked for writes' do
+ expect(subject.table_locked_for_writes?(test_table)).to eq(false)
+ end
+
+ it 'returns true for a table that is locked for writes' do
+ expect { subject.lock_writes }.to change { subject.table_locked_for_writes?(test_table) }.from(false).to(true)
+ end
+ end
+
describe '#lock_writes' do
it 'prevents any writes on the table' do
subject.lock_writes
@@ -84,11 +96,57 @@
subject.lock_writes
end.to raise_error(ActiveRecord::QueryCanceled)
end
+
+ it 'skips the operation if the table is already locked for writes' do
+ subject.lock_writes
+
+ expect(logger).to receive(:info).with("Skipping lock_writes, because #{test_table} is already locked for writes")
+ expect(connection).not_to receive(:execute).with(/CREATE TRIGGER/)
+
+ expect do
+ subject.lock_writes
+ end.not_to change {
+ number_of_triggers_on(connection, test_table)
+ }
+ end
+
+ context 'when running in dry_run mode' do
+ let(:dry_run) { true }
+
+ it 'prints the sql statement to the logger' do
+ expect(logger).to receive(:info).with("Database: 'main', Table: '#{test_table}': Lock Writes")
+ expected_sql_statement = <<~SQL
+ CREATE TRIGGER gitlab_schema_write_trigger_for_#{test_table}
+ BEFORE INSERT OR UPDATE OR DELETE OR TRUNCATE
+ ON #{test_table}
+ FOR EACH STATEMENT EXECUTE FUNCTION gitlab_schema_prevent_write();
+ SQL
+ expect(logger).to receive(:info).with(expected_sql_statement)
+
+ subject.lock_writes
+ end
+
+ it 'does not lock the tables for writes' do
+ subject.lock_writes
+
+ expect do
+ connection.execute("delete from #{test_table}")
+ connection.execute("truncate #{test_table}")
+ end.not_to raise_error
+ end
+ end
end
describe '#unlock_writes' do
before do
- subject.lock_writes
+ # Locking the table without the considering the value of dry_run
+ described_class.new(
+ table_name: test_table,
+ connection: connection,
+ database_name: 'main',
+ logger: logger,
+ dry_run: false
+ ).lock_writes
end
it 'allows writing on the table again' do
@@ -114,6 +172,28 @@
subject.unlock_writes
end
+
+ context 'when running in dry_run mode' do
+ let(:dry_run) { true }
+
+ it 'prints the sql statement to the logger' do
+ expect(logger).to receive(:info).with("Database: 'main', Table: '#{test_table}': Allow Writes")
+ expected_sql_statement = <<~SQL
+ DROP TRIGGER IF EXISTS gitlab_schema_write_trigger_for_#{test_table} ON #{test_table};
+ SQL
+ expect(logger).to receive(:info).with(expected_sql_statement)
+
+ subject.unlock_writes
+ end
+
+ it 'does not unlock the tables for writes' do
+ subject.unlock_writes
+
+ expect do
+ connection.execute("delete from #{test_table}")
+ end.to raise_error(ActiveRecord::StatementInvalid, /Table: "#{test_table}" is write protected/)
+ end
+ end
end
def number_of_triggers_on(connection, table_name)