Commit 96adae97 authored by Adam Hegyi's avatar Adam Hegyi

Merge branch '201872-partitioning-implement-cascading-deletes-without-foreign-keys-2' into 'master'

Migration helpers to manage FKs for partitioned tables

See merge request gitlab-org/gitlab!29510
parents b4b255f0 c9168736
# frozen_string_literal: true
module Gitlab
module Database
module PartitioningMigrationHelpers
include SchemaHelpers
def add_partitioned_foreign_key(from_table, to_table, column: nil, primary_key: :id, on_delete: :cascade)
cascade_delete = extract_cascade_option(on_delete)
update_foreign_keys(from_table, to_table, column, primary_key, cascade_delete) do |current_keys, existing_key, specified_key|
if existing_key.nil?
unless specified_key.save
raise "failed to create foreign key: #{specified_key.errors.full_messages.to_sentence}"
end
current_keys << specified_key
else
Rails.logger.warn "foreign key not added because it already exists: #{specified_key}" # rubocop:disable Gitlab/RailsLogger
current_keys
end
end
end
def remove_partitioned_foreign_key(from_table, to_table, column: nil, primary_key: :id)
update_foreign_keys(from_table, to_table, column, primary_key) do |current_keys, existing_key, specified_key|
if existing_key
existing_key.destroy!
current_keys.delete(existing_key)
else
Rails.logger.warn "foreign key not removed because it doesn't exist: #{specified_key}" # rubocop:disable Gitlab/RailsLogger
end
current_keys
end
end
def fk_function_name(table)
object_name(table, 'fk_cascade_function')
end
def fk_trigger_name(table)
object_name(table, 'fk_cascade_trigger')
end
private
def fk_from_spec(from_table, to_table, from_column, to_column, cascade_delete)
PartitionedForeignKey.new(from_table: from_table.to_s, to_table: to_table.to_s, from_column: from_column.to_s,
to_column: to_column.to_s, cascade_delete: cascade_delete)
end
def update_foreign_keys(from_table, to_table, from_column, to_column, cascade_delete = nil)
if transaction_open?
raise 'partitioned foreign key operations can not be run inside a transaction block, ' \
'you can disable transaction blocks by calling disable_ddl_transaction! ' \
'in the body of your migration class'
end
from_column ||= "#{to_table.to_s.singularize}_id"
specified_key = fk_from_spec(from_table, to_table, from_column, to_column, cascade_delete)
current_keys = PartitionedForeignKey.by_referenced_table(to_table).to_a
existing_key = find_existing_key(current_keys, specified_key)
final_keys = yield current_keys, existing_key, specified_key
fn_name = fk_function_name(to_table)
trigger_name = fk_trigger_name(to_table)
with_lock_retries do
drop_trigger(to_table, trigger_name, if_exists: true)
if final_keys.empty?
drop_function(fn_name, if_exists: true)
else
create_or_replace_fk_function(fn_name, final_keys)
create_function_trigger(trigger_name, fn_name, fires: "AFTER DELETE ON #{to_table}")
end
end
end
def extract_cascade_option(on_delete)
case on_delete
when :cascade then true
when :nullify then false
else raise ArgumentError, "invalid option #{on_delete} for :on_delete"
end
end
def with_lock_retries(&block)
Gitlab::Database::WithLockRetries.new({
klass: self.class,
logger: Gitlab::BackgroundMigration::Logger
}).run(&block)
end
def find_existing_key(keys, key)
keys.find { |k| k.from_table == key.from_table && k.from_column == key.from_column }
end
def create_or_replace_fk_function(fn_name, fk_specs)
create_trigger_function(fn_name, replace: true) do
cascade_statements = build_cascade_statements(fk_specs)
cascade_statements << 'RETURN OLD;'
cascade_statements.join("\n")
end
end
def build_cascade_statements(foreign_keys)
foreign_keys.map do |fks|
if fks.cascade_delete?
"DELETE FROM #{fks.from_table} WHERE #{fks.from_column} = OLD.#{fks.to_column};"
else
"UPDATE #{fks.from_table} SET #{fks.from_column} = NULL WHERE #{fks.from_column} = OLD.#{fks.to_column};"
end
end
end
end
end
end
# frozen_string_literal: true
module Gitlab
module Database
module PartitioningMigrationHelpers
class PartitionedForeignKey < ApplicationRecord
validates_with PartitionedForeignKeyValidator
scope :by_referenced_table, ->(table) { where(to_table: table) }
end
end
end
end
# frozen_string_literal: true
module Gitlab
module Database
module PartitioningMigrationHelpers
class PartitionedForeignKeyValidator < ActiveModel::Validator
def validate(record)
validate_key_part(record, :from_table, :from_column)
validate_key_part(record, :to_table, :to_column)
end
private
def validate_key_part(record, table_field, column_field)
if !connection.table_exists?(record[table_field])
record.errors.add(table_field, 'must be a valid table')
elsif !connection.column_exists?(record[table_field], record[column_field])
record.errors.add(column_field, 'must be a valid column')
end
end
def connection
ActiveRecord::Base.connection
end
end
end
end
end
# frozen_string_literal: true
module Gitlab
module Database
module SchemaHelpers
def create_trigger_function(name, replace: true)
replace_clause = optional_clause(replace, "OR REPLACE")
execute(<<~SQL)
CREATE #{replace_clause} FUNCTION #{name}()
RETURNS TRIGGER AS
$$
BEGIN
#{yield}
END
$$ LANGUAGE PLPGSQL
SQL
end
def create_function_trigger(name, fn_name, fires: nil)
execute(<<~SQL)
CREATE TRIGGER #{name}
#{fires}
FOR EACH ROW
EXECUTE PROCEDURE #{fn_name}()
SQL
end
def drop_function(name, if_exists: true)
exists_clause = optional_clause(if_exists, "IF EXISTS")
execute("DROP FUNCTION #{exists_clause} #{name}()")
end
def drop_trigger(table_name, name, if_exists: true)
exists_clause = optional_clause(if_exists, "IF EXISTS")
execute("DROP TRIGGER #{exists_clause} #{name} ON #{table_name}")
end
def object_name(table, type)
identifier = "#{table}_#{type}"
hashed_identifier = Digest::SHA256.hexdigest(identifier).first(10)
"#{type}_#{hashed_identifier}"
end
private
def optional_clause(flag, clause)
flag ? clause : ""
end
end
end
end
# frozen_string_literal: true
require 'spec_helper'
describe Gitlab::Database::PartitioningMigrationHelpers::PartitionedForeignKey do
let(:foreign_key) do
described_class.new(
to_table: 'issues',
from_table: 'issue_assignees',
from_column: 'issue_id',
to_column: 'id',
cascade_delete: true)
end
describe 'validations' do
it 'allows keys that reference valid tables and columns' do
expect(foreign_key).to be_valid
end
it 'does not allow keys without a valid to_table' do
foreign_key.to_table = 'this_is_not_a_real_table'
expect(foreign_key).not_to be_valid
expect(foreign_key.errors[:to_table].first).to eq('must be a valid table')
end
it 'does not allow keys without a valid from_table' do
foreign_key.from_table = 'this_is_not_a_real_table'
expect(foreign_key).not_to be_valid
expect(foreign_key.errors[:from_table].first).to eq('must be a valid table')
end
it 'does not allow keys without a valid to_column' do
foreign_key.to_column = 'this_is_not_a_real_fk'
expect(foreign_key).not_to be_valid
expect(foreign_key.errors[:to_column].first).to eq('must be a valid column')
end
it 'does not allow keys without a valid from_column' do
foreign_key.from_column = 'this_is_not_a_real_pk'
expect(foreign_key).not_to be_valid
expect(foreign_key.errors[:from_column].first).to eq('must be a valid column')
end
end
end
# frozen_string_literal: true
require 'spec_helper'
describe Gitlab::Database::PartitioningMigrationHelpers do
let(:model) do
ActiveRecord::Migration.new.extend(described_class)
end
let_it_be(:connection) { ActiveRecord::Base.connection }
let(:referenced_table) { :issues }
let(:function_name) { model.fk_function_name(referenced_table) }
let(:trigger_name) { model.fk_trigger_name(referenced_table) }
before do
allow(model).to receive(:puts)
end
describe 'adding a foreign key' do
before do
allow(model).to receive(:transaction_open?).and_return(false)
end
context 'when the table has no foreign keys' do
it 'creates a trigger function to handle the single cascade' do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the table already has foreign keys' do
context 'when the foreign key is from a different table' do
before do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
end
it 'creates a trigger function to handle the multiple cascades' do
model.add_partitioned_foreign_key :epic_issues, referenced_table
expect_function_to_contain(function_name,
'delete from issue_assignees where issue_id = old.id',
'delete from epic_issues where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the foreign key is from the same table' do
before do
model.add_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
end
context 'when the foreign key is from a different column' do
it 'creates a trigger function to handle the multiple cascades' do
model.add_partitioned_foreign_key :issues, referenced_table, column: :duplicated_to_id
expect_function_to_contain(function_name,
'delete from issues where moved_to_id = old.id',
'delete from issues where duplicated_to_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the foreign key is from the same column' do
it 'ignores the duplicate and properly recreates the trigger function' do
model.add_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
expect_function_to_contain(function_name, 'delete from issues where moved_to_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
end
end
context 'when the foreign key is set to nullify' do
it 'creates a trigger function that nullifies the foreign key' do
model.add_partitioned_foreign_key :issue_assignees, referenced_table, on_delete: :nullify
expect_function_to_contain(function_name, 'update issue_assignees set issue_id = null where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the referencing column is a custom value' do
it 'creates a trigger function with the correct column name' do
model.add_partitioned_foreign_key :issues, referenced_table, column: :duplicated_to_id
expect_function_to_contain(function_name, 'delete from issues where duplicated_to_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the referenced column is a custom value' do
let(:referenced_table) { :user_details }
it 'creates a trigger function with the correct column name' do
model.add_partitioned_foreign_key :user_preferences, referenced_table, column: :user_id, primary_key: :user_id
expect_function_to_contain(function_name, 'delete from user_preferences where user_id = old.user_id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the given key definition is invalid' do
it 'raises an error with the appropriate message' do
expect do
model.add_partitioned_foreign_key :issue_assignees, referenced_table, column: :not_a_real_issue_id
end.to raise_error(/From column must be a valid column/)
end
end
context 'when run inside a transaction' do
it 'raises an error' do
expect(model).to receive(:transaction_open?).and_return(true)
expect do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
end.to raise_error(/can not be run inside a transaction/)
end
end
end
context 'removing a foreign key' do
before do
allow(model).to receive(:transaction_open?).and_return(false)
end
context 'when the table has multiple foreign keys' do
before do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
model.add_partitioned_foreign_key :epic_issues, referenced_table
end
it 'creates a trigger function without the removed cascade' do
expect_function_to_contain(function_name,
'delete from issue_assignees where issue_id = old.id',
'delete from epic_issues where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
model.remove_partitioned_foreign_key :issue_assignees, referenced_table
expect_function_to_contain(function_name, 'delete from epic_issues where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when the table has only one remaining foreign key' do
before do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
end
it 'removes the trigger function altogether' do
expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
model.remove_partitioned_foreign_key :issue_assignees, referenced_table
expect(find_function_def(function_name)).to be_nil
expect(find_trigger_def(trigger_name)).to be_nil
end
end
context 'when the foreign key does not exist' do
before do
model.add_partitioned_foreign_key :issue_assignees, referenced_table
end
it 'ignores the invalid key and properly recreates the trigger function' do
expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
model.remove_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
expect_valid_function_trigger(trigger_name, function_name)
end
end
context 'when run outside a transaction' do
it 'raises an error' do
expect(model).to receive(:transaction_open?).and_return(true)
expect do
model.remove_partitioned_foreign_key :issue_assignees, referenced_table
end.to raise_error(/can not be run inside a transaction/)
end
end
end
def expect_function_to_contain(name, *statements)
return_stmt, *body_stmts = parsed_function_statements(name).reverse
expect(return_stmt).to eq('return old')
expect(body_stmts).to contain_exactly(*statements)
end
def expect_valid_function_trigger(name, fn_name)
event, activation, definition = cleaned_trigger_def(name)
expect(event).to eq('delete')
expect(activation).to eq('after')
expect(definition).to eq("execute procedure #{fn_name}()")
end
def parsed_function_statements(name)
cleaned_definition = find_function_def(name)['fn_body'].downcase.gsub(/\s+/, ' ')
statements = cleaned_definition.sub(/\A\s*begin\s*(.*)\s*end\s*\Z/, "\\1")
statements.split(';').map! { |stmt| stmt.strip.presence }.compact!
end
def find_function_def(name)
connection.execute("select prosrc as fn_body from pg_proc where proname = '#{name}';").first
end
def cleaned_trigger_def(name)
find_trigger_def(name).values_at('event', 'activation', 'definition').map!(&:downcase)
end
def find_trigger_def(name)
connection.execute(<<~SQL).first
select
string_agg(event_manipulation, ',') as event,
action_timing as activation,
action_statement as definition
from information_schema.triggers
where trigger_name = '#{name}'
group by 2, 3
SQL
end
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