Rename partitions when swapping partitioned tables

When replacing a non-partitioned table with its partitioned copy, also
rename the underlying partitions to follow the same naming scheme as the
parent table. This includes renaming the primary key constraints of
those partitions, which typically would be `"#{table_name}_pkey"`.

lib/gitlab/database/partitioning/replace_table.rb

@@ -43,38 +43,70 @@ module Gitlab
         end
 
         def sql_to_replace_table
-          @sql_to_replace_table ||= [
-            drop_default_sql(original_table, primary_key_column),
-            set_default_sql(replacement_table, primary_key_column, "nextval('#{quote_table_name(sequence)}'::regclass)"),
+          @sql_to_replace_table ||= combined_sql_statements.map(&:chomp).join(DELIMITER)
+        end
+
+        def combined_sql_statements
+          statements = []
+
+          statements << alter_column_default(original_table, primary_key_column, expression: nil)
+          statements << alter_column_default(replacement_table, primary_key_column,
+              expression: "nextval('#{quote_table_name(sequence)}'::regclass)")
+
+          statements << alter_sequence_owned_by(sequence, replacement_table, primary_key_column)
 
-            change_sequence_owner_sql(sequence, replacement_table, primary_key_column),
+          rename_table_objects(statements, original_table, replaced_table, original_primary_key, replaced_primary_key)
+          rename_table_objects(statements, replacement_table, original_table, replacement_primary_key, original_primary_key)
 
-            rename_table_sql(original_table, replaced_table),
-            rename_constraint_sql(replaced_table, original_primary_key, replaced_primary_key),
+          statements
+        end
+
+        def rename_table_objects(statements, old_table, new_table, old_primary_key, new_primary_key)
+          statements << rename_table(old_table, new_table)
+          statements << rename_constraint(new_table, old_primary_key, new_primary_key)
 
-            rename_table_sql(replacement_table, original_table),
-            rename_constraint_sql(original_table, replacement_primary_key, original_primary_key)
-          ].join(DELIMITER)
+          rename_partitions(statements, old_table, new_table)
         end
 
