Merge pull request #15 from collerek/constraints

Constraints

README.md

@@ -374,6 +374,41 @@ assert len(completed) == 3
 
 ```
 
+Since version >=0.3.6 Ormar supports unique constraints on multiple columns
+```python
+import databases
+import ormar
+import sqlalchemy
+
+database = databases.Database("sqlite:///db.sqlite")
+metadata = sqlalchemy.MetaData()
+
+
+class Product(ormar.Model):
+    class Meta:
+        tablename = "products"
+        metadata = metadata
+        database = database
+        # define your constraints in Meta class of the model
+        # it's a list that can contain multiple constraints
+        constraints = [ormar.UniqueColumns("name", "company")]
+
+    id: ormar.Integer(primary_key=True)
+    name: ormar.String(max_length=100)
+    company: ormar.String(max_length=200)
+
+await Product.objects.create(name="Cookies", company="Nestle")
+await Product.objects.create(name="Mars", company="Mars")
+await Product.objects.create(name="Mars", company="Nestle")
+
+
+# will raise error based on backend 
+# (sqlite3.IntegrityError, pymysql.IntegrityError, asyncpg.exceptions.UniqueViolationError)
+await Product.objects.create(name="Mars", company="Mars")       
+
+```
+
+
 ## Data types
 
 The following keyword arguments are supported on all field types.
@@ -394,8 +429,8 @@ All fields are required unless one of the following is set:
   * `primary key` with `autoincrement` - When a column is set to primary key and autoincrement is set on this column. 
 Autoincrement is set by default on int primary keys. 
 
-Available Model Fields:
+Available Model Fields (with required args - optional ones in docs):
-* `String(length)`
+* `String(max_length)`
 * `Text()`
 * `Boolean()`
 * `Integer()`

ormar/__init__.py

@@ -14,6 +14,7 @@ from ormar.fields import (
     Text,
     Time,
     UUID,
+    UniqueColumns,
 )
 from ormar.models import Model
 from ormar.queryset import QuerySet
@@ -51,4 +52,5 @@ __all__ = [
     "RelationType",
     "Undefined",
     "UUID",
+    "UniqueColumns",
 ]

ormar/fields/__init__.py

@@ -1,5 +1,5 @@
 from ormar.fields.base import BaseField
-from ormar.fields.foreign_key import ForeignKey
+from ormar.fields.foreign_key import ForeignKey, UniqueColumns
 from ormar.fields.many_to_many import ManyToMany, ManyToManyField
 from ormar.fields.model_fields import (
     BigInteger,
@@ -33,4 +33,5 @@ __all__ = [
     "ManyToMany",
     "ManyToManyField",
     "BaseField",
+    "UniqueColumns",
 ]

ormar/fields/foreign_key.py

@@ -1,6 +1,7 @@
 from typing import Any, Generator, List, Optional, TYPE_CHECKING, Type, Union
 
 import sqlalchemy
+from sqlalchemy import UniqueConstraint
 
 import ormar  # noqa I101
 from ormar.exceptions import RelationshipInstanceError
@@ -22,6 +23,10 @@ def create_dummy_instance(fk: Type["Model"], pk: Any = None) -> "Model":
     return fk(**init_dict)
 
 
+class UniqueColumns(UniqueConstraint):
+    pass
+
+
 def ForeignKey(  # noqa CFQ002
     to: Type["Model"],
     *,

ormar/models/metaclass.py

@@ -6,6 +6,7 @@ import pydantic
 import sqlalchemy
 from pydantic import BaseConfig
 from pydantic.fields import FieldInfo, ModelField
+from sqlalchemy.sql.schema import ColumnCollectionConstraint
 
 import ormar  # noqa I100
 from ormar import ForeignKey, ModelDefinitionError, Integer  # noqa I100
@@ -27,6 +28,7 @@ class ModelMeta:
     metadata: sqlalchemy.MetaData
     database: databases.Database
     columns: List[sqlalchemy.Column]
+    constraints: List[ColumnCollectionConstraint]
     pkname: str
     model_fields: Dict[
         str, Union[Type[BaseField], Type[ForeignKeyField], Type[ManyToManyField]]
@@ -246,7 +248,10 @@ def populate_meta_sqlalchemy_table_if_required(
 ) -> Type["Model"]:
     if not hasattr(new_model.Meta, "table"):
         new_model.Meta.table = sqlalchemy.Table(
-            new_model.Meta.tablename, new_model.Meta.metadata, *new_model.Meta.columns
+            new_model.Meta.tablename,
+            new_model.Meta.metadata,
+            *new_model.Meta.columns,
+            *new_model.Meta.constraints,
         )
     return new_model
 
@@ -304,6 +309,8 @@ class ModelMetaclass(pydantic.main.ModelMetaclass):
         )
 
         if hasattr(new_model, "Meta"):
+            if not hasattr(new_model.Meta, "constraints"):
+                new_model.Meta.constraints = []
             new_model = populate_meta_orm_model_fields(attrs, new_model)
             new_model = populate_meta_tablename_columns_and_pk(name, new_model)
             new_model = populate_meta_sqlalchemy_table_if_required(new_model)

tests/test_unique_constraints.py

@@ -0,0 +1,60 @@
+import asyncio
+import sqlite3
+
+import asyncpg
+import databases
+import pymysql
+import pytest
+import sqlalchemy
+
+import ormar
+from tests.settings import DATABASE_URL
+
+database = databases.Database(DATABASE_URL, force_rollback=True)
+metadata = sqlalchemy.MetaData()
+
+
+class Product(ormar.Model):
+    class Meta:
+        tablename = "products"
+        metadata = metadata
+        database = database
+        constraints = [ormar.UniqueColumns("name", "company")]
+
+    id: ormar.Integer(primary_key=True)
+    name: ormar.String(max_length=100)
+    company: ormar.String(max_length=200)
+
+
+@pytest.fixture(scope="module")
+def event_loop():
+    loop = asyncio.get_event_loop()
+    yield loop
+    loop.close()
+
+
+@pytest.fixture(autouse=True, scope="module")
+async def create_test_database():
+    engine = sqlalchemy.create_engine(DATABASE_URL)
+    metadata.drop_all(engine)
+    metadata.create_all(engine)
+    yield
+    metadata.drop_all(engine)
+
+
+@pytest.mark.asyncio
+async def test_unique_columns():
+    async with database:
+        async with database.transaction(force_rollback=True):
+            await Product.objects.create(name="Cookies", company="Nestle")
+            await Product.objects.create(name="Mars", company="Mars")
+            await Product.objects.create(name="Mars", company="Nestle")
+
+            with pytest.raises(
+                (
+                    sqlite3.IntegrityError,
+                    pymysql.IntegrityError,
+                    asyncpg.exceptions.UniqueViolationError,
+                )
+            ):
+                await Product.objects.create(name="Mars", company="Mars")