在 PostgreSQL 上使用 CTE() 的 SQLAlchemy 查询

Question

I have the following function (stored procedure) on PostgreSQL that calculates availability and pricing for a small boutique hotel prototype application:

-- Function that emulates Transact-SQL's IIF (if-and-only-if)
CREATE OR REPLACE FUNCTION IIF(BOOLEAN, DATE, DATE) RETURNS DATE
AS $$
    SELECT CASE $1 WHEN True THEN $2 ELSE $3 END
$$
LANGUAGE SQL IMMUTABLE;

-- Function to have together all steps that lead to availability and pricing calculation
CREATE OR REPLACE FUNCTION availability(check_in DATE, check_out DATE, guests INTEGER, room INTEGER[] DEFAULT '{}')
RETURNS TABLE (
    r_id INTEGER,
    r_floor_no INTEGER,
    r_room_no INTEGER,
    r_name VARCHAR,
    r_sgl_beds INTEGER,
    r_dbl_beds INTEGER,
    r_accommodates INTEGER,
    r_code VARCHAR,
    t_nights INTEGER,
    t_price REAL
) AS $$
BEGIN
RETURN QUERY
(
    WITH p AS (
         -- Sum of nights and prices per season (0..N)
         SELECT SUM(IIF($1 > t.date_to, t.date_to, $2) - IIF($1 > t.date_from, $1, t.date_from)) AS nights,
                SUM((IIF($2 > t.date_to, t.date_to, $2) - IIF($1 > t.date_from, $1, t.date_from)) * (t.base_price + t.bed_price * $3)) AS price
           FROM rate AS t
          WHERE (t.date_from, t.date_to) OVERLAPS ($1, $2)
            AND t.published = True
         ),
         a AS (
         -- Room availability
         SELECT r.id AS r_id,
                r.floor_no AS r_floor_no,
                r.room_no AS r_room_no,
                r.name AS r_name,
                r.sgl_beds AS r_sgl_beds,
                r.dbl_beds AS r_dbl_beds,
                r.accommodates AS r_accommodates,
                r.supplement AS r_supplement,
                r.code AS r_code
           FROM room AS r
          WHERE r.id NOT IN (
                SELECT b.id_room
                  FROM booking as b
                 WHERE (b.check_in, b.check_out) OVERLAPS ($1, $2)
                   AND b.cancelled IS NULL
                )
            AND r.accommodates >= $3
            AND CASE WHEN $4 = '{}'::INTEGER[] THEN r.id > 0 ELSE r.id = ANY($4) END
         )
  SELECT a.r_id AS r_id,
         a.r_floor_no AS r_floor_no,
         a.r_room_no AS r_room_no,
         a.r_name AS r_name,
         a.r_sgl_beds AS r_sgl_beds,
         a.r_dbl_beds AS r_dbl_beds,
         a.r_accommodates AS r_accommodates,
         a.r_code AS r_code,
         p.nights::INTEGER AS t_nights,
         (a.r_supplement * p.nights + p.price)::REAL AS t_price
    FROM a, p
ORDER BY t_price ASC, r_accommodates ASC, r_sgl_beds ASC, r_dbl_beds ASC, r_floor_no ASC, r_room_no ASC
);
END
$$ LANGUAGE plpgsql;

I am trying to migrate this code to SQLAlchemy but I cannot seem to be able to handle the double use of CTE (Common Table Expression) in the form of WITH p AS [..] and then a AS [..] on SQLAlchemy. Here you are what I have so far:

# Sum of nights and prices per season (0..N)
p = session.query(
    func.sum(Rate.date_to - Rate.date_from).label('nights'),
    (func.sum(
        case(
            [(p.check_in > Rate.date_to, Rate.date_to)],
            else_=p.check_out
        ) -
        case(
            [(p.check_in > Rate.date_from, p.check_in)],
            else_=Rate.date_from
        ) * (Rate.base_price + Rate.bed_price * p.guests)
    ).label('price'))
    ).\
    filter(
        tuple_(Rate.date_from, Rate.date_to).
        op('OVERLAPS')
        (tuple_(p.check_in, p.check_out))
    ).\
    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_(p.check_in, p.check_out))
    ).\
    filter(Booking.cancelled.is_(None)).\
    subquery('subq')

a = session.query(Room).\
    filter(Room.deleted.is_(None)).\
    filter(Room.id.notin_(subq)).\
    filter(Room.accommodates >= p.guests)
if p.rooms:
    a = a.filter(Room.id.any(p.rooms))
a = a.cte(name='a')

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

p.check_in (Date), p.check_out (Date), p.guests (int) and p.rooms (List of ints) are input parameters.

The error I am getting is:

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

On this line:

(tuple_(p.check_in, p.check_out))

which is inside the sub-query block:

# 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_(p.check_in, p.check_out))
    ).\
    filter(Booking.cancelled.is_(None)).\
    subquery('subq')

I have this feeling that SQLAlchemy expects just one call to cte() but I have not been able to figure it out from the documentation online . I've tried to build the big query block by block, then assemble them together, but with no success.

