Helpers to manage FK cascades on partitioned tables

Add new Migration Helpers which allow cascading delete/nullify behavior
from partitioned tables, which don't support foreign keys natively in
PG11.

lib/gitlab/database/partitioning_migration_helpers.rb

+# 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

lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key.rb

+# 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

lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key_validator.rb

+# 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

lib/gitlab/database/schema_helpers.rb

+# 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

spec/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key_spec.rb

+# 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

spec/lib/gitlab/database/partitioning_migration_helpers_spec.rb

+# 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