PostgreSQL 数据库上的 SQLAlchemy 查询使用 CTE()

3 阅读1分钟

我想将使用 PostgreSQL 存储过程 availability 的应用程序迁移到 SQLAlchemy。该存储过程用于计算小型精品酒店原型应用程序的空房情况和定价。availability 函数接受以下参数:

  • check_in: 入住日期
  • check_out: 退房日期
  • guests: 入住客人数量
  • rooms (可选):要检查空房情况的房间 ID 列表

此函数使用 CTE(通用表表达式)来计算空房情况和定价。第一个 CTE 名为 p,它计算每个季节的总住宿天数和价格。第二个 CTE 名为 a,它计算可用的客房。

我试图将此代码移植到 SQLAlchemy,但无法处理 WITH p AS [..]a AS [..] 形式的 CTE 的双重使用。我收到以下错误:

AttributeError: 'CTE' object has no attribute 'check_in'

解决方案

为了解决此问题,我需要使用 SQLAlchemy 的 cte() 函数来创建 CTE。我可以通过以下方式创建 pa CTE:

p = session.query(
    func.sum(
        case(
            [(check_out > Rate.date_to, Rate.date_to)],
            else_=check_out
        ) -
        case(
            [(check_in > Rate.date_from, check_in)],
            else_=Rate.date_from
        )
    ).label('nights'),
    (func.sum((
        case(
            [(check_out > Rate.date_to, Rate.date_to)],
            else_=check_out
        ) -
        case(
            [(check_in > Rate.date_from, check_in)],
            else_=Rate.date_from
        )) * (Rate.base_price + Rate.bed_price * guests)
    ).label('price'))
    ).\
    filter(
        tuple_(Rate.date_from, Rate.date_to).
        op('OVERLAPS')
        (tuple_(cast(check_in, sqlDate), cast(check_out, sqlDate)))
    ).\
    filter(Rate.published.is_(True)).\
    cte(name='p')

# Room availability using a sub-select
subq = session.query(Booking.id_room.label('id')).\
    filter(
        tuple_(Booking.check_in, Booking.check_out).
        op('OVERLAPS')
        (tuple_(cast(check_in, sqlDate), cast(check_out, sqlDate)))
    ).\
    filter(Booking.cancelled.is_(None)).\
    subquery('subq')

a = session.query(Room.id, Room.floor_no, Room.room_no, Room.name,
                  Room.sgl_beds, Room.dbl_beds,  Room.supplement,
                  Room.code, Room.number, Room.accommodates).\
    filter(Room.deleted.is_(None)).\
    filter(Room.id.notin_(subq)).\
    filter(Room.accommodates >= guests)
if rooms:
    a = a.filter(Room.id.any(rooms))
a = a.cte(name='a')

然后,我可以使用这些 CTE 来计算可用性并为每个房间获取总价格:

result = session.query(
    a.c.id, a.c.floor_no, a.c.room_no, a.c.name, a.c.sgl_beds,
    a.c.dbl_beds, a.c.code, a.c.number, a.c.accommodates,
    cast(p.c.nights, sqlInteger).label('nights'),
    cast(a.c.supplement * p.c.nights + p.c.price, sqlFloat).
    label('total_price')).\
    order_by('total_price ASC').\
    order_by(a.c.accommodates.asc()).\
    order_by(a.c.sgl_beds.asc()).\
    order_by(a.c.dbl_beds.asc()).\
    order_by(a.c.floor_no.asc()).\
    order_by(a.c.room_no.asc()).\
    all()

通过这些步骤,我能够将 availability 函数移植到 SQLAlchemy,并成功地计算了可用性并为每个房间提供了总价格。