Fix partitioning helpers for restricted migrations

Update partitioning helpers that combine DML and DDL operations to
require the correct mode of operation according to
`restrict_gitlab_migration`. Suppress any operationst that don't fit
within the required mode of operation.

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -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
 
-          min_date ||= connection.select_one(<<~SQL)['minimum'] || max_date - 1.month
-            SELECT date_trunc('MONTH', MIN(#{column_name})) AS minimum
-            FROM #{table_name}
-          SQL
+          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,8 +176,10 @@ module Gitlab
           primary_key = connection.primary_key(table_name)
           copy_missed_records(table_name, partitioned_table_name, primary_key)
 
-          disable_statement_timeout do
-            execute("VACUUM FREEZE ANALYZE #{partitioned_table_name}")
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
+            disable_statement_timeout do
+              execute("VACUUM FREEZE ANALYZE #{partitioned_table_name}")
+            end
           end
         end
 

spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb

@@ -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}")