From 2ce3e9043b61eb1b7d4c6a1f6876e84b77487882 Mon Sep 17 00:00:00 2001
From: Marius Bobin <mbobin@gitlab.com>
Date: Tue, 7 Feb 2023 13:57:20 +0200
Subject: [PATCH 1/4] Add migration helpers for async FK validation

---
 .../async_foreign_keys/migration_helpers.rb   |  62 ++++++++
 lib/gitlab/database/migration_helpers.rb      |   1 +
 .../migration_helpers_spec.rb                 | 137 ++++++++++++++++++
 3 files changed, 200 insertions(+)
 create mode 100644 lib/gitlab/database/async_foreign_keys/migration_helpers.rb
 create mode 100644 spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb

diff --git a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
new file mode 100644
index 00000000000000..034c5a4f8c1602
--- /dev/null
+++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncForeignKeys
+      module MigrationHelpers
+        # Prepares a foreign key for asynchronous validation.
+        #
+        # Stores the FK information in the postgres_async_foreign_key_validations
+        # table to be executed later.
+        #
+        def prepare_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          fk_name = name || concurrent_foreign_key_name(table_name, column_name)
+
+          raise 'Specifying foreign key name is mandatory - specify name: argument' unless fk_name
+
+          unless foreign_key_exists?(table_name, name: fk_name)
+            raise missing_schema_object_message(table_name, "foreign key", fk_name)
+          end
+
+          async_validation = PostgresAsyncForeignKeyValidation
+            .find_or_create_by!(name: fk_name, table_name: table_name)
+
+          Gitlab::AppLogger.info(
+            message: 'Prepared FK for async validation',
+            table_name: async_validation.table_name,
+            fk_name: async_validation.name)
+
+          async_validation
+        end
+
+        def unprepare_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          fk_name = name || concurrent_foreign_key_name(table_name, column_name)
+
+          raise 'Specifying foreign key name is mandatory - specify name: argument' unless fk_name
+
+          unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
+        end
+
+        def unprepare_async_foreign_key_validation_by_name(table_name, fk_name, **options)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          PostgresAsyncForeignKeyValidation.find_by(name: fk_name).try(&:destroy)
+        end
+
+        def async_fk_validation_available?
+          connection.table_exists?(:postgres_async_foreign_key_validations)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index e41107370ecb2f..9c1cb8e352cefe 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -14,6 +14,7 @@ module MigrationHelpers
       include DynamicModelHelpers
       include RenameTableHelpers
       include AsyncIndexes::MigrationHelpers
+      include AsyncForeignKeys::MigrationHelpers
 
       def define_batchable_model(table_name, connection: self.connection)
         super(table_name, connection: connection)
