try alternate pk handling

PrefectHQ · Nov 5, 2024 · d0c6934 · d0c6934
1 parent 8361e42
commit d0c6934
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/...rations/versions/sqlite/2022_04_23_114831_fd966d4ad99c_rename_block_to_blockbasis_and_.py b/...rations/versions/sqlite/2022_04_23_114831_fd966d4ad99c_rename_block_to_blockbasis_and_.py
@@ -16,22 +16,30 @@
 
 
 def upgrade():
+    # First drop the foreign key constraints
+    with op.batch_alter_table("block", schema=None) as batch_op:
+        batch_op.drop_constraint("fk_block__block_spec_id__block_spec")
+
+    # Then rename the tables
     op.rename_table("block_spec", "block_schema")
     op.rename_table("block", "block_document")
 
+    # Handle indexes and column renames for block_document
     with op.batch_alter_table("block_document", schema=None) as batch_op:
+        # Drop indexes first
         batch_op.drop_index("ix_block__is_default_storage_block")
         batch_op.drop_index("ix_block__name")
         batch_op.drop_index("ix_block__updated")
         batch_op.drop_index("uq_block__spec_id_name")
+
+        # Rename columns
         batch_op.alter_column("block_spec_id", new_column_name="block_schema_id")
         batch_op.alter_column(
             "is_default_storage_block",
             new_column_name="is_default_storage_block_document",
         )
-        batch_op.drop_constraint("fk_block__block_spec_id__block_spec")
-        batch_op.drop_constraint("pk_block_data")
 
+    # Create new indexes
     with op.batch_alter_table("block_document", schema=None) as batch_op:
         batch_op.create_index(
             batch_op.f("ix_block_document__is_default_storage_block_document"),
@@ -48,6 +56,15 @@ def upgrade():
             "uq_block__schema_id_name", ["block_schema_id", "name"], unique=True
         )
 
+        # Re-create foreign key at the end
+        batch_op.create_foreign_key(
+            batch_op.f("fk_block__block_schema_id__block_schema"),
+            "block_schema",
+            ["block_schema_id"],
+            ["id"],
+            ondelete="cascade",
+        )
+
     with op.batch_alter_table("block_schema", schema=None) as batch_op:
         batch_op.drop_index("ix_block_spec__type")
         batch_op.drop_index("ix_block_spec__updated")