1 files changed, 2 insertions, 114 deletions
diff --git a/lib/gitlab/database/partitioning_migration_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers.rb
index 55649ebbf8a..881177a195e 100644
--- a/lib/gitlab/database/partitioning_migration_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers.rb
@@ -3,120 +3,8 @@
 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
+      include ForeignKeyHelpers
+      include TableManagementHelpers
     end
   end
 end