diff --git a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
new file mode 100644
index 00000000000000..3ff35be860d94f
--- /dev/null
+++ b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::AsyncForeignKeys::MigrationHelpers, feature_category: :database do
+  let(:migration) { ActiveRecord::Migration.new.extend(Gitlab::Database::MigrationHelpers) }
+  let(:connection) { ApplicationRecord.connection }
+  let(:fk_model) { Gitlab::Database::AsyncForeignKeys::PostgresAsyncForeignKeyValidation }
+  let(:table_name) { '_test_async_fks' }
+  let(:column_name) { 'parent_id' }
+  let(:fk_name) {}
+
+  before do
+    allow(migration).to receive(:puts)
+    allow(migration).to receive(:transaction_open?).and_return(false)
+
+    connection.create_table(table_name) do |t|
+      t.integer column_name
+    end
+
+    migration.add_concurrent_foreign_key(
+      table_name, table_name,
+      column: column_name, validate: false, name: fk_name)
+  end
+
+  describe '#prepare_async_foreign_key_validation' do
+    it 'creates the record for the async FK validation' do
+      expect do
+        migration.prepare_async_foreign_key_validation(table_name, column_name)
+      end.to change { fk_model.where(table_name: table_name).count }.by(1)
+
+      record = fk_model.find_by(table_name: table_name)
+
+      expect(record.name).to start_with('fk_')
+    end
+
+    context 'when an explicit name is given' do
+      let(:fk_name) { 'my_fk_name' }
+
+      it 'creates the record with the given name' do
+        expect do
+          migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
+        end.to change { fk_model.where(name: fk_name).count }.by(1)
+
+        record = fk_model.find_by(name: fk_name)
+
+        expect(record.table_name).to eq(table_name)
+      end
+    end
+
+    context 'when the FK does not exist' do
+      it 'returns an error' do
+        expect do
+          migration.prepare_async_foreign_key_validation(table_name, name: 'no_fk')
+        end.to raise_error RuntimeError, /Could not find foreign key "no_fk" on table "_test_async_fks"/
+      end
+    end
+
+    context 'when the record already exists' do
+      let(:fk_name) { 'my_fk_name' }
+
+      it 'does attempt to create the record' do
+        create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+
+        expect do
+          migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
+        end.not_to change { fk_model.where(name: fk_name).count }
+      end
+    end
+
+    context 'when the async FK validation table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_foreign_key_validations)
+
+        expect(fk_model).not_to receive(:safe_find_or_create_by!)
+
+        expect { migration.prepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+      end
+    end
+  end
+
+  describe '#unprepare_async_foreign_key_validation' do
+    before do
+      migration.prepare_async_foreign_key_validation(table_name, column_name, name: fk_name)
+    end
+
+    it 'destroys the record' do
+      expect do
+        migration.unprepare_async_foreign_key_validation(table_name, column_name)
+      end.to change { fk_model.where(table_name: table_name).count }.by(-1)
+    end
+
+    context 'when an explicit name is given' do
+      let(:fk_name) { 'my_test_async_fk' }
+
+      it 'destroys the record' do
+        expect do
+          migration.unprepare_async_foreign_key_validation(table_name, name: fk_name)
+        end.to change { fk_model.where(name: fk_name).count }.by(-1)
+      end
+    end
+
+    context 'when the async fk validation table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_foreign_key_validations)
+
+        expect(fk_model).not_to receive(:find_by)
+
+        expect { migration.unprepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+      end
+    end
+  end
+
+  describe '#unprepare_async_foreign_key_validation_by_name' do
+    let(:fk_name) { "fk_#{table_name}_on_#{column_name}" }
+
+    before do
+      create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+    end
+
+    it 'destroys the record' do
+      expect do
+        migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
+      end.to change { fk_model.where(name: fk_name).count }.by(-1)
+    end
+
+    context 'when the async fk validation table does not exist' do
+      it 'does not raise an error' do
+        connection.drop_table(:postgres_async_foreign_key_validations)
+
+        expect(fk_model).not_to receive(:find_by)
+
+        expect { migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name) }.not_to raise_error
+      end
+    end
+  end
+end
-- 
GitLab


From e5e31e88f36a5448805e9245dd4252e1f2ae40ba Mon Sep 17 00:00:00 2001
From: Marius Bobin <mbobin@gitlab.com>
Date: Wed, 8 Feb 2023 17:05:31 +0200
Subject: [PATCH 2/4] Add helpers for partitioned tables

---
 .../async_foreign_keys/migration_helpers.rb   |  30 +++
 .../migration_helpers_spec.rb                 | 204 +++++++++++-------
 2 files changed, 159 insertions(+), 75 deletions(-)

diff --git a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
index 034c5a4f8c1602..04558e3dad1d23 100644
--- a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
+++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
@@ -53,6 +53,36 @@ def unprepare_async_foreign_key_validation_by_name(table_name, fk_name, **option
           PostgresAsyncForeignKeyValidation.find_by(name: fk_name).try(&:destroy)
         end
 
+        def prepare_partitioned_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          each_table_partition(table_name) do |partition|
+            prepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
+          end
+        end
+
+        def unprepare_partitioned_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          each_table_partition(table_name) do |partition|
+            unprepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
+          end
+        end
+
+        private
+
+        def each_table_partition(table_name, &block)
+          Gitlab::Database::PostgresPartitionedTable
+            .find_by_name_in_current_schema(table_name)
+            .postgres_partitions
+            .order(:name)
+            .each(&block)
+        end
+
         def async_fk_validation_available?
           connection.table_exists?(:postgres_async_foreign_key_validations)
         end
diff --git a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
index 3ff35be860d94f..bf3439d5489186 100644
--- a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
@@ -10,127 +10,181 @@
   let(:column_name) { 'parent_id' }
   let(:fk_name) {}
 
-  before do
-    allow(migration).to receive(:puts)
-    allow(migration).to receive(:transaction_open?).and_return(false)
+  context 'with regular tables' do
+    before do
+      allow(migration).to receive(:puts)
+      allow(migration).to receive(:transaction_open?).and_return(false)
+
+      connection.create_table(table_name) do |t|
+        t.integer column_name
+      end
 
-    connection.create_table(table_name) do |t|
-      t.integer column_name
+      migration.add_concurrent_foreign_key(
+        table_name, table_name,
+        column: column_name, validate: false, name: fk_name)
     end
 
-    migration.add_concurrent_foreign_key(
-      table_name, table_name,
-      column: column_name, validate: false, name: fk_name)
-  end
+    describe '#prepare_async_foreign_key_validation' do
+      it 'creates the record for the async FK validation' do
+        expect do
+          migration.prepare_async_foreign_key_validation(table_name, column_name)
+        end.to change { fk_model.where(table_name: table_name).count }.by(1)
 
-  describe '#prepare_async_foreign_key_validation' do
-    it 'creates the record for the async FK validation' do
-      expect do
-        migration.prepare_async_foreign_key_validation(table_name, column_name)
-      end.to change { fk_model.where(table_name: table_name).count }.by(1)
+        record = fk_model.find_by(table_name: table_name)
 
-      record = fk_model.find_by(table_name: table_name)
+        expect(record.name).to start_with('fk_')
+      end
 
-      expect(record.name).to start_with('fk_')
-    end
+      context 'when an explicit name is given' do
+        let(:fk_name) { 'my_fk_name' }
 
-    context 'when an explicit name is given' do
-      let(:fk_name) { 'my_fk_name' }
+        it 'creates the record with the given name' do
+          expect do
+            migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
+          end.to change { fk_model.where(name: fk_name).count }.by(1)
 
-      it 'creates the record with the given name' do
-        expect do
-          migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
-        end.to change { fk_model.where(name: fk_name).count }.by(1)
+          record = fk_model.find_by(name: fk_name)
 
-        record = fk_model.find_by(name: fk_name)
+          expect(record.table_name).to eq(table_name)
+        end
+      end
 
-        expect(record.table_name).to eq(table_name)
+      context 'when the FK does not exist' do
+        it 'returns an error' do
+          expect do
+            migration.prepare_async_foreign_key_validation(table_name, name: 'no_fk')
+          end.to raise_error RuntimeError, /Could not find foreign key "no_fk" on table "_test_async_fks"/
+        end
       end
-    end
 
-    context 'when the FK does not exist' do
-      it 'returns an error' do
-        expect do
-          migration.prepare_async_foreign_key_validation(table_name, name: 'no_fk')
-        end.to raise_error RuntimeError, /Could not find foreign key "no_fk" on table "_test_async_fks"/
+      context 'when the record already exists' do
+        let(:fk_name) { 'my_fk_name' }
+
+        it 'does attempt to create the record' do
+          create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+
+          expect do
+            migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
+          end.not_to change { fk_model.where(name: fk_name).count }
+        end
       end
-    end
 
-    context 'when the record already exists' do
-      let(:fk_name) { 'my_fk_name' }
+      context 'when the async FK validation table does not exist' do
+        it 'does not raise an error' do
+          connection.drop_table(:postgres_async_foreign_key_validations)
 
-      it 'does attempt to create the record' do
-        create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+          expect(fk_model).not_to receive(:safe_find_or_create_by!)
 
-        expect do
-          migration.prepare_async_foreign_key_validation(table_name, name: fk_name)
-        end.not_to change { fk_model.where(name: fk_name).count }
+          expect { migration.prepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+        end
       end
     end
 
-    context 'when the async FK validation table does not exist' do
-      it 'does not raise an error' do
-        connection.drop_table(:postgres_async_foreign_key_validations)
+    describe '#unprepare_async_foreign_key_validation' do
+      before do
+        migration.prepare_async_foreign_key_validation(table_name, column_name, name: fk_name)
+      end
+
+      it 'destroys the record' do
+        expect do
+          migration.unprepare_async_foreign_key_validation(table_name, column_name)
+        end.to change { fk_model.where(table_name: table_name).count }.by(-1)
+      end
 
-        expect(fk_model).not_to receive(:safe_find_or_create_by!)
+      context 'when an explicit name is given' do
+        let(:fk_name) { 'my_test_async_fk' }
 
-        expect { migration.prepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+        it 'destroys the record' do
+          expect do
+            migration.unprepare_async_foreign_key_validation(table_name, name: fk_name)
+          end.to change { fk_model.where(name: fk_name).count }.by(-1)
+        end
       end
-    end
-  end
 
-  describe '#unprepare_async_foreign_key_validation' do
-    before do
-      migration.prepare_async_foreign_key_validation(table_name, column_name, name: fk_name)
-    end
+      context 'when the async fk validation table does not exist' do
+        it 'does not raise an error' do
+          connection.drop_table(:postgres_async_foreign_key_validations)
+
+          expect(fk_model).not_to receive(:find_by)
 
-    it 'destroys the record' do
-      expect do
-        migration.unprepare_async_foreign_key_validation(table_name, column_name)
-      end.to change { fk_model.where(table_name: table_name).count }.by(-1)
+          expect { migration.unprepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+        end
+      end
     end
 
-    context 'when an explicit name is given' do
-      let(:fk_name) { 'my_test_async_fk' }
+    describe '#unprepare_async_foreign_key_validation_by_name' do
+      let(:fk_name) { "fk_#{table_name}_on_#{column_name}" }
+
+      before do
+        create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+      end
 
       it 'destroys the record' do
         expect do
-          migration.unprepare_async_foreign_key_validation(table_name, name: fk_name)
+          migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
         end.to change { fk_model.where(name: fk_name).count }.by(-1)
       end
-    end
 
-    context 'when the async fk validation table does not exist' do
-      it 'does not raise an error' do
-        connection.drop_table(:postgres_async_foreign_key_validations)
+      context 'when the async fk validation table does not exist' do
+        it 'does not raise an error' do
+          connection.drop_table(:postgres_async_foreign_key_validations)
 
-        expect(fk_model).not_to receive(:find_by)
+          expect(fk_model).not_to receive(:find_by)
 
-        expect { migration.unprepare_async_foreign_key_validation(table_name, column_name) }.not_to raise_error
+          expect { migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name) }.not_to raise_error
+        end
       end
     end
   end
 
-  describe '#unprepare_async_foreign_key_validation_by_name' do
-    let(:fk_name) { "fk_#{table_name}_on_#{column_name}" }
+  context 'with partitioned tables' do
+    let(:partition_schema) { 'gitlab_partitions_dynamic' }
+    let(:partition1_name) { "#{partition_schema}.#{table_name}_202001" }
+    let(:partition2_name) { "#{partition_schema}.#{table_name}_202002" }
+    let(:fk_name) { 'my_partitioned_fk_name' }
 
     before do
-      create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
+      connection.execute(<<~SQL)
+        CREATE TABLE #{table_name} (
+          id serial NOT NULL,
+          #{column_name} int NOT NULL,
+          created_at timestamptz NOT NULL,
+          PRIMARY KEY (id, created_at)
+        ) PARTITION BY RANGE (created_at);
+
+        CREATE TABLE #{partition1_name} PARTITION OF #{table_name}
+        FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+
+        CREATE TABLE #{partition2_name} PARTITION OF #{table_name}
+        FOR VALUES FROM ('2020-02-01') TO ('2020-03-01');
+      SQL
     end
 
-    it 'destroys the record' do
-      expect do
-        migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
-      end.to change { fk_model.where(name: fk_name).count }.by(-1)
+    describe '#prepare_partitioned_async_foreign_key_validation' do
+      it 'delegates to prepare_async_foreign_key_validation for each partition' do
+        expect(migration)
+          .to receive(:prepare_async_foreign_key_validation)
+          .with(partition1_name, column_name, name: fk_name)
+
+        expect(migration)
+          .to receive(:prepare_async_foreign_key_validation)
+          .with(partition2_name, column_name, name: fk_name)
+
+        migration.prepare_partitioned_async_foreign_key_validation(table_name, column_name, name: fk_name)
+      end
     end
 
-    context 'when the async fk validation table does not exist' do
-      it 'does not raise an error' do
-        connection.drop_table(:postgres_async_foreign_key_validations)
+    describe '#unprepare_partitioned_async_foreign_key_validation' do
+      it 'delegates to unprepare_async_foreign_key_validation for each partition' do
+        expect(migration)
+          .to receive(:unprepare_async_foreign_key_validation)
+          .with(partition1_name, column_name, name: fk_name)
 
-        expect(fk_model).not_to receive(:find_by)
+        expect(migration)
+          .to receive(:unprepare_async_foreign_key_validation)
+          .with(partition2_name, column_name, name: fk_name)
 
-        expect { migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name) }.not_to raise_error
+        migration.unprepare_partitioned_async_foreign_key_validation(table_name, column_name, name: fk_name)
       end
     end
   end
-- 
GitLab


From 20dfccd87624aacdefb0fe801aef5615c3f43c35 Mon Sep 17 00:00:00 2001
From: Marius Bobin <mbobin@gitlab.com>
Date: Mon, 13 Feb 2023 12:12:08 +0200
Subject: [PATCH 3/4] Apply review feedback

---
 .../async_foreign_keys/migration_helpers.rb   | 16 ++--------
 .../database/postgres_partitioned_table.rb    |  7 ++++
 .../migration_helpers_spec.rb                 |  4 +--
 .../postgres_partitioned_table_spec.rb        | 32 +++++++++++++++++++
 4 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
index 04558e3dad1d23..76beec3d357229 100644
--- a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
+++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
@@ -16,8 +16,6 @@ def prepare_async_foreign_key_validation(table_name, column_name = nil, name: ni
 
           fk_name = name || concurrent_foreign_key_name(table_name, column_name)
 
-          raise 'Specifying foreign key name is mandatory - specify name: argument' unless fk_name
-
           unless foreign_key_exists?(table_name, name: fk_name)
             raise missing_schema_object_message(table_name, "foreign key", fk_name)
           end
@@ -40,8 +38,6 @@ def unprepare_async_foreign_key_validation(table_name, column_name = nil, name:
 
           fk_name = name || concurrent_foreign_key_name(table_name, column_name)
 
-          raise 'Specifying foreign key name is mandatory - specify name: argument' unless fk_name
-
           unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
         end
 
@@ -58,7 +54,7 @@ def prepare_partitioned_async_foreign_key_validation(table_name, column_name = n
 
           return unless async_fk_validation_available?
 
-          each_table_partition(table_name) do |partition|
+          Gitlab::Database::PostgresPartitionedTable.each_partition(table_name) do |partition|
             prepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
           end
         end
@@ -68,21 +64,13 @@ def unprepare_partitioned_async_foreign_key_validation(table_name, column_name =
 
           return unless async_fk_validation_available?
 
-          each_table_partition(table_name) do |partition|
+          Gitlab::Database::PostgresPartitionedTable.each_partition(table_name) do |partition|
             unprepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
           end
         end
 
         private
 
-        def each_table_partition(table_name, &block)
-          Gitlab::Database::PostgresPartitionedTable
-            .find_by_name_in_current_schema(table_name)
-            .postgres_partitions
-            .order(:name)
-            .each(&block)
-        end
-
         def async_fk_validation_available?
           connection.table_exists?(:postgres_async_foreign_key_validations)
         end
diff --git a/lib/gitlab/database/postgres_partitioned_table.rb b/lib/gitlab/database/postgres_partitioned_table.rb
index 3bd342f940f8da..ef15f724cedfa7 100644
--- a/lib/gitlab/database/postgres_partitioned_table.rb
+++ b/lib/gitlab/database/postgres_partitioned_table.rb
@@ -19,6 +19,13 @@ def self.find_by_name_in_current_schema(name)
         find_by("identifier = concat(current_schema(), '.', ?)", name)
       end
 
+      def self.each_partition(table_name, &block)
+        find_by_name_in_current_schema(table_name)
+        .postgres_partitions
+        .order(:name)
+        .each(&block)
+      end
+
       def dynamic?
         DYNAMIC_PARTITION_STRATEGIES.include?(strategy)
       end
diff --git a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
index bf3439d5489186..d415ba1508f73a 100644
--- a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::AsyncForeignKeys::MigrationHelpers, feature_category: :database do
-  let(:migration) { ActiveRecord::Migration.new.extend(Gitlab::Database::MigrationHelpers) }
+  let(:migration) { Gitlab::Database::Migration[2.1].new }
   let(:connection) { ApplicationRecord.connection }
   let(:fk_model) { Gitlab::Database::AsyncForeignKeys::PostgresAsyncForeignKeyValidation }
   let(:table_name) { '_test_async_fks' }
@@ -13,7 +13,7 @@
   context 'with regular tables' do
     before do
       allow(migration).to receive(:puts)
-      allow(migration).to receive(:transaction_open?).and_return(false)
+      allow(migration.connection).to receive(:transaction_open?).and_return(false)
 
       connection.create_table(table_name) do |t|
         t.integer column_name
diff --git a/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb b/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb
index 23cf7329b73d58..170cc89407162f 100644
--- a/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb
+++ b/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb
@@ -60,6 +60,38 @@ def find(identifier)
     end
   end
 
+  describe '.each_partition' do
+    context 'without partitions' do
+      it 'does not yield control' do
+        expect { |b| described_class.each_partition(name, &b) }.not_to yield_control
+      end
+    end
+
+    context 'with partitions' do
+      let(:partition_schema) { 'gitlab_partitions_dynamic' }
+      let(:partition1_name) { "#{partition_schema}.#{name}_202001" }
+      let(:partition2_name) { "#{partition_schema}.#{name}_202002" }
+
+      before do
+        ActiveRecord::Base.connection.execute(<<~SQL)
+          CREATE TABLE #{partition1_name} PARTITION OF #{identifier}
+          FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+
+          CREATE TABLE #{partition2_name} PARTITION OF #{identifier}
+          FOR VALUES FROM ('2020-02-01') TO ('2020-03-01');
+        SQL
+      end
+
+      it 'yields control with partition as argument' do
+        args = Gitlab::Database::PostgresPartition
+          .where(identifier: [partition1_name, partition2_name])
+          .order(:name).to_a
+
+        expect { |b| described_class.each_partition(name, &b) }.to yield_successive_args(*args)
+      end
+    end
+  end
+
   describe '#dynamic?' do
     it 'returns true for tables partitioned by range' do
       expect(find("#{schema}.#{foo_range_table_name}")).to be_dynamic
-- 
GitLab


From 223380ced7008f22245d4b298f11b2a28bbfe088 Mon Sep 17 00:00:00 2001
From: Marius Bobin <mbobin@gitlab.com>
Date: Wed, 15 Feb 2023 10:07:28 +0200
Subject: [PATCH 4/4] Remove unprepare_async_foreign_key_validation by name

---
 .../async_foreign_keys/migration_helpers.rb   |  8 ------
 .../migration_helpers_spec.rb                 | 26 +------------------
 2 files changed, 1 insertion(+), 33 deletions(-)

diff --git a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
index 76beec3d357229..b8b9fc6d1566b4 100644
--- a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
+++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
@@ -38,14 +38,6 @@ def unprepare_async_foreign_key_validation(table_name, column_name = nil, name:
 
           fk_name = name || concurrent_foreign_key_name(table_name, column_name)
 
-          unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
-        end
-
-        def unprepare_async_foreign_key_validation_by_name(table_name, fk_name, **options)
-          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
-
-          return unless async_fk_validation_available?
-
           PostgresAsyncForeignKeyValidation.find_by(name: fk_name).try(&:destroy)
         end
 
diff --git a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
index d415ba1508f73a..0bd0e8045ff931 100644
--- a/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb
@@ -8,7 +8,7 @@
   let(:fk_model) { Gitlab::Database::AsyncForeignKeys::PostgresAsyncForeignKeyValidation }
   let(:table_name) { '_test_async_fks' }
   let(:column_name) { 'parent_id' }
-  let(:fk_name) {}
+  let(:fk_name) { nil }
 
   context 'with regular tables' do
     before do
@@ -111,30 +111,6 @@
         end
       end
     end
-
-    describe '#unprepare_async_foreign_key_validation_by_name' do
-      let(:fk_name) { "fk_#{table_name}_on_#{column_name}" }
-
-      before do
-        create(:postgres_async_foreign_key_validation, table_name: table_name, name: fk_name)
-      end
-
-      it 'destroys the record' do
-        expect do
-          migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name)
-        end.to change { fk_model.where(name: fk_name).count }.by(-1)
-      end
-
-      context 'when the async fk validation table does not exist' do
-        it 'does not raise an error' do
-          connection.drop_table(:postgres_async_foreign_key_validations)
-
-          expect(fk_model).not_to receive(:find_by)
-
-          expect { migration.unprepare_async_foreign_key_validation_by_name(table_name, fk_name) }.not_to raise_error
-        end
-      end
-    end
   end
 
   context 'with partitioned tables' do
-- 
GitLab