gitlab-org--gitlab-foss/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