Merge pull request #12 from collerek/get_or_create

add get/update or create queryset method and tests

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

ormar/queryset/queryset.py

@@ -203,6 +203,21 @@ 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)
+        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()

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')