CODEWITHSUNDEEP
pythonsqlalchemyflask-sqlalchemy
Joshmaker
Joshmaker

Reputation: 4248

SQLAlchemy select hybrid property expression as a column

I have a SQLAlchemy hyrbid property similar to what's the official docs describe as a Correlated Subquery Relationship Hybrid.

Their example code looks like:

class SavingsAccount(Base):
    __tablename__ = "account"
    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey("user.id"), nullable=False)
    balance = Column(Numeric(15, 5))


class User(Base):
    __tablename__ = "user"
    id = Column(Integer, primary_key=True)
    name = Column(String(100), nullable=False)

    accounts = relationship("SavingsAccount", backref="owner")

    @hybrid_property
    def balance(self):
        # I don't want to use this method
        return sum(acc.balance for acc in self.accounts)

    @balance.expression
    def balance(cls):
        # This is how I want to calculate the total balance
        return (
            select([func.sum(SavingsAccount.balance)])
            .where(SavingsAccount.user_id == cls.id)
            .label("total_balance")
        )

If I'm selecting a list of multiple users, I'd like to populate the value of User.balance based on the balance.expression so the aggregation is done in SQL not in Python. However, I can't find a way to do that.

The closest I've been able to come up with is:

>>> User.query.add_columns(User.balance).all()
[(<User 1>, Decimal('10.00')), (<User 2>, Decimal('15.00')), (<User 3>, Decimal('10.00'))]

Is there another way to make an aggregation query that populates a property on my model (instead of just returning a tuple? I looked at column_property but it has fewer examples and doesn't seem powerful enough for all my use cases.

Upvotes: 1

Views: 1957

Answers (1)

Joshmaker
Joshmaker

Reputation: 4248

I figured it out, thanks to this Google Groups post. Basically, I need to use a column_property, but because of all the logic I need to set that column_property in a hacky way. Ideally, I would be able to use the @declared_attr property, but that doesn't seem to be supported by SQLAlchemy at the momement.

As a result we need to define and attach it to the User class using a function with an event listener.

class SavingsAccount(ext.db.Model):
    __tablename__ = "test_account"
    id = ext.db.Column(ext.db.Integer, primary_key=True)
    user_id = ext.db.Column(ext.db.Integer, ext.db.ForeignKey("test_user.id"), nullable=False)
    balance = ext.db.Column(ext.db.Numeric(15, 5))


class User(ext.db.Model):
    __tablename__ = "test_user"
    id = ext.db.Column(ext.db.Integer, primary_key=True)
    name = ext.db.Column(ext.db.String(100), nullable=False)
    accounts = ext.db.relationship("SavingsAccount", backref="owner")


@event.listens_for(mapper, "mapper_configured")
def set_thread_count(mapper, cls) -> None:
    if not issubclass(cls, User):
        return

    total_balance_subquery = (
        select([func.sum(SavingsAccount.balance)])
        .where(SavingsAccount.user_id == cls.id)
        .label("total_balance")
    )
    cls.total_balance = column_property(total_balance_subquery, deferred=True)

Ugly, but it works. Would still be very happy to hear a cleaner alternative solution!

Upvotes: 3

Related Questions