Merge branch '241267-rename-partitions-when-swapping-tables' into 'master'

Rename partitions when swapping partitioned tables

See merge request gitlab-org/gitlab!46489

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