1 files changed, 58 insertions, 45 deletions

diff --git a/doc/development/database/table_partitioning.md b/doc/development/database/table_partitioning.md
index 30131fc0347..0d5e3c233f6 100644
--- a/doc/development/database/table_partitioning.md
+++ b/doc/development/database/table_partitioning.md
@@ -307,44 +307,12 @@ class PrepareIndexesForPartitioning < Gitlab::Database::Migration[2.1]
 end
 ```
 
-### Step 3 - Swap primary key
+### Step 3 - Enforce unique constraint
 
-Swap the primary key including the partitioning key column. For example, in a rails migration:
-
-```ruby
-class PreparePrimaryKeyForPartitioning < Gitlab::Database::Migration[2.1]
-  disable_ddl_transaction!
-
-  TABLE_NAME = :table_name
-  PRIMARY_KEY = :primary_key
-  OLD_INDEX_NAME = :old_index_name
-  NEW_INDEX_NAME = :new_index_name
-
-  def up
-    swap_primary_key(TABLE_NAME, PRIMARY_KEY, NEW_INDEX_NAME)
-  end
-
-  def down
-    add_concurrent_index(TABLE_NAME, :id, unique: true, name: OLD_INDEX_NAME)
-    add_concurrent_index(TABLE_NAME, [:id, :partition_id], unique: true, name: NEW_INDEX_NAME)
-
-    unswap_primary_key(TABLE_NAME, PRIMARY_KEY, OLD_INDEX_NAME)
-  end
-end
-```
-
-NOTE:
-Do not forget to set the primary key explicitly in your model as `ActiveRecord` does not support composite primary keys.
-
-```ruby
-class Model < ApplicationRecord
-  self.primary_key = :id
-end
-```
-
-### Step 4 - Enforce unique constraint
-
-Enforce unique indexes including the partitioning key column. For example, in a rails migration:
+Change all unique indexes to include the partitioning key column,
+including the primary key index. You can start by adding an unique
+index on `[primary_key_column, :partition_id]`, which will be
+required for the next two steps. For example, in a rails migration:
 
 ```ruby
 class PrepareUniqueContraintForPartitioning < Gitlab::Database::Migration[2.1]
@@ -355,20 +323,20 @@ class PrepareUniqueContraintForPartitioning < Gitlab::Database::Migration[2.1]
   NEW_UNIQUE_INDEX_NAME = :new_index_name
 
   def up
-    add_concurrent_index(TABLE_NAME, [:some_column, :partition_id], unique: true, name: NEW_UNIQUE_INDEX_NAME)
+    add_concurrent_index(TABLE_NAME, [:id, :partition_id], unique: true, name: NEW_UNIQUE_INDEX_NAME)
 
     remove_concurrent_index_by_name(TABLE_NAME, OLD_UNIQUE_INDEX_NAME)
   end
 
   def down
-    add_concurrent_index(TABLE_NAME, :some_column, unique: true, name: OLD_UNIQUE_INDEX_NAME)
+    add_concurrent_index(TABLE_NAME, :id, unique: true, name: OLD_UNIQUE_INDEX_NAME)
 
     remove_concurrent_index_by_name(TABLE_NAME, NEW_UNIQUE_INDEX_NAME)
   end
 end
 ```
 
-### Step 5 - Enforce foreign key constraint
+### Step 4 - Enforce foreign key constraint
 
 Enforce foreign keys including the partitioning key column. For example, in a rails migration:
 
@@ -380,7 +348,7 @@ class PrepareForeignKeyForPartitioning < Gitlab::Database::Migration[2.1]
   TARGET_TABLE_NAME = :target_table_name
   COLUMN = :foreign_key_id
   TARGET_COLUMN = :id
-  CONSTRAINT_NAME = :fk_365d1db505_p
+  FK_NAME = :fk_365d1db505_p
   PARTITION_COLUMN = :partition_id
 
   def up
@@ -391,14 +359,17 @@ class PrepareForeignKeyForPartitioning < Gitlab::Database::Migration[2.1]
       target_column: [PARTITION_COLUMN, TARGET_COLUMN],
       validate: false,
       on_update: :cascade,
-      name: CONSTRAINT_NAME
+      name: FK_NAME
     )
 
-    validate_foreign_key(TARGET_TABLE_NAME, CONSTRAINT_NAME)
+    # This should be done in a separate post migration when dealing with a high traffic table
+    validate_foreign_key(TABLE_NAME, [PARTITION_COLUMN, COLUMN], name: FK_NAME)
   end
 
   def down
-    drop_constraint(TARGET_TABLE_NAME, CONSTRAINT_NAME)
+    with_lock_retries do
+      remove_foreign_key_if_exists(SOURCE_TABLE_NAME, name: FK_NAME)
+    end
   end
 end
 ```
@@ -410,6 +381,48 @@ result in a `Key is still referenced from table ...` error and updating the
 partition column on the source table would raise a
 `Key is not present in table ...` error.
 
+This migration can be automatically generated using:
+
+```shell
+./scripts/partitioning/generate-fk --source source_table_name --target target_table_name
+```
+
+### Step 5 - Swap primary key
+
+Swap the primary key including the partitioning key column. This can be done only after
+including the partition key for all references foreign keys. For example, in a rails migration:
+
+```ruby
+class PreparePrimaryKeyForPartitioning < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  TABLE_NAME = :table_name
+  PRIMARY_KEY = :primary_key
+  OLD_INDEX_NAME = :old_index_name
+  NEW_INDEX_NAME = :new_index_name
+
+  def up
+    swap_primary_key(TABLE_NAME, PRIMARY_KEY, NEW_INDEX_NAME)
+  end
+
+  def down
+    add_concurrent_index(TABLE_NAME, :id, unique: true, name: OLD_INDEX_NAME)
+    add_concurrent_index(TABLE_NAME, [:id, :partition_id], unique: true, name: NEW_INDEX_NAME)
+
+    unswap_primary_key(TABLE_NAME, PRIMARY_KEY, OLD_INDEX_NAME)
+  end
+end
+```
+
+NOTE:
+Do not forget to set the primary key explicitly in your model as `ActiveRecord` does not support composite primary keys.
+
+```ruby
+class Model < ApplicationRecord
+  self.primary_key = :id
+end
+```
+
 ### Step 6 - Create parent table and attach existing table as the initial partition
 
 You can now create the parent table attaching the existing table as the initial
@@ -465,7 +478,7 @@ class ConvertTableToListPartitioning < Gitlab::Database::Migration[2.1]
       table_name: TABLE_NAME,
       partitioning_column: PARTITION_COLUMN,
       parent_table_name: PARENT_TABLE_NAME,
-      initial_partitioning_value: FIRST_PARTITION
+      initial_partitioning_value: FIRST_PARTITION,
       lock_tables: [TABLE_FK, TABLE_NAME]
     )
   end