diff --git a/.coverage b/.coverage index de993f7..6da7884 100644 Binary files a/.coverage and b/.coverage differ diff --git a/README.md b/README.md index 5c5701a..b8f3f26 100644 --- a/README.md +++ b/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: -* `String(length)` +Available Model Fields (with required args - optional ones in docs): +* `String(max_length)` * `Text()` * `Boolean()` * `Integer()` diff --git a/ormar/__init__.py b/ormar/__init__.py index c3d1ed7..1616879 100644 --- a/ormar/__init__.py +++ b/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", ] diff --git a/ormar/fields/__init__.py b/ormar/fields/__init__.py index 0035a4f..325fcf6 100644 --- a/ormar/fields/__init__.py +++ b/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", ] diff --git a/ormar/fields/foreign_key.py b/ormar/fields/foreign_key.py index 84d9110..957c9e5 100644 --- a/ormar/fields/foreign_key.py +++ b/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"], *, diff --git a/ormar/models/metaclass.py b/ormar/models/metaclass.py index 0ba7840..30d3071 100644 --- a/ormar/models/metaclass.py +++ b/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) diff --git a/tests/test_unique_constraints.py b/tests/test_unique_constraints.py new file mode 100644 index 0000000..7908ae4 --- /dev/null +++ b/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")