To help contextualise, here you are the data inside the room table:

 id | floor_no | room_no |                           name                           | sgl_beds | dbl_beds | supplement |  code  | deleted 
----+----------+---------+----------------------------------------------------------+----------+----------+------------+--------+---------
  1 |        1 |       1 | Normal bedroom with two single beds                      |        2 |        0 |         20 | pink   | 
  2 |        1 |       2 | Large bedroom with two single and one double beds        |        2 |        1 |         40 | black  | 
  3 |        1 |       3 | Very large bedroom with three single and one double beds |        3 |        1 |         50 | white  | 
  4 |        1 |       4 | Very large bedroom with four single beds                 |        4 |        0 |         40 | purple | 
  5 |        1 |       5 | Large bedroom with three single beds                     |        3 |        0 |         30 | blue   | 
  6 |        1 |       6 | Normal bedroom with one double bed                       |        0 |        1 |         20 | brown  | 

accommodates is now an hybrid property in the Room model class but it used to be a column in the table (and could be reverted to it, updated by a trigger).

And this is the rate table:

 id | date_from  |  date_to   | base_price | bed_price | published 
----+------------+------------+------------+-----------+-----------
  1 | 2017-03-01 | 2017-04-30 |         10 |        19 | t
  2 | 2017-05-01 | 2017-06-30 |         20 |        29 | t
  3 | 2017-07-01 | 2017-08-31 |         30 |        39 | t
  4 | 2017-09-01 | 2017-10-31 |         20 |        29 | t
  5 | 2018-03-01 | 2018-04-30 |         10 |        21 | t
  6 | 2018-05-01 | 2018-06-30 |         20 |        31 | t
  7 | 2018-07-01 | 2018-08-31 |         30 |        41 | t
  8 | 2018-09-01 | 2018-10-31 |         20 |        31 | t
  9 | 2019-03-01 | 2019-04-30 |         10 |        20 | t
 10 | 2019-05-01 | 2019-06-30 |         20 |        30 | t
 11 | 2019-07-01 | 2019-08-31 |         30 |        40 | t
 12 | 2019-09-01 | 2019-10-31 |         20 |        30 | t

Finally, this is a fragment of the booking table:

 id | id_guest | id_room |      reserved       | guests |  check_in  | check_out  | checked_in | checked_out | cancelled | base_price | taxes_percentage | taxes_value | total_price | locator | pin  |  status   |    meal_plan    | additional_services |                 uuid                 | deleted 
----+----------+---------+---------------------+--------+------------+------------+------------+-------------+-----------+------------+------------------+-------------+-------------+---------+------+-----------+-----------------+---------------------+--------------------------------------+---------
  1 |        1 |       1 | 2016-12-25 17:00:04 |      2 | 2017-05-05 | 2017-05-09 |            |             |           |        200 |               10 |          20 |         220 | AAAAA   | 1234 | Confirmed | BedAndBreakfast | "PoolKit"=>"1"      | 4df783c9-9375-47d6-8a9d-3309aa2c0a10 | 
  2 |        2 |       2 | 2016-12-26 09:03:54 |      3 | 2017-04-01 | 2017-04-11 |            |             |           |        500 |               10 |          50 |         550 | AAAAB   | 1234 | Confirmed | BedAndBreakfast | "PoolKit"=>"1"      | 0428692a-267a-46e7-871f-a7a20c8e9406 | 
  3 |        3 |       3 | 2016-01-25 14:43:00 |      3 | 2017-06-02 | 2017-06-12 |            |             |           |        500 |               10 |          50 |         550 | AAAAC   | 1234 | Confirmed | BedAndBreakfast | "PoolKit"=>"1"      | 12deeb14-1568-4b70-9247-5df2df433359 | 
  4 |        4 |       4 | 2016-01-25 14:43:00 |      3 | 2017-06-01 | 2017-06-10 |            |             |           |        500 |               10 |          50 |         550 | AAAAD   | 1234 | Confirmed | BedAndBreakfast | "PoolKit"=>"1"      | b3453b07-5ec7-4c15-be72-998e451998c6 | 
  5 |        5 |       5 | 2016-01-25 14:43:00 |      3 | 2017-06-08 | 2017-06-18 |            |             |           |        500 |               10 |          50 |         550 | AAAAE   | 1234 | Confirmed | BedAndBreakfast | "PoolKit"=>"1"      | 02a5c8f8-1d4c-45d6-9698-50bfa6d47b42 |

I am using the latest versions of both SQLAlchemy and PostgreSQL, so no restrictions there.

Everything you see here does not necessarily have to make all the sense in the world since this is just a prototype to test a number of features of a combination of technologies.

Thanks in advance.

Answer 1

So after Ilja figuring out the variable naming conflict, I kept on working on the query and this is the final, working result:

from sqlalchemy import func, tuple_, case, cast
from sqlalchemy import Integer as sqlInteger
from sqlalchemy import Float as sqlFloat
from sqlalchemy import Date as sqlDate

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')

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()

Please note that now input parameters are located inside the check_in , check_out , guests and rooms variables.