Merge pull request #11 from collerek/queryset_level_ops

Queryset level ops

README.md

@@ -264,6 +264,52 @@ await news.posts.clear()
 
 ```
 
+Since version >=0.3.4 Ormar supports also queryset level delete and update statements
+```python
+import databases
+import ormar
+import sqlalchemy
+
+database = databases.Database("sqlite:///db.sqlite")
+metadata = sqlalchemy.MetaData()
+
+class Book(ormar.Model):
+    class Meta:
+        tablename = "books"
+        metadata = metadata
+        database = database
+
+    id: ormar.Integer(primary_key=True)
+    title: ormar.String(max_length=200)
+    author: ormar.String(max_length=100)
+    genre: ormar.String(max_length=100, default='Fiction', choices=['Fiction', 'Adventure', 'Historic', 'Fantasy'])
+
+await Book.objects.create(title='Tom Sawyer', author="Twain, Mark", genre='Adventure')
+await Book.objects.create(title='War and Peace', author="Tolstoy, Leo", genre='Fiction')
+await Book.objects.create(title='Anna Karenina', author="Tolstoy, Leo", genre='Fiction')
+await Book.objects.create(title='Harry Potter', author="Rowling, J.K.", genre='Fantasy')
+await Book.objects.create(title='Lord of the Rings', author="Tolkien, J.R.", genre='Fantasy')
+
+# update accepts kwargs that are used to update queryset model
+# all other arguments are ignored (argument names not in own model table)
+await Book.objects.filter(author="Tolstoy, Leo").update(author="Lenin, Vladimir") # update all Tolstoy's books
+all_books = await Book.objects.filter(author="Lenin, Vladimir").all()
+assert len(all_books) == 2
+
+# delete accepts kwargs that will be used in filter
+# acting in same way as queryset.filter(**kwargs).delete()
+await Book.objects.delete(genre='Fantasy') # delete all fantasy books
+all_books = await Book.objects.all()
+assert len(all_books) == 3
+
+# queryset needs to be filtered before deleting to prevent accidental overwrite
+# to update whole database table each=True needs to be provided as a safety switch
+await Book.objects.update(each=True, genre='Fiction')
+all_books = await Book.objects.filter(genre='Fiction').all()
+assert len(all_books) == 3
+
+```
+
 ## Data types
 
 The following keyword arguments are supported on all field types.

ormar/__init__.py

@@ -26,7 +26,7 @@ class UndefinedType:  # pragma no cover
 
 Undefined = UndefinedType()
 
-__version__ = "0.3.3"
+__version__ = "0.3.4"
 __all__ = [
     "Integer",
     "BigInteger",

ormar/models/modelproxy.py

@@ -25,6 +25,14 @@ class ModelTableProxy:
         self_fields = {k: v for k, v in self.dict().items() if k not in related_names}
         return self_fields
 
+    @classmethod
+    def extract_db_own_fields(cls) -> set:
+        related_names = cls._extract_related_names()
+        self_fields = {
+            name for name in cls.Meta.model_fields.keys() if name not in related_names
+        }
+        return self_fields
+
     @classmethod
     def substitute_models_with_pks(cls, model_dict: dict) -> dict:
         for field in cls._extract_related_names():

ormar/queryset/queryset.py

@@ -5,6 +5,7 @@ import sqlalchemy
 
 import ormar  # noqa I100
 from ormar import MultipleMatches, NoMatch
+from ormar.exceptions import QueryDefinitionError
 from ormar.queryset import FilterQuery
 from ormar.queryset.clause import QueryClause
 from ormar.queryset.query import Query
@@ -135,10 +136,28 @@ class QuerySet:
         expr = sqlalchemy.func.count().select().select_from(expr)
         return await self.database.fetch_val(expr)
 
-    async def delete(self, **kwargs: Any) -> int:
+    async def update(self, each: bool = False, **kwargs: Any) -> int:
+        self_fields = self.model_cls.extract_db_own_fields()
+        updates = {k: v for k, v in kwargs.items() if k in self_fields}
+        if not each and not self.filter_clauses:
+            raise QueryDefinitionError(
+                "You cannot update without filtering the queryset first. "
+                "If you want to update all rows use update(each=True, **kwargs)"
+            )
+        expr = FilterQuery(filter_clauses=self.filter_clauses).apply(
+            self.table.update().values(**updates)
+        )
+        return await self.database.execute(expr)
+
+    async def delete(self, each: bool = False, **kwargs: Any) -> int:
         if kwargs:
             return await self.filter(**kwargs).delete()
-        expr = FilterQuery(filter_clauses=self.filter_clauses,).apply(
+        if not each and not self.filter_clauses:
+            raise QueryDefinitionError(
+                "You cannot delete without filtering the queryset first. "
+                "If you want to delete all rows use delete(each=True)"
+            )
+        expr = FilterQuery(filter_clauses=self.filter_clauses).apply(
             self.table.delete()
         )
         return await self.database.execute(expr)

tests/test_many_to_many.py

@@ -71,10 +71,10 @@ async def create_test_database():
 async def cleanup():
     yield
     async with database:
-        await PostCategory.objects.delete()
+        await PostCategory.objects.delete(each=True)
-        await Post.objects.delete()
+        await Post.objects.delete(each=True)
-        await Category.objects.delete()
+        await Category.objects.delete(each=True)
-        await Author.objects.delete()
+        await Author.objects.delete(each=True)
 
 
 @pytest.mark.asyncio

tests/test_queryset_level_methods.py

@@ -0,0 +1,75 @@
+import databases
+import pytest
+import sqlalchemy
+
+import ormar
+from ormar.exceptions import QueryDefinitionError
+from tests.settings import DATABASE_URL
+
+database = databases.Database(DATABASE_URL, force_rollback=True)
+metadata = sqlalchemy.MetaData()
+
+
+class Book(ormar.Model):
+    class Meta:
+        tablename = "books"
+        metadata = metadata
+        database = database
+
+    id: ormar.Integer(primary_key=True)
+    title: ormar.String(max_length=200)
+    author: ormar.String(max_length=100)
+    genre: ormar.String(max_length=100, default='Fiction', choices=['Fiction', 'Adventure', 'Historic', 'Fantasy'])
+
+
+@pytest.fixture(autouse=True, scope="module")
+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_delete_and_update():
+    async with database:
+        async with database.transaction(force_rollback=True):
+            await Book.objects.create(title='Tom Sawyer', author="Twain, Mark", genre='Adventure')
+            await Book.objects.create(title='War and Peace', author="Tolstoy, Leo", genre='Fiction')
+            await Book.objects.create(title='Anna Karenina', author="Tolstoy, Leo", genre='Fiction')
+            await Book.objects.create(title='Harry Potter', author="Rowling, J.K.", genre='Fantasy')
+            await Book.objects.create(title='Lord of the Rings', author="Tolkien, J.R.", genre='Fantasy')
+
+            all_books = await Book.objects.all()
+            assert len(all_books) == 5
+
+            await Book.objects.filter(author="Tolstoy, Leo").update(author="Lenin, Vladimir")
+            all_books = await Book.objects.filter(author="Lenin, Vladimir").all()
+            assert len(all_books) == 2
+
+            historic_books = await Book.objects.filter(genre='Historic').all()
+            assert len(historic_books) == 0
+
+            with pytest.raises(QueryDefinitionError):
+                await Book.objects.update(genre='Historic')
+
+            await Book.objects.filter(author="Lenin, Vladimir").update(genre='Historic')
+
+            historic_books = await Book.objects.filter(genre='Historic').all()
+            assert len(historic_books) == 2
+
+            await Book.objects.delete(genre='Fantasy')
+            all_books = await Book.objects.all()
+            assert len(all_books) == 3
+
+            await Book.objects.update(each=True, genre='Fiction')
+            all_books = await Book.objects.filter(genre='Fiction').all()
+            assert len(all_books) == 3
+
+            with pytest.raises(QueryDefinitionError):
+                await Book.objects.delete()
+
+            await Book.objects.delete(each=True)
+            all_books = await Book.objects.all()
+            assert len(all_books) == 0