Merge pull request #34 from codemation/add-unique-model-constraint

Add unique model constraint

README.md

@@ -26,11 +26,11 @@ $ pip install pydbantic[postgres]
 
 ```python
 from typing import List, Optional, Union
-from pydbantic import DataBaseModel, PrimaryKey
+from pydbantic import DataBaseModel, PrimaryKey, Unique
 
 class Department(DataBaseModel):
     department_id: str = PrimaryKey()
-    name: str
+    name: str = Unique()
     company: str
     is_sensitive: bool = False
     positions: List[Optional['Positions']] = []  # One to Many

docs/model-usage.md

@@ -6,7 +6,7 @@
 ```python
 from typing import List, Optional
 from pydantic import BaseModel, Field
-from pydbantic import DataBaseModel, PrimaryKey
+from pydbantic import DataBaseModel, PrimaryKey, Unique
 
 class Coordinates(BaseModel):
 
@@ -30,6 +30,7 @@ class EmployeeInfo(DataBaseModel):
 
 class Employee(DataBaseModel):
     id: str = PrimaryKey()
+    bio_id: str = Unique()
     employee_info: EmployeeInfo
     position: Positions
     salary: float

pydbantic/__init__.py

@@ -1,2 +1,2 @@
 from .database import Database
-from .core import DataBaseModel, PrimaryKey, Default
+from .core import DataBaseModel, PrimaryKey, Default, Unique

pydbantic/core.py

@@ -71,6 +71,11 @@ def Default(default=...):
         return Field(default_factory=default)
     return Field(default=default)
 
+def Unique(default=...):
+    if isinstance(default, type(lambda x: x)):
+        return Field(default_factory=default, unique=True)
+    return Field(unique=True)
+
 class DataBaseModelCondition:
     def __init__(
         self, 
@@ -420,6 +425,7 @@ def convert_fields_to_columns(
             include = [f for f in cls.__fields__]
 
         primary_key = None
+        unique_keys = set()
         array_fields = set()
 
         field_properties = {}
@@ -442,6 +448,8 @@ def convert_fields_to_columns(
                 if primary_key:
                     raise Exception(f"Duplicate Primary Key Specified for {cls.__name__}")
                 primary_key = field_property
+            if 'unique' in config:
+                unique_keys.add(field_property)
             if 'type' in config and config['type'] == 'array' and not primary_key == field_property:
                 array_fields.add(field_property)
 
@@ -525,7 +533,8 @@ def convert_fields_to_columns(
                         *column_type_config['args'], 
                         **column_type_config['kwargs']
                     ),
-                    primary_key = field['name'] == primary_key
+                    primary_key = field['name'] == primary_key,
+                    unique = field['name'] in unique_keys,
                 )
             )
 

requirements-test.txt

@@ -1,5 +1,5 @@
-pytest
-pytest-asyncio
+pytest==6.2.5
+pytest-asyncio==0.16.0
 uvloop==0.16.0
 nest-asyncio==1.5.4
 fastapi==0.68.1

tests/models/__init__.py

@@ -1,24 +1,28 @@
 from uuid import uuid4
 from datetime import datetime
 from typing import List, Optional, Union
-from pydbantic import DataBaseModel, PrimaryKey
+from pydbantic import DataBaseModel, PrimaryKey, Unique
+
+def uuid_str():
+    return str(uuid4())
 
 class Department(DataBaseModel):
-    department_id: str 
+    department_id: str = PrimaryKey()
     name: str
     company: str
     is_sensitive: bool = False
     positions: List[Optional['Positions']] = []
 
 class Positions(DataBaseModel):
-    position_id: str
+    position_id: str = PrimaryKey()
     name: str
     department: Department = None
     employees: List[Optional['Employee']] = []
 
 class EmployeeInfo(DataBaseModel):
     #__renamed__= [{'old_name': 'first_name', 'new_name': 'first_names'}]
-    ssn: str
+    ssn: str = PrimaryKey()
+    bio_id: str = Unique(default=uuid_str)
     first_name: str
     last_name: str
     address: str

tests/test_database.py

@@ -1,7 +1,7 @@
 import time
 import pytest
 from pydbantic import Database
-from tests.models import Employee, Positions, Department
+from tests.models import Employee, EmployeeInfo, Positions, Department
 
 @pytest.mark.asyncio
 async def test_database(db_url):
@@ -40,6 +40,7 @@ async def test_database(db_url):
 
     employee = Employee(**new_employee)
     await employee.insert()
+    emp_info = await EmployeeInfo.filter(bio_id=employee.employee_info.bio_id)
 
     result = await Employee.select('*', where={'employee_id': employee.employee_id})
 
@@ -57,16 +58,28 @@ async def test_database(db_url):
         i = int(time.time()) + i
         position = new_employee['position'][0]
         position['employee_id'] = f'p{i}'
-        row = Employee(
+        emp = Employee(
             employee_id=f'a{i}',
             position=[position],
             is_employed=True, 
             salary=new_employee['salary'], 
             employee_info=new_employee['employee_info'],
         )
-        await row.insert()
-    employees = await Employee.select('*', where={'employee_id': row.employee_id})
-
+        await emp.insert()
+        emp = await Employee.get(employee_id=f'a{i}')
+        emp_info: EmployeeInfo = await EmployeeInfo.filter(bio_id=emp.employee_info.bio_id)
+        assert len(emp_info) == 1
+
+        try:
+            emp_info = emp_info[0]
+            # this should fail due to unique constraint on bio_id
+            emp_info.ssn = 1234567890
+            await emp_info.save()
+
+            assert False, f"This should have thrown an Integrity Exception for Unique field"
+        except Exception:
+            pass
+
     filtered_employee = await Employee.filter(is_employed=True)
     assert len(filtered_employee) > 0