adding support for altering foreign key constraints

Author: dantownsend

piccolo/columns/__init__.py

@@ -9,5 +9,5 @@
     ForeignKey,
     UUID,
 )
-from .base import Column, Selectable  # noqa
+from .base import Column, Selectable, OnDelete  # noqa
 from .combination import And, Or, Where  # noqa

piccolo/columns/base.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
+from enum import Enum
 import typing as t
 
 from piccolo.columns.operators import (
@@ -28,6 +29,24 @@
     from .column_types import ForeignKey  # noqa
 
 
+class OnDelete(str, Enum):
+    cascade = "CASCADE"
+    restrict = "RESTRICT"
+    no_action = "NO ACTION"
+    set_null = "SET NULL"
+    set_default = "SET DEFAULT"
+
+
+OnUpdate = OnDelete
+
+
+@dataclass
+class ForeignKeyMeta:
+    references: t.Type[Table]
+    on_delete: OnDelete
+    proxy_columns: t.List[Column] = field(default_factory=list)
+
+
 @dataclass
 class ColumnMeta:
     """
@@ -46,7 +65,7 @@ class ColumnMeta:
 
     # Set by the Table Metaclass:
     _name: t.Optional[str] = None
-    _table: t.Optional[Table] = None
+    _table: t.Optional[t.Type[Table]] = None
 
     # Used by Foreign Keys:
     call_chain: t.List["ForeignKey"] = field(default_factory=lambda: [])
@@ -219,10 +238,13 @@ def querystring(self) -> QueryString:
         if not self._meta.null:
             query += " NOT NULL"
 
-        foreign_key_meta = getattr(self, "_foreign_key_meta", None)
+        foreign_key_meta: t.Optional[ForeignKeyMeta] = getattr(
+            self, "_foreign_key_meta", None
+        )
         if foreign_key_meta:
-            references = foreign_key_meta.references
-            query += f" REFERENCES {references._meta.tablename}"
+            tablename = foreign_key_meta.references._meta.tablename
+            on_delete = foreign_key_meta.on_delete.value
+            query += f" REFERENCES {tablename} ON DELETE {on_delete}"
 
         return QueryString(query)
 

piccolo/columns/column_types.py

@@ -2,11 +2,10 @@
 import copy
 from dataclasses import dataclass, field
 from datetime import datetime
-from enum import Enum
 import typing as t
 import uuid
 
-from piccolo.columns.base import Column
+from piccolo.columns.base import Column, ForeignKeyMeta, OnDelete
 from piccolo.querystring import Unquoted
 
 if t.TYPE_CHECKING:
@@ -126,18 +125,6 @@ def __init__(self, default: bool = False, **kwargs) -> None:
         super().__init__(**kwargs)
 
 
-@dataclass
-class ForeignKeyMeta:
-    references: t.Type[Table]
-    proxy_columns: t.List[Column] = field(default_factory=list)
-
-
-class OnDelete(str, Enum):
-    cascade = "CASCADE"
-    restrict = "RESTRICT"
-    no_action = "NO ACTION"
-
-
 class ForeignKey(Integer):
     """
     Returns an integer, representing the referenced row's ID.
@@ -169,10 +156,17 @@ class ForeignKey(Integer):
 
     column_type = "INTEGER"
 
-    def __init__(self, references: t.Type[Table], **kwargs) -> None:
+    def __init__(
+        self,
+        references: t.Type[Table],
+        on_delete: OnDelete = OnDelete.cascade,
+        **kwargs,
+    ) -> None:
         kwargs.update({"references": references})
         super().__init__(**kwargs)
-        self._foreign_key_meta = ForeignKeyMeta(references=references)
+        self._foreign_key_meta = ForeignKeyMeta(
+            references=references, on_delete=on_delete
+        )
 
         # Allow columns on the referenced table to be accessed via auto
         # completion.

piccolo/columns/combination.py

@@ -28,7 +28,7 @@ def __init__(self, first: Combinable, second: Combinable) -> None:
         self.second = second
 
     @property
-    def querystring(self):
+    def querystring(self) -> QueryString:
         return QueryString(
             "({} " + self.operator + " {})",
             self.first.querystring,

piccolo/engine/postgres.py

@@ -49,7 +49,8 @@ async def __anext__(self):
     async def __aenter__(self):
         self._transaction = self.connection.transaction()
         await self._transaction.start()
-        template, template_args = self.query.querystring.compile_string()
+        querystring = self.query.querystring[0]
+        template, template_args = querystring.compile_string()
 
         self._cursor = await self.connection.cursor(template, *template_args)
         return self
@@ -84,10 +85,11 @@ def add(self, *query: Query):
     async def _run_queries(self, connection):
         async with connection.transaction():
             for query in self.queries:
-                _query, args = query.querystring.compile_string(
-                    engine_type=self.engine.engine_type
-                )
-                await connection.execute(_query, *args)
+                for querystring in query.querystring:
+                    _query, args = querystring.compile_string(
+                        engine_type=self.engine.engine_type
+                    )
+                    await connection.execute(_query, *args)
 
         self.queries = []
 

piccolo/engine/sqlite.py

@@ -41,7 +41,8 @@ async def __anext__(self):
         return response
 
     async def __aenter__(self):
-        template, template_args = self.query.querystring.compile_string()
+        querystring = self.query.querystring[0]
+        template, template_args = querystring.compile_string()
 
         self._cursor = await self.connection.execute(template, *template_args)
         return self

piccolo/query/base.py

@@ -90,11 +90,20 @@ async def run(self, in_pool=True):
                 f"Table {self.table._meta.tablename} has no db defined in _meta"
             )
 
-        results = await engine.run_querystring(
-            self.querystring, in_pool=in_pool
-        )
-
-        return await self._process_results(results)
+        if len(self.querystring) == 1:
+            results = await engine.run_querystring(
+                self.querystring[0], in_pool=in_pool
+            )
+            return await self._process_results(results)
+        else:
+            responses = []
+            # TODO - run in a transaction
+            for querystring in self.querystring:
+                results = await engine.run_querystring(
+                    querystring, in_pool=in_pool
+                )
+                responses.append(await self._process_results(results))
+            return responses
 
     def run_sync(self, *args, **kwargs):
         """
@@ -112,19 +121,19 @@ async def response_handler(self, response):
     ###########################################################################
 
     @property
-    def sqlite_querystring(self) -> QueryString:
+    def sqlite_querystring(self) -> t.Sequence[QueryString]:
         raise NotImplementedError
 
     @property
-    def postgres_querystring(self) -> QueryString:
+    def postgres_querystring(self) -> t.Sequence[QueryString]:
         raise NotImplementedError
 
     @property
-    def default_querystring(self) -> QueryString:
+    def default_querystring(self) -> t.Sequence[QueryString]:
         raise NotImplementedError
 
     @property
-    def querystring(self) -> QueryString:
+    def querystring(self) -> t.Sequence[QueryString]:
         """
         Calls the correct underlying method, depending on the current engine.
         """
@@ -147,4 +156,4 @@ def querystring(self) -> QueryString:
     ###########################################################################
 
     def __str__(self) -> str:
-        return self.querystring.__str__()
+        return "; ".join([i.__str__() for i in self.querystring])