Merge branch '201872-partitioning-implement-cascading-deletes-without-foreign-keys-3' into 'master'

Dump/Load FK data for partitioned tables

See merge request gitlab-org/gitlab!29525

db/gitlab_structure.sql

+-- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables
+-- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872
+SET search_path=public;
+

lib/gitlab/database/custom_structure.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class CustomStructure
+      CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'
+
+      def dump
+        File.open(self.class.custom_dump_filepath, 'wb') do |io|
+          io << "-- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables\n"
+          io << "-- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872\n"
+          io << "SET search_path=public;\n\n"
+
+          dump_partitioned_foreign_keys(io) if partitioned_foreign_keys_exist?
+        end
+      end
+
+      def self.custom_dump_filepath
+        Rails.root.join(CUSTOM_DUMP_FILE)
+      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

@@ -92,9 +92,42 @@ 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|
+      Gitlab::Database::CustomStructure.new.dump
+
+      # 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']
+
+      command = 'psql'
+      dump_filepath = Gitlab::Database::CustomStructure.custom_dump_filepath.to_path
+      args = ['-v', 'ON_ERROR_STOP=1', '-q', '-X', '-f', dump_filepath, configuration['database']]
+
+      unless Kernel.system(command, *args)
+        raise "failed to execute:\n#{command} #{args.join(' ')}\n\n" \
+          "Please ensure `#{command}` is installed in your PATH and has proper permissions.\n\n"
+      end
+    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_it_be(:filepath) { Rails.root.join(described_class::CUSTOM_DUMP_FILE) }
+  let_it_be(:file_header) do
+    <<~DATA
+      -- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables
+      -- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872
+      SET search_path=public;
+    DATA
+  end
+
+  let(:io) { StringIO.new }
+
+  before do
+    allow(File).to receive(:open).with(filepath, anything).and_yield(io)
+  end
+
+  context 'when there are no partitioned_foreign_keys' do
+    it 'dumps a valid structure file' do
+      structure.dump
+
+      expect(io.string).to eq("#{file_header}\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 current keys' do
+      structure.dump
+
+      expect(io.string).to eq(<<~DATA)
+        #{file_header}
+        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
+
+      first_fk.destroy
+      io.truncate(0)
+      io.rewind
+
+      structure.dump
+
+      expect(io.string).to eq(<<~DATA)
+        #{file_header}
+        COPY partitioned_foreign_keys (id, cascade_delete, from_table, from_column, to_table, to_column) FROM STDIN;
+        #{second_fk.id}\tfalse\tissues\tmoved_to_id\tissues\tid
+        \\.
+      DATA
+    end
+  end
+end

spec/tasks/gitlab/db_rake_spec.rb

@@ -115,19 +115,52 @@ describe 'gitlab:db namespace rake task' do
     end
 
     it 'can be executed multiple times within another rake task' do
-      Rake::Task.define_task(test_task_name => :environment) do
+      expect_multiple_executions_of_task(test_task_name, clean_rake_task) do
         expect_next_instance_of(Gitlab::Database::SchemaCleaner) do |cleaner|
           expect(cleaner).to receive(:clean).with(output)
         end
-        Rake::Task[clean_rake_task].invoke
+      end
+    end
+  end
 
-        expect_next_instance_of(Gitlab::Database::SchemaCleaner) do |cleaner|
-          expect(cleaner).to receive(:clean).with(output)
+  describe 'load_custom_structure' do
+    let_it_be(:db_config) { Rails.application.config_for(:database) }
+    let_it_be(:custom_load_task) { 'gitlab:db:load_custom_structure' }
+    let_it_be(:custom_filepath) { Pathname.new('db/directory') }
+
+    it 'uses the psql command to load the custom structure file' do
+      expect(Gitlab::Database::CustomStructure).to receive(:custom_dump_filepath).and_return(custom_filepath)
+
+      expect(Kernel).to receive(:system)
+        .with('psql', any_args, custom_filepath.to_path, db_config['database']).and_return(true)
+
+      run_rake_task(custom_load_task)
+    end
+
+    it 'raises an error when the call to the psql command fails' do
+      expect(Gitlab::Database::CustomStructure).to receive(:custom_dump_filepath).and_return(custom_filepath)
+
+      expect(Kernel).to receive(:system)
+        .with('psql', any_args, custom_filepath.to_path, db_config['database']).and_return(nil)
+
+      expect { run_rake_task(custom_load_task) }.to raise_error(/failed to execute:\s*psql/)
+    end
+  end
+
+  describe 'dump_custom_structure' do
+    let_it_be(:test_task_name) { 'gitlab:db:_test_multiple_task_executions' }
+    let_it_be(:custom_dump_task) { 'gitlab:db:dump_custom_structure' }
+
+    after do
+      Rake::Task[test_task_name].clear if Rake::Task.task_defined?(test_task_name)
+    end
+
+    it 'can be executed multiple times within another rake task' do
+      expect_multiple_executions_of_task(test_task_name, custom_dump_task) do
+        expect_next_instance_of(Gitlab::Database::CustomStructure) do |custom_structure|
+          expect(custom_structure).to receive(:dump)
         end
-        Rake::Task[clean_rake_task].invoke
       end
-
-      run_rake_task(test_task_name)
     end
   end
 
@@ -135,4 +168,16 @@ describe 'gitlab:db namespace rake task' do
     Rake::Task[task_name].reenable
     Rake.application.invoke_task task_name
   end
+
+  def expect_multiple_executions_of_task(test_task_name, task_to_invoke, count: 2)
+    Rake::Task.define_task(test_task_name => :environment) do
+      count.times do
+        yield
+
+        Rake::Task[task_to_invoke].invoke
+      end
+    end
+
+    run_rake_task(test_task_name)
+  end
 end