KeyError 'hour' in annotate function tortoise ORM query

Question

Need to run the below query using tortoise ORM.

SELECT date_trunc('hour',"created") "group_name",COUNT("id") "count" FROM "abc" WHERE "user_id"='xyz@gmail.com' AND "active"=true GROUP BY "group_name" ORDER BY group_name DESC;

Trying to use in-built function date_trunc as group_by in below query.
Getting KeyError 'hour'.

ABC.filter(user_id=user_id, active=True) \
        .annotate(count=Count("id"), group_name=DateTrunc("hour", "created")) \
        .group_by('group_name').order_by('group_name').values_list('group_name', 'count').sql()


class ABC(AbstractBaseABC, ModelUtilMixin):
    class Meta:
        table = 'abc'


class AbstractBaseABC(Model):
    serializable_keys = {'id', 'user_id', 'active', 'created', 'updated'}
    id = fields.BigIntField(pk=True)
    user_id = CITextField(index=True)
    active = fields.BooleanField(default=True)
    created = fields.DatetimeField(index=True)
    updated = fields.DatetimeField()

    class Meta:
        abstract = True

class DateTrunc(Function):
    database_func = CustomFunction("date_trunc", ["name", "text"])

score 0 · Answer 1 · answered Dec 28 '22 at 16:01

I didn't work with Tortoise ORM, but it looks like orm thinks that the function argument is a table column (or model field). You can try to override resolve(self, model: 'Type[Model]', table: Table) method:

from typing import Type, Union, Any

from pypika import Table
from pypika.terms import ArithmeticExpression
from pypika.terms import Function as PypikaFunc
from tortoise import Model, fields, Tortoise, run_async
from tortoise.expressions import Function, F
from tortoise.fields import TextField
from tortoise.functions import Count


class User(Model):
    serializable_keys = {'id', 'user_id', 'active', 'created', 'updated'}
    id = fields.BigIntField(pk=True)
    user_id = TextField()
    active = fields.BooleanField(default=True)
    created = fields.DatetimeField()
    updated = fields.DatetimeField()

    class Meta:
        table = 'abc'


class DateTrunc(PypikaFunc):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__('date_trunc', *args, **kwargs)


class OrmDateTrunc(Function):
    def __init__(
        self,
        interval: str,
        field: Union[str, F, ArithmeticExpression, 'Function'],
        *default_values: Any
    ) -> None:
        super().__init__(field, *default_values)
        self.interval = interval

    def resolve(self, model: 'Type[Model]', table: Table) -> dict:
        # just an example
        function = self._resolve_field_for_model(model, table, self.field)
        function['field'] = DateTrunc(self.interval, function['field'])
        return function


async def init():
    await Tortoise.init(
        config={
            'connections': {
                'default': {
                    'engine': 'tortoise.backends.asyncpg',
                    'credentials': {
                        'database': 'DB_NAME',
                        'host': 'HOST_NAME',
                        'password': 'PASSWORD',
                        'port': 5432,
                        'user': 'DB_USER',
                    }
                }
            },
            'apps': {
                'models': {
                    'models': ['__main__'],
                    'default_connection': 'default',
                }
            },
        }
    )

    await Tortoise.generate_schemas()


run_async(init())

query = (
    User
    .filter(user_id='xyz@gmail.com', active=True)
    .annotate(group_name=OrmDateTrunc('hour', 'created'), count=Count('id'))
    .group_by('group_name')
    .order_by('-group_name')
    .values('group_name', 'count')
)

print(query.sql())
# SELECT date_trunc('hour',"created") "group_name",
#        COUNT("id") "count"
#   FROM "abc"
#  WHERE "user_id"='xyz@gmail.com' AND "active"=true
#  GROUP BY "group_name" 
#  ORDER BY date_trunc('hour',"created") DESC

KeyError 'hour' in annotate function tortoise ORM query

1 Answers1