Altering Ondelete and OnUpdate with MigrationManager (#1183)

* fix OnDelete and OnUpdate in migrations

* fix test docstring

* remove for loop

* refactor `add_foreign_key_constraint`

If a `ForeignKey` is passed in, we can get most of the info we need from it, rather than requiring additional params to be passed in.

* add `get_fk_constraint_name`

* rename pk constraint on table renaming in migrations

* `RenamePkConstraint` -> `RenameConstraint`

* skip renaming pk constraint in SQLite

* revert changes

* update param names in `rename_constraint` method

* add docstring

* add `get_fk_constraint_rules` function

So the logic for getting constraints is centralised.

* add TODO for supporting SQLite

* add TODO for moving logic from `generate.py` to `constraints.py`

* remove redundant lookup in `generate.py`

We should refactor a bunch of this - I came across this dict lookup, which isn't required because we can just call the enum directly.

---------

Co-authored-by: Daniel Townsend <[email protected]>

Author: sinisaos

piccolo/apps/migrations/auto/migration_manager.py

@@ -19,6 +19,7 @@
 from piccolo.engine import engine_finder
 from piccolo.query import Query
 from piccolo.query.base import DDL
+from piccolo.query.constraints import get_fk_constraint_name
 from piccolo.schema import SchemaDDLBase
 from piccolo.table import Table, create_table_class, sort_table_classes
 from piccolo.utils.warnings import colored_warning
@@ -423,8 +424,8 @@ async def _print_query(query: t.Union[DDL, Query, SchemaDDLBase]):
 
     async def _run_query(self, query: t.Union[DDL, Query, SchemaDDLBase]):
         """
-        If MigrationManager is not in the preview mode,
-         executes the queries. else, prints the query.
+        If MigrationManager is in preview mode then it just print the query
+        instead of executing it.
         """
         if self.preview:
             await self._print_query(query)
@@ -534,6 +535,39 @@ async def _run_alter_columns(self, backwards: bool = False):
 
                 ###############################################################
 
+                on_delete = params.get("on_delete")
+                on_update = params.get("on_update")
+                if on_delete is not None or on_update is not None:
+                    existing_table = await self.get_table_from_snapshot(
+                        table_class_name=table_class_name,
+                        app_name=self.app_name,
+                    )
+
+                    fk_column = existing_table._meta.get_column_by_name(
+                        alter_column.column_name
+                    )
+
+                    assert isinstance(fk_column, ForeignKey)
+
+                    # First drop the existing foreign key constraint
+                    constraint_name = await get_fk_constraint_name(
+                        column=fk_column
+                    )
+                    await self._run_query(
+                        _Table.alter().drop_constraint(
+                            constraint_name=constraint_name
+                        )
+                    )
+
+                    # Then add a new foreign key constraint
+                    await self._run_query(
+                        _Table.alter().add_foreign_key_constraint(
+                            column=fk_column,
+                            on_delete=on_delete,
+                            on_update=on_update,
+                        )
+                    )
+
                 null = params.get("null")
                 if null is not None:
                     await self._run_query(

piccolo/apps/schema/commands/generate.py

@@ -501,6 +501,8 @@ async def get_fk_triggers(
         Any Table subclass - just used to execute raw queries on the database.
 
     """
+    # TODO - Move this query to `piccolo.query.constraints` or use:
+    # `piccolo.query.constraints.referential_constraints`
     triggers = await table_class.raw(
         (
             "SELECT tc.constraint_name, "
@@ -537,23 +539,6 @@ async def get_fk_triggers(
     )
 
 
-ONDELETE_MAP = {
-    "NO ACTION": OnDelete.no_action,
-    "RESTRICT": OnDelete.restrict,
-    "CASCADE": OnDelete.cascade,
-    "SET NULL": OnDelete.set_null,
-    "SET DEFAULT": OnDelete.set_default,
-}
-
-ONUPDATE_MAP = {
-    "NO ACTION": OnUpdate.no_action,
-    "RESTRICT": OnUpdate.restrict,
-    "CASCADE": OnUpdate.cascade,
-    "SET NULL": OnUpdate.set_null,
-    "SET DEFAULT": OnUpdate.set_default,
-}
-
-
 async def get_constraints(
     table_class: t.Type[Table], tablename: str, schema_name: str = "public"
 ) -> TableConstraints:
@@ -765,8 +750,8 @@ async def create_table_class_from_db(
                     column_name, constraint_table.name
                 )
                 if trigger:
-                    kwargs["on_update"] = ONUPDATE_MAP[trigger.on_update]
-                    kwargs["on_delete"] = ONDELETE_MAP[trigger.on_delete]
+                    kwargs["on_update"] = OnUpdate(trigger.on_update)
+                    kwargs["on_delete"] = OnDelete(trigger.on_delete)
                 else:
                     output_schema.trigger_warnings.append(
                         f"{tablename}.{column_name}"

piccolo/query/constraints.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass
+
+from piccolo.columns import ForeignKey
+from piccolo.columns.base import OnDelete, OnUpdate
+
+
+async def get_fk_constraint_name(column: ForeignKey) -> str:
+    """
+    Checks what the foreign key constraint is called in the database.
+    """
+
+    table = column._meta.table
+
+    if table._meta.db.engine_type == "sqlite":
+        # TODO - add the query for SQLite
+        raise ValueError("SQLite isn't currently supported.")
+
+    schema = table._meta.schema or "public"
+    table_name = table._meta.tablename
+    column_name = column._meta.db_column_name
+
+    constraints = await table.raw(
+        """
+        SELECT
+            kcu.constraint_name AS fk_constraint_name
+        FROM
+            information_schema.referential_constraints AS rc
+        INNER JOIN
+            information_schema.key_column_usage AS kcu
+            ON kcu.constraint_catalog = rc.constraint_catalog
+            AND kcu.constraint_schema = rc.constraint_schema
+            AND kcu.constraint_name = rc.constraint_name
+        WHERE
+            kcu.table_schema = {} AND
+            kcu.table_name = {} AND
+            kcu.column_name = {}
+        """,
+        schema,
+        table_name,
+        column_name,
+    )
+
+    return constraints[0]["fk_constraint_name"]
+
+
+@dataclass
+class ConstraintRules:
+    on_delete: OnDelete
+    on_update: OnUpdate
+
+
+async def get_fk_constraint_rules(column: ForeignKey) -> ConstraintRules:
+    """
+    Checks the constraint rules for this foreign key in the database.
+    """
+    table = column._meta.table
+
+    if table._meta.db.engine_type == "sqlite":
+        # TODO - add the query for SQLite
+        raise ValueError("SQLite isn't currently supported.")
+
+    schema = table._meta.schema or "public"
+    table_name = table._meta.tablename
+    column_name = column._meta.db_column_name
+
+    constraints = await table.raw(
+        """
+        SELECT
+            kcu.constraint_name,
+            kcu.table_name,
+            kcu.column_name,
+            rc.update_rule,
+            rc.delete_rule
+        FROM
+            information_schema.key_column_usage AS kcu
+        INNER JOIN
+            information_schema.referential_constraints AS rc
+            ON kcu.constraint_name = rc.constraint_name
+        WHERE
+            kcu.table_schema = {} AND
+            kcu.table_name = {} AND
+            kcu.column_name = {}
+        """,
+        schema,
+        table_name,
+        column_name,
+    )
+
+    return ConstraintRules(
+        on_delete=OnDelete(constraints[0]["delete_rule"]),
+        on_update=OnUpdate(constraints[0]["update_rule"]),
+    )

piccolo/query/methods/alter.py

@@ -36,6 +36,18 @@ def ddl(self) -> str:
         return f"RENAME TO {self.new_name}"
 
 
+@dataclass
+class RenameConstraint(AlterStatement):
+    __slots__ = ("old_name", "new_name")
+
+    old_name: str
+    new_name: str
+
+    @property
+    def ddl(self) -> str:
+        return f"RENAME CONSTRAINT {self.old_name} TO {self.new_name}"
+
+
 @dataclass
 class AlterColumnStatement(AlterStatement):
     __slots__ = ("column",)
@@ -194,16 +206,17 @@ class AddForeignKeyConstraint(AlterStatement):
         "constraint_name",
         "foreign_key_column_name",
         "referenced_table_name",
+        "referenced_column_name",
         "on_delete",
         "on_update",
     )
 
     constraint_name: str
     foreign_key_column_name: str
     referenced_table_name: str
+    referenced_column_name: str
     on_delete: t.Optional[OnDelete]
     on_update: t.Optional[OnUpdate]
-    referenced_column_name: str = "id"
 
     @property
     def ddl(self) -> str:
@@ -273,8 +286,8 @@ def ddl(self) -> str:
 
 class Alter(DDL):
     __slots__ = (
-        "_add_foreign_key_constraint",
         "_add",
+        "_add_foreign_key_constraint",
         "_drop_constraint",
         "_drop_default",
         "_drop_table",
@@ -288,6 +301,7 @@ class Alter(DDL):
         "_set_null",
         "_set_schema",
         "_set_unique",
+        "_rename_constraint",
     )
 
     def __init__(self, table: t.Type[Table], **kwargs):
@@ -307,6 +321,7 @@ def __init__(self, table: t.Type[Table], **kwargs):
         self._set_null: t.List[SetNull] = []
         self._set_schema: t.List[SetSchema] = []
         self._set_unique: t.List[SetUnique] = []
+        self._rename_constraint: t.List[RenameConstraint] = []
 
     def add_column(self: Self, name: str, column: Column) -> Self:
         """
@@ -372,6 +387,24 @@ def rename_table(self, new_name: str) -> Alter:
         self._rename_table = [RenameTable(new_name=new_name)]
         return self
 
+    def rename_constraint(self, old_name: str, new_name: str) -> Alter:
+        """
+        Rename a constraint on the table::
+
+            >>> await Band.alter().rename_constraint(
+            ...     'old_constraint_name',
+            ...     'new_constraint_name',
+            ... )
+
+        """
+        self._rename_constraint = [
+            RenameConstraint(
+                old_name=old_name,
+                new_name=new_name,
+            )
+        ]
+        return self
+
     def rename_column(
         self, column: t.Union[str, Column], new_name: str
     ) -> Alter:
@@ -488,7 +521,7 @@ def set_length(self, column: t.Union[str, Varchar], length: int) -> Alter:
     def _get_constraint_name(self, column: t.Union[str, ForeignKey]) -> str:
         column_name = AlterColumnStatement(column=column).column_name
         tablename = self.table._meta.tablename
-        return f"{tablename}_{column_name}_fk"
+        return f"{tablename}_{column_name}_fkey"
 
     def drop_constraint(self, constraint_name: str) -> Alter:
         self._drop_constraint.append(
@@ -500,37 +533,58 @@ def drop_foreign_key_constraint(
         self, column: t.Union[str, ForeignKey]
     ) -> Alter:
         constraint_name = self._get_constraint_name(column=column)
-        return self.drop_constraint(constraint_name=constraint_name)
+        self._drop_constraint.append(
+            DropConstraint(constraint_name=constraint_name)
+        )
+        return self
 
     def add_foreign_key_constraint(
         self,
         column: t.Union[str, ForeignKey],
-        referenced_table_name: str,
+        referenced_table_name: t.Optional[str] = None,
+        referenced_column_name: t.Optional[str] = None,
+        constraint_name: t.Optional[str] = None,
         on_delete: t.Optional[OnDelete] = None,
         on_update: t.Optional[OnUpdate] = None,
-        referenced_column_name: str = "id",
     ) -> Alter:
         """
         Add a new foreign key constraint::
 
             >>> await Band.alter().add_foreign_key_constraint(
             ...     Band.manager,
-            ...     referenced_table_name='manager',
             ...     on_delete=OnDelete.cascade
             ... )
 
         """
-        constraint_name = self._get_constraint_name(column=column)
+        constraint_name = constraint_name or self._get_constraint_name(
+            column=column
+        )
         column_name = AlterColumnStatement(column=column).column_name
 
+        if referenced_column_name is None:
+            if isinstance(column, ForeignKey):
+                referenced_column_name = (
+                    column._foreign_key_meta.resolved_target_column._meta.db_column_name  # noqa: E501
+                )
+            else:
+                raise ValueError("Please pass in `referenced_column_name`.")
+
+        if referenced_table_name is None:
+            if isinstance(column, ForeignKey):
+                referenced_table_name = (
+                    column._foreign_key_meta.resolved_references._meta.tablename  # noqa: E501
+                )
+            else:
+                raise ValueError("Please pass in `referenced_table_name`.")
+
         self._add_foreign_key_constraint.append(
             AddForeignKeyConstraint(
                 constraint_name=constraint_name,
                 foreign_key_column_name=column_name,
                 referenced_table_name=referenced_table_name,
+                referenced_column_name=referenced_column_name,
                 on_delete=on_delete,
                 on_update=on_update,
-                referenced_column_name=referenced_column_name,
             )
         )
         return self
@@ -579,9 +633,12 @@ def default_ddl(self) -> t.Sequence[str]:
             i.ddl
             for i in itertools.chain(
                 self._add,
+                self._add_foreign_key_constraint,
                 self._rename_columns,
                 self._rename_table,
+                self._rename_constraint,
                 self._drop,
+                self._drop_constraint,
                 self._drop_default,
                 self._set_column_type,
                 self._set_unique,

piccolo/utils/sync.py

@@ -2,10 +2,14 @@
 
 import asyncio
 import typing as t
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
 
+ReturnType = t.TypeVar("ReturnType")
 
-def run_sync(coroutine: t.Coroutine):
+
+def run_sync(
+    coroutine: t.Coroutine[t.Any, t.Any, ReturnType],
+) -> ReturnType:
     """
     Run the coroutine synchronously - trying to accommodate as many edge cases
     as possible.
@@ -20,5 +24,5 @@ def run_sync(coroutine: t.Coroutine):
     except RuntimeError:
         # An event loop already exists.
         with ThreadPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(asyncio.run, coroutine)
+            future: Future = executor.submit(asyncio.run, coroutine)
             return future.result()