125 lines
3.9 KiB
Ruby
125 lines
3.9 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
namespace :gitlab do
|
|
namespace :db do
|
|
TRIGGER_FUNCTION_NAME = 'gitlab_schema_prevent_write'
|
|
|
|
desc "GitLab | DB | Install prevent write triggers on all databases"
|
|
task lock_writes: [:environment, 'gitlab:db:validate_config'] do
|
|
Gitlab::Database::EachDatabase.each_database_connection do |connection, database_name|
|
|
create_write_trigger_function(connection)
|
|
|
|
schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
|
|
Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
|
|
# TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
|
|
next if schema_name == :gitlab_geo
|
|
|
|
if schemas_for_connection.include?(schema_name.to_sym)
|
|
drop_write_trigger(database_name, connection, table_name)
|
|
else
|
|
create_write_trigger(database_name, connection, table_name)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
desc "GitLab | DB | Remove all triggers that prevents writes from all databases"
|
|
task unlock_writes: :environment do
|
|
Gitlab::Database::EachDatabase.each_database_connection do |connection, database_name|
|
|
Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
|
|
# TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
|
|
next if schema_name == :gitlab_geo
|
|
|
|
drop_write_trigger(database_name, connection, table_name)
|
|
end
|
|
drop_write_trigger_function(connection)
|
|
end
|
|
end
|
|
|
|
def create_write_trigger_function(connection)
|
|
sql = <<-SQL
|
|
CREATE OR REPLACE FUNCTION #{TRIGGER_FUNCTION_NAME}()
|
|
RETURNS TRIGGER AS
|
|
$$
|
|
BEGIN
|
|
RAISE EXCEPTION 'Table: "%" is write protected within this Gitlab database.', TG_TABLE_NAME
|
|
USING ERRCODE = 'modifying_sql_data_not_permitted',
|
|
HINT = 'Make sure you are using the right database connection';
|
|
END
|
|
$$ LANGUAGE PLPGSQL
|
|
SQL
|
|
|
|
connection.execute(sql)
|
|
end
|
|
|
|
def drop_write_trigger_function(connection)
|
|
sql = <<-SQL
|
|
DROP FUNCTION IF EXISTS #{TRIGGER_FUNCTION_NAME}()
|
|
SQL
|
|
|
|
connection.execute(sql)
|
|
end
|
|
|
|
def create_write_trigger(database_name, connection, table_name)
|
|
puts "#{database_name}: '#{table_name}'... Lock Writes".color(:yellow)
|
|
sql = <<-SQL
|
|
DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name};
|
|
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
|
|
end
|
|
|
|
def drop_write_trigger(database_name, connection, table_name)
|
|
puts "#{database_name}: '#{table_name}'... Allow Writes".color(:green)
|
|
sql = <<-SQL
|
|
DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name}
|
|
SQL
|
|
|
|
with_retries(connection) do
|
|
connection.execute(sql)
|
|
end
|
|
end
|
|
|
|
def with_retries(connection, &block)
|
|
with_statement_timeout_retries do
|
|
with_lock_retries(connection) do
|
|
yield
|
|
end
|
|
end
|
|
end
|
|
|
|
def with_statement_timeout_retries(times = 5)
|
|
current_iteration = 1
|
|
begin
|
|
yield
|
|
rescue ActiveRecord::QueryCanceled => err
|
|
puts "Retrying after #{err.message}"
|
|
|
|
if current_iteration <= times
|
|
current_iteration += 1
|
|
retry
|
|
else
|
|
raise err
|
|
end
|
|
end
|
|
end
|
|
|
|
def with_lock_retries(connection, &block)
|
|
Gitlab::Database::WithLockRetries.new(
|
|
klass: "gitlab:db:lock_writes",
|
|
logger: Gitlab::AppLogger,
|
|
connection: connection
|
|
).run(&block)
|
|
end
|
|
|
|
def write_trigger_name(table_name)
|
|
"gitlab_schema_write_trigger_for_#{table_name}"
|
|
end
|
|
end
|
|
end
|