diff --git a/.coverage b/.coverage index 3ff348c..f743d7d 100644 Binary files a/.coverage and b/.coverage differ diff --git a/README.md b/README.md index 4cfb463..ab8f88f 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 7bbef37..92da04b 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -203,6 +203,22 @@ class QuerySet: 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 259a8bd..9611ad6 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')