From a13c13550d6f13c1de0b32d12bc31d9ccfc12db5 Mon Sep 17 00:00:00 2001 From: collerek Date: Thu, 24 Sep 2020 17:17:10 +0200 Subject: [PATCH 1/2] add get/update or create queryset method and tests --- .coverage | Bin 53248 -> 53248 bytes README.md | 20 ++++++++++++++++- ormar/queryset/queryset.py | 16 ++++++++++++++ tests/test_queryset_level_methods.py | 31 +++++++++++++++++++++++++++ 4 files changed, 66 insertions(+), 1 deletion(-) diff --git a/.coverage b/.coverage index 3ff348ccb033a4b80c34f4fd99c01ff839278819..f743d7d929f8d1ea74f68ee8e7e632a6db4124c2 100644 GIT binary patch delta 188 zcmV;t07L(PpaX!Q1F%3Z2RAw~GCDCZvqvw#PywKm2Y@6mv<{IDehy|1T@FzWKn^hu z9u5r-01f61%?-s3w+(>}VGTnK2@LZL-3-tSy9}VS5fDZUlc` z{I}cM-Oji3Z3nZ|j1~bJ7##!w2{Ik_C6o8{-~RsXe|z6!f9SX0djXRJj` z{I}cM-43(Rj1~bI7aarv2{0Y@C6o94x4(b;-`@AwANuY0UXqTqBq9<70SQJD+AsG1 h+uQf;clWQpw?Dgo-utF+{qxV=-Mgzd|FheUB|#hrOw#}W diff --git a/README.md b/README.md index 4cfb46304..ab8f88f16 100644 --- a/README.md +++ b/README.md @@ -264,7 +264,8 @@ await news.posts.clear() ``` -Since version >=0.3.4 Ormar supports also queryset level delete and update statements +Since version >=0.3.4 Ormar supports also queryset level delete and update statements, +as well as get_or_create and update_or_create ```python import databases import ormar @@ -308,6 +309,23 @@ await Book.objects.update(each=True, genre='Fiction') all_books = await Book.objects.filter(genre='Fiction').all() assert len(all_books) == 3 +# helper get/update or create methods of queryset +# if not exists it will be created +vol1 = await Book.objects.get_or_create(title="Volume I", author='Anonymous', genre='Fiction') +assert await Book.objects.count() == 1 + +# if exists it will be returned +assert await Book.objects.get_or_create(title="Volume I", author='Anonymous', genre='Fiction') == vol1 +assert await Book.objects.count() == 1 + +# if not exist the instance will be persisted in db +vol2 = await Book.objects.update_or_create(title="Volume II", author='Anonymous', genre='Fiction') +assert await Book.objects.count() == 1 + +# if pk or pkname passed in kwargs (like id here) the object will be updated +assert await Book.objects.update_or_create(id=vol2.id, genre='Historic') +assert await Book.objects.count() == 1 + ``` ## Data types diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index 7bbef37ab..92da04b99 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -203,6 +203,22 @@ async def get(self, **kwargs: Any) -> "Model": self.check_single_result_rows_count(rows) return rows[0] + async def get_or_create(self, **kwargs: Any) -> "Model": + try: + return await self.get(**kwargs) + except NoMatch: + return await self.create(**kwargs) + + async def update_or_create(self, **kwargs: Any) -> "Model": + pk_name = self.model_cls.Meta.pkname + if "pk" in kwargs: + kwargs[pk_name] = kwargs.pop("pk") + if pk_name not in kwargs or kwargs.get(pk_name) is None: + return await self.create(**kwargs) + else: + model = await self.get(pk=kwargs[pk_name]) + return await model.update(**kwargs) + async def all(self, **kwargs: Any) -> List["Model"]: # noqa: A003 if kwargs: return await self.filter(**kwargs).all() diff --git a/tests/test_queryset_level_methods.py b/tests/test_queryset_level_methods.py index 259a8bdb6..9611ad6b1 100644 --- a/tests/test_queryset_level_methods.py +++ b/tests/test_queryset_level_methods.py @@ -73,3 +73,34 @@ async def test_delete_and_update(): await Book.objects.delete(each=True) all_books = await Book.objects.all() assert len(all_books) == 0 + + +@pytest.mark.asyncio +async def test_get_or_create(): + async with database: + tom = await Book.objects.get_or_create(title="Volume I", author='Anonymous', genre='Fiction') + assert await Book.objects.count() == 1 + + assert await Book.objects.get_or_create(title="Volume I", author='Anonymous', genre='Fiction') == tom + assert await Book.objects.count() == 1 + + assert await Book.objects.create(title="Volume I", author='Anonymous', genre='Fiction') + with pytest.raises(ormar.exceptions.MultipleMatches): + await Book.objects.get_or_create(title="Volume I", author='Anonymous', genre='Fiction') + + +@pytest.mark.asyncio +async def test_update_or_create(): + async with database: + tom = await Book.objects.update_or_create(title="Volume I", author='Anonymous', genre='Fiction') + assert await Book.objects.count() == 1 + + assert await Book.objects.update_or_create(id=tom.id, genre='Historic') + assert await Book.objects.count() == 1 + + assert await Book.objects.update_or_create(pk=tom.id, genre='Fantasy') + assert await Book.objects.count() == 1 + + assert await Book.objects.create(title="Volume I", author='Anonymous', genre='Fantasy') + with pytest.raises(ormar.exceptions.MultipleMatches): + await Book.objects.get(title="Volume I", author='Anonymous', genre='Fantasy') From df3bb9587da0f74b36b0d4b883d7eb837dd83fff Mon Sep 17 00:00:00 2001 From: collerek Date: Thu, 24 Sep 2020 17:20:44 +0200 Subject: [PATCH 2/2] remove unnec else --- ormar/queryset/queryset.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index 92da04b99..6c7f657b2 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -215,9 +215,8 @@ async def update_or_create(self, **kwargs: Any) -> "Model": kwargs[pk_name] = kwargs.pop("pk") if pk_name not in kwargs or kwargs.get(pk_name) is None: return await self.create(**kwargs) - else: - model = await self.get(pk=kwargs[pk_name]) - return await model.update(**kwargs) + model = await self.get(pk=kwargs[pk_name]) + return await model.update(**kwargs) async def all(self, **kwargs: Any) -> List["Model"]: # noqa: A003 if kwargs: