Commit d214b234 authored by Thong Kuah's avatar Thong Kuah

Merge branch '356276-validate-partitioning-helpers-with-restricted-schema' into 'master'

Fix partitioning helpers for restricted migrations

See merge request gitlab-org/gitlab!83342
parents 6d0090cd cb313041
......@@ -40,16 +40,20 @@ module Gitlab
# 1. The minimum value for the partitioning column in the table
# 2. If no data is present yet, the current month
def partition_table_by_date(table_name, column_name, min_date: nil, max_date: nil)
Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
assert_table_is_allowed(table_name)
assert_not_in_transaction_block(scope: ERROR_SCOPE)
max_date ||= Date.today + 1.month
Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
min_date ||= connection.select_one(<<~SQL)['minimum'] || max_date - 1.month
SELECT date_trunc('MONTH', MIN(#{column_name})) AS minimum
FROM #{table_name}
SQL
end
raise "max_date #{max_date} must be greater than min_date #{min_date}" if min_date >= max_date
......@@ -154,6 +158,8 @@ module Gitlab
# finalize_backfilling_partitioned_table :audit_events
#
def finalize_backfilling_partitioned_table(table_name)
Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_dml_mode!
assert_table_is_allowed(table_name)
assert_not_in_transaction_block(scope: ERROR_SCOPE)
......@@ -170,10 +176,12 @@ module Gitlab
primary_key = connection.primary_key(table_name)
copy_missed_records(table_name, partitioned_table_name, primary_key)
Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
disable_statement_timeout do
execute("VACUUM FREEZE ANALYZE #{partitioned_table_name}")
end
end
end
# Replaces a non-partitioned table with its partitioned copy. This is the final step in a partitioning
# migration, which makes the partitioned table ready for use by the application. The partitioned copy should be
......
......@@ -125,6 +125,17 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
expect_table_partitioned_by(partitioned_table, [partition_column])
end
it 'requires the migration helper to be run in DDL mode' do
expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_ddl_mode!)
migration.partition_table_by_date source_table, partition_column, min_date: min_date, max_date: max_date
expect(connection.table_exists?(partitioned_table)).to be(true)
expect(connection.primary_key(partitioned_table)).to eq(new_primary_key)
expect_table_partitioned_by(partitioned_table, [partition_column])
end
it 'changes the primary key datatype to bigint' do
migration.partition_table_by_date source_table, partition_column, min_date: min_date, max_date: max_date
......@@ -191,6 +202,8 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
end
it 'creates a partition spanning over each month from the first record' do
expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:with_suppressed).and_yield
migration.partition_table_by_date source_table, partition_column, max_date: max_date
expect_range_partitions_for(partitioned_table, {
......@@ -206,6 +219,8 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
context 'without data' do
it 'creates the catchall partition plus two actual partition' do
expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:with_suppressed).and_yield
migration.partition_table_by_date source_table, partition_column, max_date: max_date
expect_range_partitions_for(partitioned_table, {
......@@ -536,6 +551,16 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
migration.finalize_backfilling_partitioned_table source_table
end
it 'requires the migration helper to execute in DML mode' do
expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:require_dml_mode!)
expect(Gitlab::BackgroundMigration).to receive(:steal)
.with(described_class::MIGRATION_CLASS_NAME)
.and_yield(background_job)
migration.finalize_backfilling_partitioned_table source_table
end
end
context 'when there is missed data' do
......@@ -627,6 +652,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
allow(backfill).to receive(:perform).and_return(1)
end
expect(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas).to receive(:with_suppressed).and_yield
expect(migration).to receive(:disable_statement_timeout).and_call_original
expect(migration).to receive(:execute).with("VACUUM FREEZE ANALYZE #{partitioned_table}")
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment