postgresql 注解函数龟式ORM查询中的KeyError 'hour'

rslzwgfq  于 2022-12-29  发布在  PostgreSQL
关注(0)|答案(1)|浏览(156)

需要运行下面的查询使用龟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;

尝试将内置函数date_trunc用作以下查询中的group_by
正在获取密钥错误“小时”。

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"])
mzsu5hc0

mzsu5hc01#

我没有使用过Tortoise ORM,但是看起来orm认为函数参数是一个表列(或模型字段)。您可以尝试覆盖resolve(self, model: 'Type[Model]', table: Table)方法:

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

相关问题