Possible to eager load in select using hybrid expression?

[dusktreader]

I have an ORM class derived from Base with a 1-M relationship to another. Futher, I've defined a hybrid attribute (with an analog expression) to compute a sum based on a numerical column in the related class. I am able to get a list of matching ORM objects to load, but without calling refresh() on each one, the hybrid attribute won't load. Ideally, I would like to be able to load the hybrid attribute eagerly at the first query execution time.

Example code:

# stub of https://github.com/omnivector-solutions/license-manager/blob/main/backend/lm_backend/api/models/feature.py
class Feature(Base):
    __tablename__ = "features"
    id = Column(Integer, primary_key=True)
    bookings = relationship("Booking", back_populates="feature", lazy="selectin")

    @hybrid_property
    def booked_total(self) -> int:
        return sum(b.quantity for b in self.bookings)

    @booked_total.inplace.expression
    @classmethod
    def _booked_total_expression(cls) -> SQLColumnExpression[int]:
        return select(func.sum(Booking.quantity)).where(Booking.feature_id == cls.id).label("booked_total")

# stub of https://github.com/omnivector-solutions/license-manager/blob/main/backend/lm_backend/api/models/booking.py
class Booking(Base):
    __tablename__ = "bookings"
    id = Column(Integer, primary_key=True)
    quantity = Column(Integer, CheckConstraint("quantity>=0"), nullable=False)

Now, when I execute a query without refreshing the resulting the resulting objects, the booked_total attribute shows 0:

# stub of https://github.com/omnivector-solutions/license-manager/blob/main/backend/lm_backend/api/cruds/generic.py

async def read_all(self, db_session: AsyncSession, force_refresh: bool = False):
    stmt = select(self.model)
    query = await db_session.scalars(stmt)
    db_objs = list(query.all())
    if force_refresh:
        await gather(*(db_session.refresh(db_obj) for db_obj in db_objs))
    return db_objs

When I analyzed the query that's produced from the select, it does not join against bookings. However, if I include the hybrid attribute as a column in the select, it will do the join. The scalars() method doesn't work, in this case, though because executing returns a tuple of the Feature object and a scalar booking_total.

The only way I have found to get the booked_total to populate is to use the force_refresh parameter to the method and refresh each individual result.

In a case like this where I know I want the booked_total to be included in the result of a query, I would love to be able to direct sqlalchemy to use the hybrid expression eagerly and to load the results in the resulting ORM objects.

Based on some investigation (see: https://groups.google.com/g/sqlalchemy/c/XQW8gxX93hU), it seems like it might be a limitation of hybrid attributes that a select can't load them into the dict of the resulting objects. Is there a way to work around this? I would really like to reap the performance gains of including the sub-select in the generated query instead of issuing a bunch of axillary queries to load the computed value.

Thanks!!

[zzzeek]

hi -

a hybrid attribute, when you call it from a loaded instance, is just a python function, a method on the instance. I dont see any example use here in the code given, but if you have code that says assert my_instance.booked_total > 0 you could replace that with assert sum(b.quantity for b in my_instance.bookings) > 0 and that is equivalent.

When you say "However, if I include the hybrid attribute as a column in the select, it will do the join." that doesnt match the code illustrated because you have lazy="selectin", which normally should just be SELECTing from "bookings" alone, with a list of primary key identifiers from "features". if the relationship is a simple one-to-many it does not need to emit a JOIN to the parent table.

overall, I would be able to provide better guidance if the above fragments were worked into a self-contained MCVE that I can copy-and-paste here and run.

[zzzeek]

Here's an MCVE. See if you can modify this to show the mis-behavior you are seeing

import asyncio

from sqlalchemy import CheckConstraint
from sqlalchemy import Column
from sqlalchemy import ForeignKey
from sqlalchemy import func
from sqlalchemy import Integer
from sqlalchemy import select
from sqlalchemy import SQLColumnExpression
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import relationship

Base = declarative_base()


class Feature(Base):
    __tablename__ = "features"
    id = Column(Integer, primary_key=True)
    bookings = relationship(
        "Booking", back_populates="feature", lazy="selectin"
    )

    @hybrid_property
    def booked_total(self) -> int:
        return sum(b.quantity for b in self.bookings)

    @booked_total.inplace.expression
    @classmethod
    def _booked_total_expression(cls) -> SQLColumnExpression[int]:
        return (
            select(func.sum(Booking.quantity))
            .where(Booking.feature_id == cls.id)
            .label("booked_total")
        )


class Booking(Base):
    __tablename__ = "bookings"
    id = Column(Integer, primary_key=True)
    quantity = Column(Integer, CheckConstraint("quantity>=0"), nullable=False)
    feature_id = Column(ForeignKey("features.id"))
    feature = relationship("Feature", back_populates="bookings")


async def read_all(db_session: AsyncSession):
    stmt = select(Feature)
    query = await db_session.scalars(stmt)
    db_objs = list(query.all())
    return db_objs


async def generate_data(db_session: AsyncSession):
    db_session.add_all(
        [
            Feature(
                bookings=[
                    Booking(quantity=5),
                    Booking(quantity=10),
                ]
            ),
            Feature(
                bookings=[
                    Booking(quantity=8),
                    Booking(quantity=19),
                ]
            ),
        ]
    )
    await db_session.commit()


async def main():
    engine = create_async_engine(
        "postgresql+asyncpg://scott:tiger@localhost/test", echo=True
    )

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)

    async with AsyncSession(engine) as sess:
        await generate_data(sess)

    async with AsyncSession(engine) as sess:
        result = await read_all(sess)

        for obj in result:
            print(f"Object: {obj} booked total: {obj.booked_total}")
            assert obj.booked_total > 0


asyncio.run(main())

output within the read_all() function:

2023-07-19 09:45:29,035 INFO sqlalchemy.engine.Engine BEGIN (implicit)
2023-07-19 09:45:29,036 INFO sqlalchemy.engine.Engine SELECT features.id 
FROM features
2023-07-19 09:45:29,036 INFO sqlalchemy.engine.Engine [generated in 0.00024s] ()
2023-07-19 09:45:29,043 INFO sqlalchemy.engine.Engine SELECT bookings.feature_id AS bookings_feature_id, bookings.id AS bookings_id, bookings.quantity AS bookings_quantity 
FROM bookings 
WHERE bookings.feature_id IN ($1::INTEGER, $2::INTEGER)
2023-07-19 09:45:29,043 INFO sqlalchemy.engine.Engine [generated in 0.00028s] (1, 2)
Object: <__main__.Feature object at 0x7fcb1727a910> booked total: 15
Object: <__main__.Feature object at 0x7fcb1727a990> booked total: 27
2023-07-19 09:45:29,046 INFO sqlalchemy.engine.Engine ROLLBACK

[zzzeek]

if these objects were already loaded, perhaps you want to add populate_existing ?

async def read_all(db_session: AsyncSession):
    stmt = select(Feature).execution_options(populate_existing=True)
    query = await db_session.scalars(stmt)
    db_objs = list(query.all())
    return db_objs

[dusktreader]

Thank you for your response on this. I'm sorry I got busy and never got back to you with an MCVE. Thank you for taking the time to explain this and provide your own example.