123 lines
4.3 KiB
Ruby
123 lines
4.3 KiB
Ruby
|
# 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
|