๐ Issue: `Model.annotate(...).update(...)` not working
DavideGalilei opened this issue ยท 0 comments
DavideGalilei commented
Description
Model.annotate(...).update(...)
does not work, giving an error.
tortoise.exceptions.FieldError: There is no non-virtual field bonus_salary on Model Employee
Example code
from tortoise import Tortoise, Model, fields, run_async
from datetime import datetime, timedelta
from tortoise.expressions import Case, When, Q, F
class Employee(Model):
id = fields.IntField(pk=True)
name = fields.TextField()
employed_at = fields.DatetimeField(auto_now_add=True)
salary = fields.BigIntField()
async def main():
await Tortoise.init(
db_url='postgres://username:password@localhost:5432/postgres',
modules={'models': ['__main__']},
)
await Tortoise.generate_schemas()
bob = await Employee.create(name='Bob', employed_at=datetime.now() - timedelta(days=30), salary=1000)
# If Bob was employed before yesterday, he gets a 10% bonus
# Doesn't work
await bob.filter(id=bob.id).select_for_update().annotate(
bonus_salary=Case(
When(Q(employed_at__lt=datetime.now() - timedelta(days=1)), then=F('salary') * 1.1),
default=F('salary'),
)
).update(salary=F('bonus_salary'))
await bob.refresh_from_db()
print(f"NEW SALARY:", bob.salary) # Bob's salary is now 110% of the original value
run_async(main())
Traceback
Traceback (most recent call last):
File "/root/.local/lib/python3.12/site-packages/tortoise/expressions.py", line 53, in resolver_arithmetic_expression
arithmetic_expression_or_field.name = model._meta.fields_db_projection[name]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
KeyError: 'bonus_salary'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/root/tortoise_bug/test.py", line 66, in <module>
run_async(main())
File "/root/.local/lib/python3.12/site-packages/tortoise/__init__.py", line 624, in run_async
loop.run_until_complete(coro)
File "/usr/lib64/python3.12/asyncio/base_events.py", line 687, in run_until_complete
return future.result()
^^^^^^^^^^^^^^^
File "/root/tortoise_bug/test.py", line 34, in main
await bob.filter(id=bob.id).select_for_update().annotate(
File "/root/.local/lib/python3.12/site-packages/tortoise/queryset.py", line 1159, in __await__
self._make_query()
File "/root/.local/lib/python3.12/site-packages/tortoise/queryset.py", line 1144, in _make_query
value = F.resolver_arithmetic_expression(self.model, value)[0]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.local/lib/python3.12/site-packages/tortoise/expressions.py", line 65, in resolver_arithmetic_expression
raise FieldError(f"There is no non-virtual field {name} on Model {model.__name__}")
tortoise.exceptions.FieldError: There is no non-virtual field bonus_salary on Model Employee
Workaround (using raw sql with pypika)
Click to expand
import pypika
from pypika.queries import QueryBuilder
from tortoise import Tortoise, Model, fields, run_async
from datetime import datetime, timedelta
from tortoise.expressions import Case, When, Q, F
class Employee(Model):
id = fields.IntField(pk=True)
name = fields.TextField()
employed_at = fields.DatetimeField(auto_now_add=True)
salary = fields.BigIntField()
async def main():
await Tortoise.init(
db_url='postgres://username:password@localhost:5432/postgres',
modules={'models': ['__main__']},
)
await Tortoise.generate_schemas()
bob = await Employee.create(name='Bob', employed_at=datetime.now() - timedelta(days=30), salary=1000)
EmployeePypika = pypika.Table(Employee._meta.db_table)
where_condition = EmployeePypika.id == bob.id
select_query = pypika.Query.from_(EmployeePypika).select(EmployeePypika.id).where(where_condition).for_update()
update_query: QueryBuilder = pypika.Query.update(EmployeePypika).where(where_condition)
update_query = update_query.set(
EmployeePypika.salary,
pypika.Case().when(EmployeePypika.employed_at < datetime.now() - timedelta(days=1), EmployeePypika.salary * 1.1)
.else_(EmployeePypika.salary),
)
print("SELECT:", select_query.get_sql())
print("UPDATE:", update_query.get_sql())
async with in_transaction() as conn:
await conn.execute_query(str(select_query))
await conn.execute_query(str(update_query))
await bob.refresh_from_db()
print(f"NEW SALARY:", bob.salary) # Bob's salary is now 110% of the original value
run_async(main())
Expected behavior
Either update()
should accept tortoise.expressions.Case()
or it should consider annotated values.
Also, this is an issue with .annotate().update()
directly, which in itself caused my Case(...)
use case to fail.
Additional context
The following snippet also gives an error:
await bob.filter(id=bob.id).select_for_update().update(salary=Case(
When(Q(employed_at__lt=datetime.now() - timedelta(days=1)), then=F('salary') * 1.1),
default=F('salary'),
))
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'Case'