tortoise/tortoise-orm

๐Ÿ› Issue: `Model.annotate(...).update(...)` not working

DavideGalilei opened this issue ยท 0 comments

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'