Enhance Rake tasks to dump custom structure.sql

Add additional logic to buitlin Rails tasks that dump and load the
structure.sql file, to also dump and load custom data needed track
partitioned foreign keys.

db/gitlab_structure.sql

+SET search_path=public;
+

lib/gitlab/database/custom_structure.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class CustomStructure
+      def dump(io)
+        io << "SET search_path=public;\n\n"
+
+        dump_partitioned_foreign_keys(io) if partitioned_foreign_keys_exist?
+      end
+
+      private
+
+      def dump_partitioned_foreign_keys(io)
+        io << "COPY partitioned_foreign_keys (#{partitioned_fk_columns.join(", ")}) FROM STDIN;\n"
+
+        PartitioningMigrationHelpers::PartitionedForeignKey.find_each do |fk|
+          io << fk.attributes.values_at(*partitioned_fk_columns).join("\t") << "\n"
+        end
+        io << "\\.\n"
+      end
+
+      def partitioned_foreign_keys_exist?
+        return false unless PartitioningMigrationHelpers::PartitionedForeignKey.table_exists?
+
+        PartitioningMigrationHelpers::PartitionedForeignKey.exists?
+      end
+
+      def partitioned_fk_columns
+        @partitioned_fk_columns ||= PartitioningMigrationHelpers::PartitionedForeignKey.column_names
+      end
+    end
+  end
+end

lib/tasks/gitlab/db.rake

 namespace :gitlab do
   namespace :db do
+    CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'.freeze
+
     desc 'GitLab | DB | Manually insert schema migration version'
     task :mark_migration_complete, [:version] => :environment do |_, args|
       unless args[:version]
@@ -92,9 +94,39 @@ namespace :gitlab do
       Rake::Task[task_name].reenable
     end
 
-    # Inform Rake that gitlab:schema:clean_structure_sql should be run every time rake db:structure:dump is run
+    desc 'This dumps GitLab specific database details - it runs after db:structure:dump'
+    task :dump_custom_structure do |task_name|
+      File.open(CUSTOM_DUMP_FILE, 'wb+') do |io|
+        Gitlab::Database::CustomStructure.new.dump(io)
+      end
+
+      # Allow this task to be called multiple times, as happens when running db:migrate:redo
+      Rake::Task[task_name].reenable
+    end
+
+    desc 'This loads Gitlab specific database details - runs after db:structure:dump'
+    task :load_custom_structure do
+      configuration = Rails.application.config_for(:database)
+
+      ENV['PGHOST']     = configuration['host']          if configuration['host']
+      ENV['PGPORT']     = configuration['port'].to_s     if configuration['port']
+      ENV['PGPASSWORD'] = configuration['password'].to_s if configuration['password']
+      ENV['PGUSER']     = configuration['username'].to_s if configuration['username']
+
+      args = ['-v', 'ON_ERROR_STOP=1', '-q', '-X', '-f', CUSTOM_DUMP_FILE]
+      args << configuration['database']
+      system('psql', *args)
+    end
+
+    # Inform Rake that custom tasks should be run every time rake db:structure:dump is run
     Rake::Task['db:structure:dump'].enhance do
       Rake::Task['gitlab:db:clean_structure_sql'].invoke
+      Rake::Task['gitlab:db:dump_custom_structure'].invoke
+    end
+
+    # Inform Rake that custom tasks should be run every time rake db:structure:load is run
+    Rake::Task['db:structure:load'].enhance do
+      Rake::Task['gitlab:db:load_custom_structure'].invoke
     end
   end
 end

spec/lib/gitlab/database/custom_structure_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Database::CustomStructure do
+  let_it_be(:structure) { described_class.new }
+  let(:io) { StringIO.new }
+
+  context 'when there are no partitioned_foreign_keys' do
+    it 'dumps a valid structure file' do
+      structure.dump(io)
+
+      expect(io.string).to eq("SET search_path=public;\n\n")
+    end
+  end
+
+  context 'when there are partitioned_foreign_keys' do
+    let!(:first_fk) do
+      Gitlab::Database::PartitioningMigrationHelpers::PartitionedForeignKey.create(
+        cascade_delete: true, from_table: 'issues', from_column: 'project_id', to_table: 'projects', to_column: 'id')
+    end
+    let!(:second_fk) do
+      Gitlab::Database::PartitioningMigrationHelpers::PartitionedForeignKey.create(
+        cascade_delete: false, from_table: 'issues', from_column: 'moved_to_id', to_table: 'issues', to_column: 'id')
+    end
+
+    it 'dumps a file with the command to restore the keys' do
+      structure.dump(io)
+
+      expect(io.string).to eq(<<~DATA)
+        SET search_path=public;
+
+        COPY partitioned_foreign_keys (id, cascade_delete, from_table, from_column, to_table, to_column) FROM STDIN;
+        #{first_fk.id}\ttrue\tissues\tproject_id\tprojects\tid
+        #{second_fk.id}\tfalse\tissues\tmoved_to_id\tissues\tid
+        \\.
+      DATA
+    end
+  end
+end