-        def drop_default_sql(table, column)
-          "ALTER TABLE #{quote_table_name(table)} ALTER COLUMN #{quote_column_name(column)} DROP DEFAULT"
+        def rename_partitions(statements, old_table_name, new_table_name)
+          Gitlab::Database::PostgresPartition.for_parent_table(old_table_name).each do |partition|
+            new_partition_name = partition.name.sub(/#{old_table_name}/, new_table_name)
+            old_primary_key = default_primary_key(partition.name)
+            new_primary_key = default_primary_key(new_partition_name)
+
+            statements << rename_constraint(partition.identifier, old_primary_key, new_primary_key)
+            statements << rename_table(partition.identifier, new_partition_name)
+          end
         end
 
-        def set_default_sql(table, column, expression)
-          "ALTER TABLE #{quote_table_name(table)} ALTER COLUMN #{quote_column_name(column)} SET DEFAULT #{expression}"
+        def alter_column_default(table_name, column_name, expression:)
+          default_clause = expression.nil? ? 'DROP DEFAULT' : "SET DEFAULT #{expression}"
+
+          <<~SQL
+            ALTER TABLE #{quote_table_name(table_name)}
+            ALTER COLUMN #{quote_column_name(column_name)} #{default_clause}
+          SQL
         end
 
-        def change_sequence_owner_sql(sequence, table, column)
-          "ALTER SEQUENCE #{quote_table_name(sequence)} OWNED BY #{quote_table_name(table)}.#{quote_column_name(column)}"
+        def alter_sequence_owned_by(sequence_name, table_name, column_name)
+          <<~SQL
+            ALTER SEQUENCE #{quote_table_name(sequence_name)}
+            OWNED BY #{quote_table_name(table_name)}.#{quote_column_name(column_name)}
+          SQL
         end
 
-        def rename_table_sql(old_name, new_name)
-          "ALTER TABLE #{quote_table_name(old_name)} RENAME TO #{quote_table_name(new_name)}"
+        def rename_table(old_name, new_name)
+          <<~SQL
+            ALTER TABLE #{quote_table_name(old_name)}
+            RENAME TO #{quote_table_name(new_name)}
+          SQL
         end
 
-        def rename_constraint_sql(table, old_name, new_name)
-          "ALTER TABLE #{quote_table_name(table)} RENAME CONSTRAINT #{quote_column_name(old_name)} TO #{quote_column_name(new_name)}"
+        def rename_constraint(table_name, old_name, new_name)
+          <<~SQL
+            ALTER TABLE #{quote_table_name(table_name)}
+            RENAME CONSTRAINT #{quote_column_name(old_name)} TO #{quote_column_name(new_name)}
+          SQL
         end
       end
     end

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -179,7 +179,8 @@ module Gitlab
         # 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
         # replaced with the original table in such a way that it appears seamless to any database clients. The replaced
-        # table will be renamed to "#{replaced_table}_archived"
+        # table will be renamed to "#{replaced_table}_archived". Partitions and primary key constraints will also be
+        # renamed to match the naming scheme of the parent table.
         #
         # **NOTE** This method should only be used after all other migration steps have completed successfully.
         # There are several limitations to this method that MUST be handled before, or during, the swap migration:
@@ -415,7 +416,7 @@ module Gitlab
         end
 
         def replace_table(original_table_name, replacement_table_name, replaced_table_name, primary_key_name)
-          replace_table = Gitlab::Database::Partitioning::ReplaceTable.new(original_table_name,
+          replace_table = Gitlab::Database::Partitioning::ReplaceTable.new(original_table_name.to_s,
               replacement_table_name, replaced_table_name, primary_key_name)
 
           with_lock_retries do

spec/lib/gitlab/database/partitioning/replace_table_spec.rb

@@ -30,9 +30,6 @@ RSpec.describe Gitlab::Database::Partitioning::ReplaceTable, '#perform' do
         created_at timestamptz NOT NULL,
         PRIMARY KEY (id, created_at))
         PARTITION BY RANGE (created_at);
-
-      CREATE TABLE #{replacement_table}_202001 PARTITION OF #{replacement_table}
-      FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
     SQL
   end
 
@@ -56,13 +53,58 @@ RSpec.describe Gitlab::Database::Partitioning::ReplaceTable, '#perform' do
   end
 
   it 'renames the primary key constraints to match the new table names' do
-    expect(primary_key_constraint_name(original_table)).to eq(original_primary_key)
-    expect(primary_key_constraint_name(replacement_table)).to eq(replacement_primary_key)
+    expect_primary_keys_after_tables([original_table, replacement_table])
 
     expect_table_to_be_replaced { replace_table }
 
-    expect(primary_key_constraint_name(original_table)).to eq(original_primary_key)
-    expect(primary_key_constraint_name(archived_table)).to eq(archived_primary_key)
+    expect_primary_keys_after_tables([original_table, archived_table])
+  end
+
+  context 'when the table has partitions' do
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE gitlab_partitions_dynamic.#{replacement_table}_202001 PARTITION OF #{replacement_table}
+        FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+
+        CREATE TABLE gitlab_partitions_dynamic.#{replacement_table}_202002 PARTITION OF #{replacement_table}
+        FOR VALUES FROM ('2020-02-01') TO ('2020-03-01');
+      SQL
+    end
+
+    it 'renames the partitions to match the new table name' do
+      expect(partitions_for_parent_table(original_table).count).to eq(0)
+      expect(partitions_for_parent_table(replacement_table).count).to eq(2)
+
+      expect_table_to_be_replaced { replace_table }
+
+      expect(partitions_for_parent_table(archived_table).count).to eq(0)
+
+      partitions = partitions_for_parent_table(original_table).all
+
+      expect(partitions.size).to eq(2)
+
+      expect(partitions[0]).to have_attributes(
+        identifier: "gitlab_partitions_dynamic.#{original_table}_202001",
+        condition: "FOR VALUES FROM ('2020-01-01 00:00:00+00') TO ('2020-02-01 00:00:00+00')")
+
+      expect(partitions[1]).to have_attributes(
+        identifier: "gitlab_partitions_dynamic.#{original_table}_202002",
+        condition: "FOR VALUES FROM ('2020-02-01 00:00:00+00') TO ('2020-03-01 00:00:00+00')")
+    end
+
+    it 'renames the primary key constraints to match the new partition names' do
+      original_partitions = ["#{replacement_table}_202001", "#{replacement_table}_202002"]
+      expect_primary_keys_after_tables(original_partitions, schema: 'gitlab_partitions_dynamic')
+
+      expect_table_to_be_replaced { replace_table }
+
+      renamed_partitions = ["#{original_table}_202001", "#{original_table}_202002"]
+      expect_primary_keys_after_tables(renamed_partitions, schema: 'gitlab_partitions_dynamic')
+    end
+  end
+
+  def partitions_for_parent_table(table)
+    Gitlab::Database::PostgresPartition.for_parent_table(table)
   end
 
   def expect_table_to_be_replaced(&block)

spec/support/helpers/table_schema_helpers.rb

@@ -24,6 +24,14 @@ module TableSchemaHelpers
     expect(index_exists_by_name(name, schema: schema)).to be_nil
   end
 
+  def expect_primary_keys_after_tables(tables, schema: nil)
+    tables.each do |table|
+      primary_key = primary_key_constraint_name(table, schema: schema)
+
+      expect(primary_key).to eq("#{table}_pkey")
+    end
+  end
+
   def table_oid(name)
     connection.select_value(<<~SQL)
       SELECT oid
@@ -75,13 +83,15 @@ module TableSchemaHelpers
     SQL
   end
 
-  def primary_key_constraint_name(table_name)
+  def primary_key_constraint_name(table_name, schema: nil)
+    table_name = schema ? "#{schema}.#{table_name}" : table_name
+
     connection.select_value(<<~SQL)
       SELECT
         conname AS constraint_name
       FROM pg_catalog.pg_constraint
-      WHERE conrelid = '#{table_name}'::regclass
-        AND contype = 'p'
+      WHERE pg_constraint.conrelid = '#{table_name}'::regclass
+        AND pg_constraint.contype = 'p'
     SQL
   end