diff options
Diffstat (limited to 'lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb')
| -rw-r--r-- | lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb | 140 |
1 files changed, 140 insertions, 0 deletions
diff --git a/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb new file mode 100644 index 00000000000..9e687009cd7 --- /dev/null +++ b/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb @@ -0,0 +1,140 @@ +# frozen_string_literal: true + +module Gitlab + module Database + module PartitioningMigrationHelpers + module ForeignKeyHelpers + include ::Gitlab::Database::SchemaHelpers + + # Creates a "foreign key" that references a partitioned table. Because foreign keys referencing partitioned + # tables are not supported in PG11, this does not create a true database foreign key, but instead implements the + # same functionality at the database level by using triggers. + # + # Example: + # + # add_partitioned_foreign_key :issues, :projects + # + # Available options: + # + # :column - name of the referencing column (otherwise inferred from the referenced table name) + # :primary_key - name of the primary key in the referenced table (defaults to id) + # :on_delete - supports either :cascade for ON DELETE CASCADE or :nullify for ON DELETE SET NULL + # + 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 + + # Drops a "foreign key" that references a partitioned table. This method ONLY applies to foreign keys previously + # created through the `add_partitioned_foreign_key` method. Standard database foreign keys should be managed + # through the familiar Rails helpers. + # + # Example: + # + # remove_partitioned_foreign_key :issues, :projects + # + # Available options: + # + # :column - name of the referencing column (otherwise inferred from the referenced table name) + # :primary_key - name of the primary key in the referenced table (defaults to id) + # + 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.delete + 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 + + private + + def fk_function_name(table) + object_name(table, 'fk_cascade_function') + end + + def fk_trigger_name(table) + object_name(table, 'fk_cascade_trigger') + end + + 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) + assert_not_in_transaction_block(scope: 'partitioned foreign key') + + 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_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 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 +end |