Home > Mobile >  SQLAlchemy - Eager load subquery as column for an ORM class
SQLAlchemy - Eager load subquery as column for an ORM class

Time:10-16

I have two tables. One is mapped as an ORM class like so:

from typing import Union
from dataclasses import dataclass
from api.models import db, SerializerMixin

@dataclass
class Company(db.Model, SerializerMixin):
    __tablename__ = "Company"
    company_id: int = db.Column(db.Integer, primary_key = True)
    name: str = db.Column(db.String(50), nullable = False, unique = True)
    primary_contact: Union[int, None] = db.Column(db.Integer, db.ForeignKey("User.user_id"))
    website: Union[str, None] = db.Column(db.Text)

The second one is written as just a plain SQL table (wrote it this way since I don't plan to query from this table directly, only used for JOINS and COUNTS):

user_company_table = db.Table("UserCompany",
    db.Column("user_id", db.Integer, db.ForeignKey("User.user_id"), primary_key = True, unique = True),
    db.Column("company_id", db.Integer, db.ForeignKey("Company.company_id"), primary_key = True),
    db.Column("assigned", db.DateTime, server_default = text("CURRENT_TIMESTAMP"), onupdate = datetime.utcnow),
    db.Column("approved", db.Boolean, nullable = False, default = False)
)

I need to fetch all columns from Company while at the same time getting a COUNT of all users who have approved = true and are assigned to one specific company. The resulting SQL would look like this:

SELECT Company.company_id, Company.name, Company.primary_contact, Company.website, (SELECT COUNT(1) FROM UserCompany WHERE approved = true and company_id = Company.company_id) AS user_count FROM Company;

This should give me this sample result: enter image description here

So user_count is a result set from a subquery that was appended to the Company table. How would I go about adding a user_count property to the Company class that eager loads this subquery and attaches the result as a column whenever I run Company.query.all()?

Notes: using Flask 2.0 , Flask-SQLAlchemy, Python 3.9 (that explains the @dataclass and the type hints in the Model properties)

CodePudding user response:

I think you can use a correlated subquery relationship hybrid for this:

class Company:
  ...

  @hybrid_property
  def user_count(self):
    return db.session.query(user_company_table)\
      .filter_by(company_id=self.id, approved=True)\
      .count()

  @user_count.expression
  def user_count(cls):
    return select(func.count(1))\
      .where(user_company_table.c.company_id==cls.id)\
      .where(user_company_table.c.approved==True)\
      .label('user_count')

You could even opt for the very short column_property with a scalar subquery: https://docs.sqlalchemy.org/en/14/orm/mapped_sql_expr.html#using-column-property

  • Related