encode · aminalaee · Sep 29, 2021 · Sep 26, 2021 · Sep 26, 2021 · Sep 28, 2021
diff --git a/docs/relationships.md b/docs/relationships.md
@@ -1,5 +1,7 @@
 ## ForeignKey
 
+### Defining and querying relationships
+
 ORM supports loading and filtering across foreign keys.
 
 Let's say you have the following models defined:
@@ -79,3 +81,50 @@ assert len(tracks) == 2
 tracks = Track.objects.filter(album__name__iexact="fantasies")
 assert len(tracks) == 2
 ```
+
+### ForeignKey constraints
+
+`ForeigknKey` supports specfiying a constraint through `on_delete` argument.
+
+This will result in a SQL `ON DELETE` query being generated when the referenced object is removed.
+
+With the following definition:
+
+```python
+class Album(orm.Model):
+    tablename = "albums"
+    registry = models
+    fields = {
+        "id": orm.Integer(primary_key=True),
+        "name": orm.String(max_length=100),
+    }
+
+
+class Track(orm.Model):
+    tablename = "tracks"
+    registry = models
+    fields = {
+        "id": orm.Integer(primary_key=True),
+        "album": orm.ForeignKey(Album, on_delete=orm.CASCADE),
+        "title": orm.String(max_length=100),
+    }
+```
+
+`Track` model defines `orm.ForeignKey(Album, on_delete=orm.CASCADE)` so whenever an `Album` object is removed,
+all `Track` objects referencing that `Album`  will also be removed.
+
+Available options for `on_delete` are:
+
+* `CASCADE`
+
+This will remove all referencing objects.
+
+* `RESTRICT`
+
+This will restrict removing referenced object, if there are objects referencing it.
+A database driver exception will be raised.
+
+* `SET NULL`
+
+This will set referencing objects `ForeignKey` column to `NULL`.
+The `ForeignKey` defined here should also have `allow_null=True`.
diff --git a/orm/__init__.py b/orm/__init__.py
@@ -1,3 +1,4 @@
+from orm.constants import CASCADE, RESTRICT, SET_NULL
 from orm.exceptions import MultipleMatches, NoMatch
 from orm.fields import (
     JSON,
@@ -19,6 +20,9 @@
 
 __version__ = "0.2.1"
 __all__ = [
+    "CASCADE",
+    "RESTRICT",
+    "SET_NULL",
     "NoMatch",
     "MultipleMatches",
     "BigInteger",

diff --git a/orm/constants.py b/orm/constants.py
@@ -0,0 +1,3 @@
+CASCADE = "CASCADE"
+RESTRICT = "RESTRICT"
+SET_NULL = "SET NULL"
diff --git a/orm/fields.py b/orm/fields.py
@@ -137,9 +137,12 @@ class ForeignKeyValidator(typesystem.Field):
         def validate(self, value):
             return value.pk
 
-    def __init__(self, to, allow_null: bool = False):
+    def __init__(
+        self, to, allow_null: bool = False, on_delete: typing.Optional[str] = None
+    ):
         super().__init__(allow_null=allow_null)
         self.to = to
+        self.on_delete = on_delete
 
     @property
     def target(self):
@@ -159,7 +162,10 @@ def get_column(self, name: str) -> sqlalchemy.Column:
 
         column_type = to_field.get_column_type()
         constraints = [
-            sqlalchemy.schema.ForeignKey(f"{target.tablename}.{target.pkname}")
+            sqlalchemy.schema.ForeignKey(
+                f"{target.tablename}.{target.pkname}",
+                ondelete=self.on_delete,
+            )
         ]
         return sqlalchemy.Column(
             name,

diff --git a/tests/test_foreignkey.py b/tests/test_foreignkey.py
@@ -1,4 +1,6 @@
+import asyncpg
 import databases
+import pymysql
 import pytest
 
 import orm
@@ -22,7 +24,7 @@ class Track(orm.Model):
     registry = models
     fields = {
         "id": orm.Integer(primary_key=True),
-        "album": orm.ForeignKey("Album"),
+        "album": orm.ForeignKey("Album", on_delete=orm.CASCADE),
         "title": orm.String(max_length=100),
         "position": orm.Integer(),
     }
@@ -40,7 +42,7 @@ class Team(orm.Model):
     registry = models
     fields = {
         "id": orm.Integer(primary_key=True),
-        "org": orm.ForeignKey(Organisation),
+        "org": orm.ForeignKey(Organisation, on_delete=orm.RESTRICT),
         "name": orm.String(max_length=100),
     }
 
@@ -49,7 +51,7 @@ class Member(orm.Model):
     registry = models
     fields = {
         "id": orm.Integer(primary_key=True),
-        "team": orm.ForeignKey(Team),
+        "team": orm.ForeignKey(Team, on_delete=orm.SET_NULL, allow_null=True),
         "email": orm.String(max_length=100),
     }
 
@@ -188,3 +190,42 @@ async def test_queryset_update_with_fk():
     await Track.objects.filter(album=malibu).update(album=wall)
     assert await Track.objects.filter(album=malibu).count() == 0
     assert await Track.objects.filter(album=wall).count() == 1
+
+
+@pytest.mark.skipif(database.url.dialect == "sqlite", reason="Not supported on SQLite")
+async def test_on_delete_cascade():
+    album = await Album.objects.create(name="The Wall")
+    await Track.objects.create(album=album, title="Hey You", position=1)
+    await Track.objects.create(album=album, title="Breathe", position=2)
+
+    assert await Track.objects.count() == 2
+
+    await album.delete()
+
+    assert await Track.objects.count() == 0
+
+
+@pytest.mark.skipif(database.url.dialect == "sqlite", reason="Not supported on SQLite")
+async def test_on_delete_retstrict():
+    organisation = await Organisation.objects.create(ident="Encode")
+    await Team.objects.create(org=organisation, name="Maintainers")
+
+    exceptions = (
+        asyncpg.exceptions.ForeignKeyViolationError,
+        pymysql.err.IntegrityError,
+    )
+
+    with pytest.raises(exceptions):
+        await organisation.delete()
+
+
+@pytest.mark.skipif(database.url.dialect == "sqlite", reason="Not supported on SQLite")
+async def test_on_delete_set_null():
+    organisation = await Organisation.objects.create(ident="Encode")
+    team = await Team.objects.create(org=organisation, name="Maintainers")
+    await Member.objects.create(email="member@encode.io", team=team)
+
+    await team.delete()
+
+    member = await Member.objects.first()
+    assert member.team.pk is None