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 0000000000000000000000000000000000000000..b8b9fc6d1566b4a88382982bd588b10edded1476 --- /dev/null +++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb @@ -0,0 +1,72 @@ +# 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) + + 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) + + 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? + + Gitlab::Database::PostgresPartitionedTable.each_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? + + Gitlab::Database::PostgresPartitionedTable.each_partition(table_name) do |partition| + unprepare_async_foreign_key_validation(partition.identifier, column_name, name: name) + end + end + + private + + 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 e41107370ecb2fff1dc84c59e347243f058809fd..9c1cb8e352cefeaf44cce7c3da7b346a7298af15 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/lib/gitlab/database/postgres_partitioned_table.rb b/lib/gitlab/database/postgres_partitioned_table.rb index 3bd342f940f8da7015e4724c2a7601901e383bd3..ef15f724cedfa7a0812c94dc00d85f71bf68c6f1 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 new file mode 100644 index 0000000000000000000000000000000000000000..0bd0e8045ff9313d4a0bd5377770e5d27402df5c --- /dev/null +++ b/spec/lib/gitlab/database/async_foreign_keys/migration_helpers_spec.rb @@ -0,0 +1,167 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::AsyncForeignKeys::MigrationHelpers, feature_category: :database do + 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' } + let(:column_name) { 'parent_id' } + let(:fk_name) { nil } + + context 'with regular tables' do + before do + allow(migration).to receive(:puts) + allow(migration.connection).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 + end + + 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 + 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 + + 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 + + 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(migration) + .to receive(:unprepare_async_foreign_key_validation) + .with(partition2_name, column_name, name: fk_name) + + migration.unprepare_partitioned_async_foreign_key_validation(table_name, column_name, name: fk_name) + end + end + end +end diff --git a/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb b/spec/lib/gitlab/database/postgres_partitioned_table_spec.rb index 23cf7329b73d5852e9316ac8d317b177d51eabf4..170cc89407162f82b876b8bbef101d899c070542 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