diff --git a/.coverage b/.coverage index 17ef238..3ff348c 100644 Binary files a/.coverage and b/.coverage differ diff --git a/ormar/models/modelproxy.py b/ormar/models/modelproxy.py index 9052b71..5534dcb 100644 --- a/ormar/models/modelproxy.py +++ b/ormar/models/modelproxy.py @@ -25,6 +25,12 @@ 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(): @@ -51,9 +57,9 @@ class ModelTableProxy: related_names = set() for name, field in cls.Meta.model_fields.items(): if ( - inspect.isclass(field) - and issubclass(field, ForeignKeyField) - and not field.virtual + inspect.isclass(field) + and issubclass(field, ForeignKeyField) + and not field.virtual ): related_names.add(name) return related_names @@ -65,9 +71,9 @@ class ModelTableProxy: related_names = set() for name, field in cls.Meta.model_fields.items(): if ( - inspect.isclass(field) - and issubclass(field, ForeignKeyField) - and field.nullable + inspect.isclass(field) + and issubclass(field, ForeignKeyField) + and field.nullable ): related_names.add(name) return related_names @@ -95,7 +101,7 @@ class ModelTableProxy: @staticmethod def resolve_relation_field( - item: Union["Model", Type["Model"]], related: Union["Model", Type["Model"]] + item: Union["Model", Type["Model"]], related: Union["Model", Type["Model"]] ) -> Type[Field]: name = ModelTableProxy.resolve_relation_name(item, related) to_field = item.Meta.model_fields.get(name) @@ -121,12 +127,12 @@ class ModelTableProxy: for field in one.Meta.model_fields.keys(): current_field = getattr(one, field) if isinstance(current_field, list) and not isinstance( - current_field, ormar.Model + current_field, ormar.Model ): setattr(other, field, current_field + getattr(other, field)) elif ( - isinstance(current_field, ormar.Model) - and current_field.pk == getattr(other, field).pk + isinstance(current_field, ormar.Model) + and current_field.pk == getattr(other, field).pk ): setattr( other, diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index 7a44fec..639a803 100644 --- a/ormar/queryset/queryset.py +++ b/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 @@ -15,13 +16,13 @@ if TYPE_CHECKING: # pragma no cover class QuerySet: def __init__( # noqa CFQ002 - self, - model_cls: Type["Model"] = None, - filter_clauses: List = None, - exclude_clauses: List = None, - select_related: List = None, - limit_count: int = None, - offset: int = None, + self, + model_cls: Type["Model"] = None, + filter_clauses: List = None, + exclude_clauses: List = None, + select_related: List = None, + limit_count: int = None, + offset: int = None, ) -> None: self.model_cls = model_cls self.filter_clauses = [] if filter_clauses is None else filter_clauses @@ -52,7 +53,7 @@ class QuerySet: pkname = self.model_cls.Meta.pkname pk = self.model_cls.Meta.model_fields[pkname] if new_kwargs.get(pkname, ormar.Undefined) is None and ( - pk.nullable or pk.autoincrement + pk.nullable or pk.autoincrement ): del new_kwargs[pkname] return new_kwargs @@ -135,10 +136,25 @@ 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) + ) + # print(expr.compile(compile_kwargs={"literal_binds": True})) + 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) diff --git a/tests/test_many_to_many.py b/tests/test_many_to_many.py index 10347bf..2859b71 100644 --- a/tests/test_many_to_many.py +++ b/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 Post.objects.delete() - await Category.objects.delete() - await Author.objects.delete() + await PostCategory.objects.delete(each=True) + await Post.objects.delete(each=True) + await Category.objects.delete(each=True) + await Author.objects.delete(each=True) @pytest.mark.asyncio diff --git a/tests/test_queryset_level_methods.py b/tests/test_queryset_level_methods.py new file mode 100644 index 0000000..259a8bd --- /dev/null +++ b/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