add one-to-one field

encode · aminalaee · Sep 30, 2021 · Sep 27, 2021 · Sep 27, 2021 · Sep 29, 2021
commit 769ea34b70aac6c773c645ae7eafad5a499eee65
diff --git a/orm/__init__.py b/orm/__init__.py
@@ -11,6 +11,7 @@
     Float,
     ForeignKey,
     Integer,
+    OneToOne,
     String,
     Text,
     Time,
@@ -28,13 +29,14 @@
     "Decimal",
     "Enum",
     "Float",
+    "ForeignKey",
     "Integer",
     "JSON",
+    "OneToOne",
     "String",
     "Text",
     "Time",
     "UUID",
-    "ForeignKey",
     "Model",
     "ModelRegistry",
 ]
diff --git a/orm/fields.py b/orm/fields.py
@@ -175,6 +175,25 @@ def expand_relationship(self, value):
         return target(pk=value)
 
 
+class OneToOne(ForeignKey):
+    def get_column(self, name: str) -> sqlalchemy.Column:
+        target = self.target
+        to_field = target.fields[target.pkname]
+
+        column_type = to_field.get_column_type()
+        constraints = [
+            sqlalchemy.schema.ForeignKey(f"{target.tablename}.{target.pkname}"),
+        ]
+
+        return sqlalchemy.Column(
+            name,
+            column_type,
+            *constraints,
+            nullable=self.allow_null,
+            unique=True,
+        )
+
+
 class Enum(ModelField):
     def __init__(self, enum, **kwargs):
         super().__init__(**kwargs)

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
@@ -54,6 +56,23 @@ class Member(orm.Model):
     }
 
 
+class Profile(orm.Model):
+    registry = models
+    fields = {
+        "id": orm.Integer(primary_key=True),
+        "website": orm.String(max_length=100),
+    }
+
+
+class Person(orm.Model):
+    registry = models
+    fields = {
+        "id": orm.Integer(primary_key=True),
+        "email": orm.String(max_length=100),
+        "profile": orm.OneToOne(Profile),
+    }
+
+
 @pytest.fixture(autouse=True, scope="module")
 def create_test_database():
     models.create_all()
@@ -188,3 +207,23 @@ 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
+
+
+async def test_one_to_one_crud():
+    profile = await Profile.objects.create(website="https://encode.io")
+    await Person.objects.create(email="[email protected]", profile=profile)
+
+    person = await Person.objects.get(email="[email protected]")
+    assert person.profile.pk == profile.pk
+    assert not hasattr(person.profile, "website")
+
+    await person.profile.load()
+    assert person.profile.website == "https://encode.io"
+
+    exceptions = (
+        asyncpg.exceptions.UniqueViolationError,
+        pymysql.err.IntegrityError,
+    )
+
+    with pytest.raises(exceptions):
+        await Person.objects.create(email="[email protected]", profile=profile)