diff --git a/.flake8 b/.flake8 index abf40d7..59b6009 100644 --- a/.flake8 +++ b/.flake8 @@ -1,5 +1,5 @@ [flake8] -ignore = ANN101, ANN102, W503, S101, CFQ004 +ignore = ANN101, ANN102, W503, S101, CFQ004, S311 max-complexity = 8 max-line-length = 88 import-order-style = pycharm diff --git a/docs/models/index.md b/docs/models/index.md index 428b389..c99a71f 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -266,6 +266,44 @@ But for now you cannot change the ManyToMany column names as they go through oth --8<-- "../docs_src/models/docs010.py" ``` +## Overwriting the default QuerySet + +If you want to customize the queries run by ormar you can define your own queryset class (that extends the ormar `QuerySet`) in your model class, default one is simply the `QuerySet` + +You can provide a new class in `Meta` configuration of your class as `queryset_class` parameter. + +```python +import ormar +from ormar.queryset.queryset import QuerySet +from fastapi import HTTPException + + +class MyQuerySetClass(QuerySet): + + async def first_or_404(self, *args, **kwargs): + entity = await self.get_or_none(*args, **kwargs) + if entity is None: + # in fastapi or starlette + raise HTTPException(404) + + +class Book(ormar.Model): + + class Meta(ormar.ModelMeta): + metadata = metadata + database = database + tablename = "book" + queryset_class = MyQuerySetClass + + id: int = ormar.Integer(primary_key=True) + name: str = ormar.String(max_length=32) + + +# when book not found, raise `404` in your view. +book = await Book.objects.first_or_404(name="123") + +``` + ### Type Hints & Legacy Before version 0.4.0 `ormar` supported only one way of defining `Fields` on a `Model` using python type hints as pydantic. diff --git a/docs/models/inheritance.md b/docs/models/inheritance.md index 5282247..d0f4be6 100644 --- a/docs/models/inheritance.md +++ b/docs/models/inheritance.md @@ -570,4 +570,4 @@ class Category(CreateDateFieldsModel, AuditCreateModel): code: int = ormar.Integer() ``` -That way you can inherit from both create and update classes if needed, and only one of them otherwise. \ No newline at end of file +That way you can inherit from both create and update classes if needed, and only one of them otherwise. diff --git a/docs/signals.md b/docs/signals.md index 3a85d52..50bfd0a 100644 --- a/docs/signals.md +++ b/docs/signals.md @@ -234,7 +234,7 @@ Send for `Model.relation_name.remove()` method for `ManyToMany` relations and re ### post_bulk_update -`post_bulk_update(sender: Type["Model"], instances: List["Model"], **kwargs), +`post_bulk_update(sender: Type["Model"], instances: List["Model"], **kwargs)`, Send for `Model.objects.bulk_update(List[objects])` method. diff --git a/ormar/__init__.py b/ormar/__init__.py index fd52f2b..4187d60 100644 --- a/ormar/__init__.py +++ b/ormar/__init__.py @@ -25,12 +25,12 @@ except ImportError: # pragma: no cover from importlib_metadata import version # type: ignore from ormar.protocols import QuerySetProtocol, RelationProtocol # noqa: I100 from ormar.decorators import ( # noqa: I100 + post_bulk_update, post_delete, post_relation_add, post_relation_remove, post_save, post_update, - post_bulk_update, pre_delete, pre_relation_add, pre_relation_remove, diff --git a/ormar/decorators/signals.py b/ormar/decorators/signals.py index f70b5b7..3e10c1d 100644 --- a/ormar/decorators/signals.py +++ b/ormar/decorators/signals.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Type, TYPE_CHECKING, Union +from typing import Callable, List, TYPE_CHECKING, Type, Union if TYPE_CHECKING: # pragma: no cover from ormar import Model diff --git a/ormar/models/descriptors/descriptors.py b/ormar/models/descriptors/descriptors.py index 3d168e6..35f4194 100644 --- a/ormar/models/descriptors/descriptors.py +++ b/ormar/models/descriptors/descriptors.py @@ -1,10 +1,7 @@ import base64 from typing import Any, TYPE_CHECKING, Type -try: - import orjson as json -except ImportError: # pragma: no cover - import json # type: ignore +from ormar.fields.parsers import encode_json if TYPE_CHECKING: # pragma: no cover from ormar import Model @@ -40,9 +37,7 @@ class JsonDescriptor: return value def __set__(self, instance: "Model", value: Any) -> None: - if not isinstance(value, str): - value = json.dumps(value) - value = value.decode("utf-8") if isinstance(value, bytes) else value + value = encode_json(value) instance._internal_set(self.name, value) instance.set_save_status(False) diff --git a/ormar/models/helpers/models.py b/ormar/models/helpers/models.py index aa32fee..4f55aeb 100644 --- a/ormar/models/helpers/models.py +++ b/ormar/models/helpers/models.py @@ -47,18 +47,18 @@ def populate_default_options_values( # noqa: CCR001 :param model_fields: dict of model fields :type model_fields: Union[Dict[str, type], Dict] """ - if not hasattr(new_model.Meta, "constraints"): - new_model.Meta.constraints = [] - if not hasattr(new_model.Meta, "model_fields"): - new_model.Meta.model_fields = model_fields - if not hasattr(new_model.Meta, "abstract"): - new_model.Meta.abstract = False - if not hasattr(new_model.Meta, "extra"): - new_model.Meta.extra = Extra.forbid - if not hasattr(new_model.Meta, "orders_by"): - new_model.Meta.orders_by = [] - if not hasattr(new_model.Meta, "exclude_parent_fields"): - new_model.Meta.exclude_parent_fields = [] + defaults = { + "queryset_class": ormar.QuerySet, + "constraints": [], + "model_fields": model_fields, + "abstract": False, + "extra": Extra.forbid, + "orders_by": [], + "exclude_parent_fields": [], + } + for key, value in defaults.items(): + if not hasattr(new_model.Meta, key): + setattr(new_model.Meta, key, value) if any( is_field_an_forward_ref(field) for field in new_model.Meta.model_fields.values() diff --git a/ormar/models/metaclass.py b/ormar/models/metaclass.py index cc1dede..bc22727 100644 --- a/ormar/models/metaclass.py +++ b/ormar/models/metaclass.py @@ -85,6 +85,7 @@ class ModelMeta: orders_by: List[str] exclude_parent_fields: List[str] extra: Extra + queryset_class: Type[QuerySet] def add_cached_properties(new_model: Type["Model"]) -> None: @@ -622,7 +623,7 @@ class ModelMetaclass(pydantic.main.ModelMetaclass): f"ForwardRefs. \nBefore using the model you " f"need to call update_forward_refs()." ) - return QuerySet(model_cls=cls) + return cls.Meta.queryset_class(model_cls=cls) def __getattr__(self, item: str) -> Any: """ diff --git a/ormar/models/mixins/save_mixin.py b/ormar/models/mixins/save_mixin.py index 89ddbfe..7617a26 100644 --- a/ormar/models/mixins/save_mixin.py +++ b/ormar/models/mixins/save_mixin.py @@ -12,15 +12,11 @@ from typing import ( cast, ) -try: - import orjson as json -except ImportError: # pragma: no cover - import json # type: ignore - import pydantic import ormar # noqa: I100, I202 from ormar.exceptions import ModelPersistenceError +from ormar.fields.parsers import encode_json from ormar.models.mixins import AliasMixin from ormar.models.mixins.relation_mixin import RelationMixin @@ -207,8 +203,8 @@ class SavePrepareMixin(RelationMixin, AliasMixin): :rtype: Dict """ for key, value in model_dict.items(): - if key in cls._json_fields and not isinstance(value, str): - model_dict[key] = json.dumps(value) + if key in cls._json_fields: + model_dict[key] = encode_json(value) return model_dict @classmethod diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index c576547..beae8d7 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -30,9 +30,9 @@ except ImportError: # pragma: no cover import ormar # noqa I100 from ormar import MultipleMatches, NoMatch from ormar.exceptions import ( + ModelListEmptyError, ModelPersistenceError, QueryDefinitionError, - ModelListEmptyError, ) from ormar.queryset import FieldAccessor, FilterQuery, SelectAction from ormar.queryset.actions.order_action import OrderAction @@ -1065,8 +1065,8 @@ class QuerySet(Generic[T]): :type objects: List[Model] """ ready_objects = [obj.prepare_model_to_save(obj.dict()) for obj in objects] - - # don't use execute_many, as in databases it's executed in a loop + + # don't use execute_many, as in databases it's executed in a loop # instead of using execute_many from drivers expr = self.table.insert().values(ready_objects) await self.database.execute(expr) diff --git a/ormar/signals/signal.py b/ormar/signals/signal.py index 2b76787..868b494 100644 --- a/ormar/signals/signal.py +++ b/ormar/signals/signal.py @@ -1,6 +1,6 @@ import asyncio import inspect -from typing import Any, Callable, Dict, Tuple, Type, TYPE_CHECKING, Union +from typing import Any, Callable, Dict, TYPE_CHECKING, Tuple, Type, Union from ormar.exceptions import SignalDefinitionError diff --git a/poetry.lock b/poetry.lock index 51b35fc..5aef2eb 100644 --- a/poetry.lock +++ b/poetry.lock @@ -789,7 +789,7 @@ i18n = ["babel (>=2.9.0)"] [[package]] name = "mkdocs-material" -version = "8.1.6" +version = "8.1.7" description = "A Material Design theme for MkDocs" category = "dev" optional = false @@ -938,12 +938,20 @@ category = "dev" optional = false python-versions = ">=3.6.0,<4.0.0" +[[package]] +name = "orjson" +version = "3.6.1" +description = "Fast, correct Python JSON library supporting dataclasses, datetimes, and numpy" +category = "dev" +optional = false +python-versions = ">=3.6" + [[package]] name = "orjson" version = "3.6.5" description = "Fast, correct Python JSON library supporting dataclasses, datetimes, and numpy" category = "main" -optional = true +optional = false python-versions = ">=3.7" [[package]] @@ -1002,7 +1010,7 @@ testing = ["pytest", "pytest-benchmark"] [[package]] name = "pre-commit" -version = "2.16.0" +version = "2.17.0" description = "A framework for managing and maintaining multi-language pre-commit hooks." category = "dev" optional = false @@ -1369,7 +1377,7 @@ python-versions = "*" [[package]] name = "types-cryptography" -version = "3.3.12" +version = "3.3.14" description = "Typing stubs for cryptography" category = "dev" optional = false @@ -1389,7 +1397,7 @@ python-versions = "*" [[package]] name = "types-enum34" -version = "1.1.2" +version = "1.1.8" description = "Typing stubs for enum34" category = "dev" optional = false @@ -1397,7 +1405,7 @@ python-versions = "*" [[package]] name = "types-ipaddress" -version = "1.0.2" +version = "1.0.7" description = "Typing stubs for ipaddress" category = "dev" optional = false @@ -1557,7 +1565,7 @@ sqlite = [] [metadata] lock-version = "1.1" python-versions = "^3.6.2" -content-hash = "03ae3c3b1c4029b9ddd4df1c99c6c53d346863e2ff28b997286e9c076b2d4fff" +content-hash = "c13ac746ee85d4f2d04fce44b0218c3981a27c09c75d73a02b34db9e1a2f7ca4" [metadata.files] aiocontextvars = [ @@ -2083,8 +2091,8 @@ mkdocs = [ {file = "mkdocs-1.2.3.tar.gz", hash = "sha256:89f5a094764381cda656af4298727c9f53dc3e602983087e1fe96ea1df24f4c1"}, ] mkdocs-material = [ - {file = "mkdocs-material-8.1.6.tar.gz", hash = "sha256:12eb74faf018950f51261a773f9bea12cc296ec4bdbb2c8cf74102ee35b6df79"}, - {file = "mkdocs_material-8.1.6-py2.py3-none-any.whl", hash = "sha256:b2303413e3154502759f90ee2720b451be8855f769c385d8fb06a93ce54aafe2"}, + {file = "mkdocs-material-8.1.7.tar.gz", hash = "sha256:16a50e3f08f1e41bdc3115a00045d174e7fd8219c26917d0d0b48b2cc9d5a18f"}, + {file = "mkdocs_material-8.1.7-py2.py3-none-any.whl", hash = "sha256:71bcac6795b22dcf8bab8b9ad3fe462242c4cd05d28398281902425401f23462"}, ] mkdocs-material-extensions = [ {file = "mkdocs-material-extensions-1.0.3.tar.gz", hash = "sha256:bfd24dfdef7b41c312ede42648f9eb83476ea168ec163b613f9abd12bbfddba2"}, @@ -2163,6 +2171,33 @@ nodeenv = [ {file = "nr.utils.re-0.3.1.tar.gz", hash = "sha256:7e4539313620f87ef5361f0b69f85d825247b833a10be8445db75ade54611091"}, ] orjson = [ + {file = "orjson-3.6.1-cp310-cp310-manylinux_2_24_aarch64.whl", hash = "sha256:ee75753d1929ddd84702ac75d146083c501c7b1978acb35561a25093446b7f5a"}, + {file = "orjson-3.6.1-cp310-cp310-manylinux_2_24_x86_64.whl", hash = "sha256:52bd32016e9cc55ca89ce5678196e5d55fec72ded9d9bd2e1e10745b9144562f"}, + {file = "orjson-3.6.1-cp36-cp36m-macosx_10_7_x86_64.whl", hash = "sha256:3954406cc8890f08632dd6f2fabc11fd93003ff843edc4aa1c02bfe326d8e7db"}, + {file = "orjson-3.6.1-cp36-cp36m-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:8e4052206bc63267d7a578e66d6f1bf560573a408fbd97b748f468f7109159e9"}, + {file = "orjson-3.6.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:97dc56a8edbe5c3df807b3fcf67037184938262475759ac3038f1287909303ec"}, + {file = "orjson-3.6.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bcf28d08fd0e22632e165c6961054a2e2ce85fbf55c8f135d21a391b87b8355a"}, + {file = "orjson-3.6.1-cp36-cp36m-manylinux_2_24_x86_64.whl", hash = "sha256:0f707c232d1d99d9812b81aac727be5185e53df7c7847dabcbf2d8888269933c"}, + {file = "orjson-3.6.1-cp36-none-win_amd64.whl", hash = "sha256:6c32b0fdc96d22a9eb086afc362e51e9be8433741d73c1b5850b929815aa722c"}, + {file = "orjson-3.6.1-cp37-cp37m-macosx_10_7_x86_64.whl", hash = "sha256:a173b436d43707ba8e6d11d073b95f0992b623749fd135ebd04489f6b656aeb9"}, + {file = "orjson-3.6.1-cp37-cp37m-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:2c7ba86aff33ca9cfd5f00f3a2a40d7d40047ad848548cb13885f60f077fd44c"}, + {file = "orjson-3.6.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:33e0be636962015fbb84a203f3229744e071e1ef76f48686f76cb639bdd4c695"}, + {file = "orjson-3.6.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa7f9c3e8db204ff9e9a3a0ff4558c41f03f12515dd543720c6b0cebebcd8cbc"}, + {file = "orjson-3.6.1-cp37-cp37m-manylinux_2_24_x86_64.whl", hash = "sha256:a89c4acc1cd7200fd92b68948fdd49b1789a506682af82e69a05eefd0c1f2602"}, + {file = "orjson-3.6.1-cp37-none-win_amd64.whl", hash = "sha256:a4810a875f56e0c0eb521fd84ab084f75026e5be8fd2163d08216796f473b552"}, + {file = "orjson-3.6.1-cp38-cp38-macosx_10_7_x86_64.whl", hash = "sha256:310d95d3abfe1d417fcafc592a1b6ce4b5618395739d701eb55b1361a0d93391"}, + {file = "orjson-3.6.1-cp38-cp38-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:62fb8f8949d70cefe6944818f5ea410520a626d5a4b33a090d5a93a6d7c657a3"}, + {file = "orjson-3.6.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b9eb1d8b15779733cf07df61d74b3a8705fe0f0156392aff1c634b83dba19b8a"}, + {file = "orjson-3.6.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4723120784a50cbf3defb65b5eb77ea0b17d3633ade7ce2cd564cec954fd6fd0"}, + {file = "orjson-3.6.1-cp38-cp38-manylinux_2_24_x86_64.whl", hash = "sha256:1575700c542b98f6149dc5783e28709dccd27222b07ede6d0709a63cd08ec557"}, + {file = "orjson-3.6.1-cp38-none-win_amd64.whl", hash = "sha256:76d82b2c5c9f87629069f7b92053c64417fc5a42fdba08fece1d94c4483c5050"}, + {file = "orjson-3.6.1-cp39-cp39-macosx_10_7_x86_64.whl", hash = "sha256:cb84f10b816ed0cb8040e0d07bfe260549798f8929e9ab88b07622924d1a215f"}, + {file = "orjson-3.6.1-cp39-cp39-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:7e6211e515dd4bd5fbb09e6de6202c106619c059221ac29da41bc77a78812bb0"}, + {file = "orjson-3.6.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f15267d2e7195331b9823e278f953058721f0feaa5e6f2a7f62a8768858eed3b"}, + {file = "orjson-3.6.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:973e67cf4b8da44c02c3d1b0e68fb6c18630f67a20e1f7f59e4f005e0df622a0"}, + {file = "orjson-3.6.1-cp39-cp39-manylinux_2_24_x86_64.whl", hash = "sha256:1cdeda055b606c308087c5492f33650af4491a67315f89829d8680db9653137c"}, + {file = "orjson-3.6.1-cp39-none-win_amd64.whl", hash = "sha256:cd0dea1eb5fc48e441e4bfd6a26baa21a5ab44c3081025f5ce9248e38d89fbfa"}, + {file = "orjson-3.6.1.tar.gz", hash = "sha256:5ee598ce6e943afeb84d5706dc604bf90f74e67dc972af12d08af22249bd62d6"}, {file = "orjson-3.6.5-cp310-cp310-macosx_10_7_x86_64.whl", hash = "sha256:6c444edc073eb69cf85b28851a7a957807a41ce9bb3a9c14eefa8b33030cf050"}, {file = "orjson-3.6.5-cp310-cp310-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:432c6da3d8d4630739f5303dcc45e8029d357b7ff8e70b7239be7bd047df6b19"}, {file = "orjson-3.6.5-cp310-cp310-manylinux_2_24_aarch64.whl", hash = "sha256:0fa32319072fadf0732d2c1746152f868a1b0f83c8cce2cad4996f5f3ca4e979"}, @@ -2209,8 +2244,8 @@ pluggy = [ {file = "pluggy-1.0.0.tar.gz", hash = "sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159"}, ] pre-commit = [ - {file = "pre_commit-2.16.0-py2.py3-none-any.whl", hash = "sha256:758d1dc9b62c2ed8881585c254976d66eae0889919ab9b859064fc2fe3c7743e"}, - {file = "pre_commit-2.16.0.tar.gz", hash = "sha256:fe9897cac830aa7164dbd02a4e7b90cae49630451ce88464bca73db486ba9f65"}, + {file = "pre_commit-2.17.0-py2.py3-none-any.whl", hash = "sha256:725fa7459782d7bec5ead072810e47351de01709be838c2ce1726b9591dad616"}, + {file = "pre_commit-2.17.0.tar.gz", hash = "sha256:c1a8040ff15ad3d648c70cc3e55b93e4d2d5b687320955505587fd79bbaed06a"}, ] psycopg2-binary = [ {file = "psycopg2-binary-2.9.3.tar.gz", hash = "sha256:761df5313dc15da1502b21453642d7599d26be88bff659382f8f9747c7ebea4e"}, @@ -2505,20 +2540,20 @@ types-aiofiles = [ {file = "types_aiofiles-0.8.3-py3-none-any.whl", hash = "sha256:e261d6c0fafc3303c40cab64872609af8c702f6ec6590dc9f04a9bb8aaccc7b2"}, ] types-cryptography = [ - {file = "types-cryptography-3.3.12.tar.gz", hash = "sha256:aab189d3a63453fba48a9e5937f354ed8d4d2151c0aef44dc813cdcce631f375"}, - {file = "types_cryptography-3.3.12-py3-none-any.whl", hash = "sha256:1a81e18c2456f996d10b2410d16fbdc6f19653131e4ce0e09978015e9207c476"}, + {file = "types-cryptography-3.3.14.tar.gz", hash = "sha256:be07857ab3e52f254bf5559f8019d85b2200391c3e1f008b98f26ebedac027a8"}, + {file = "types_cryptography-3.3.14-py3-none-any.whl", hash = "sha256:c90ec0031c3d5262660990b62d9ec076bf5ed55eebfbe0011a53778e8aa62b9d"}, ] types-dataclasses = [ {file = "types-dataclasses-0.6.4.tar.gz", hash = "sha256:2f7ab6c565cf05cc7f27f31a4c2fcc803384e319aab292807b857ddf1473429f"}, {file = "types_dataclasses-0.6.4-py3-none-any.whl", hash = "sha256:fef6ed4742ca27996530c6d549cd704772a4a86e4781841c9bb387001e369ec3"}, ] types-enum34 = [ - {file = "types-enum34-1.1.2.tar.gz", hash = "sha256:22a08eacf89a1b71b2b770321b6940abe12afd6214c12917c4f119c935ff732a"}, - {file = "types_enum34-1.1.2-py3-none-any.whl", hash = "sha256:a1e1dcb80ae9d5a86c69ac7fcd65aec529541faaedffb3b2c840b0cbed8fbd61"}, + {file = "types-enum34-1.1.8.tar.gz", hash = "sha256:6f9c769641d06d73a55e11c14d38ac76fcd37eb545ce79cebb6eec9d50a64110"}, + {file = "types_enum34-1.1.8-py3-none-any.whl", hash = "sha256:05058c7a495f6bfaaca0be4aeac3cce5cdd80a2bad2aab01fd49a20bf4a0209d"}, ] types-ipaddress = [ - {file = "types-ipaddress-1.0.2.tar.gz", hash = "sha256:b3f29a5e1dabab9ec00c75654b53b07251f731d57295097c72c864524a31034d"}, - {file = "types_ipaddress-1.0.2-py3-none-any.whl", hash = "sha256:cb5eb3ad21acea538a1b404bbe2c43a7ba918e56d94c7399730cfece01b0a947"}, + {file = "types-ipaddress-1.0.7.tar.gz", hash = "sha256:b9a2322caf093553abecb630bb6fb4b84035ea8354649278b0a67b73ec2edf36"}, + {file = "types_ipaddress-1.0.7-py3-none-any.whl", hash = "sha256:fb7d4ce36b9037e2c5c34abbbc47fc9a116d7fcf45b6f4b015334d4be6ee73be"}, ] types-orjson = [ {file = "types-orjson-3.6.2.tar.gz", hash = "sha256:cf9afcc79a86325c7aff251790338109ed6f6b1bab09d2d4262dd18c85a3c638"}, diff --git a/pyproject.toml b/pyproject.toml index 3a5ea56..f5443db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,7 @@ aiomysql = ">=0.0.21,<0.0.23" aiosqlite = "^0.17.0" aiopg = "^1.3.3" asyncpg = ">=0.24,<0.26" +orjson = "*" # Sync database drivers for standard tooling around setup/teardown/migrations. psycopg2-binary = "^2.9.1" diff --git a/tests/test_queries/test_queryset_level_methods.py b/tests/test_queries/test_queryset_level_methods.py index 449c7df..7176010 100644 --- a/tests/test_queries/test_queryset_level_methods.py +++ b/tests/test_queries/test_queryset_level_methods.py @@ -6,6 +6,7 @@ import pytest import sqlalchemy import ormar +from ormar import QuerySet from ormar.exceptions import ( ModelPersistenceError, QueryDefinitionError, @@ -42,6 +43,7 @@ class ToDo(ormar.Model): id: int = ormar.Integer(primary_key=True) text: str = ormar.String(max_length=500) completed: bool = ormar.Boolean(default=False) + pairs: pydantic.Json = ormar.JSON(default=[]) class Category(ormar.Model): @@ -76,6 +78,26 @@ class ItemConfig(ormar.Model): pairs: pydantic.Json = ormar.JSON(default=["2", "3"]) +class QuerySetCls(QuerySet): + async def first_or_404(self, *args, **kwargs): + entity = await self.get_or_none(*args, **kwargs) + if not entity: + # maybe HTTPException in fastapi + raise ValueError("customer not found") + return entity + + +class Customer(ormar.Model): + class Meta: + metadata = metadata + database = database + tablename = "customer" + queryset_class = QuerySetCls + + id: Optional[int] = ormar.Integer(primary_key=True) + name: str = ormar.String(max_length=32) + + @pytest.fixture(autouse=True, scope="module") def create_test_database(): engine = sqlalchemy.create_engine(DATABASE_URL) @@ -348,3 +370,14 @@ async def test_bulk_operations_with_json(): await ItemConfig.objects.bulk_update(items) items = await ItemConfig.objects.all() assert all(x.pairs == ["1"] for x in items) + + +@pytest.mark.asyncio +async def test_custom_queryset_cls(): + async with database: + with pytest.raises(ValueError): + await Customer.objects.first_or_404(id=1) + + await Customer(name="test").save() + c = await Customer.objects.first_or_404(name="test") + assert c.name == "test"