如何让with_loader_criteria在查询中作用于多个模型？

问题场景 #

我使用SQLAlchemy的with_loader_criteria结合do_orm_execute钩子替代1.4版本中废弃的before_compile钩子，目的是给所有继承Versioned混入类的模型自动添加WHERE条件，但实际执行时只有查询中的第一个模型生效。

现有代码 #

钩子处理函数：

@event.listens_for(Session, "do_orm_execute")
def receive_do_orm_execute(execute_state):
    if execute_state.is_select and not execute_state.is_column_load:
        print(execute_state.statement)
        execute_state.statement = execute_state.statement.options(
            with_loader_criteria(
                Versioned,
                lambda cls: cls.valid_range.contains(datetime.now(timezone.utc)),
                include_aliases=True,
                propagate_to_loaders=True,
            )
        )
        print(execute_state.statement)

查询语句：

db_session.query(A, B).join(B, A.id == B.a_id)

执行结果 #

原始SQL：

SELECT
    a.id AS a_id,
    a.valid_range AS a_valid_range,
    b.id AS b_id,
    b.a_id AS b_a_id,
    b.valid_range AS b_valid_range
FROM
    a
JOIN
    b
ON
    a.id = b.a_id

修改后的SQL（仅A模型添加了条件）：

SELECT
    a.id AS a_id,
    a.valid_range AS a_valid_range,
    b.id AS b_id,
    b.a_id AS b_a_id,
    b.valid_range AS b_valid_range
FROM
    a
JOIN
    b
ON
    a.id = b.a_id
WHERE
    a.valid_range @> :current_datetime_1

注：A、B均继承自Versioned混入类，预期两个模型都会添加valid_range过滤条件。

解决方案 #

问题根源：with_loader_criteria默认仅对查询的主实体生效，当使用query(A, B)时，只有A被识别为主实体，B作为额外查询实体不会自动应用条件。需要遍历查询中的所有实体，对每个继承Versioned的模型单独应用过滤规则。

修改后的钩子函数 #

from sqlalchemy import inspect
import datetime
from sqlalchemy import event
from sqlalchemy.orm import with_loader_criteria

@event.listens_for(Session, "do_orm_execute")
def receive_do_orm_execute(execute_state):
    if execute_state.is_select and not execute_state.is_column_load:
        stmt = execute_state.statement
        current_time = datetime.datetime.now(datetime.timezone.utc)
        
        # 遍历查询中的所有实体，逐个应用条件
        for entity in stmt.entities:
            entity_cls = entity.entity_class
            # 检查当前实体是否继承自Versioned
            if inspect(entity_cls).is_constructable and issubclass(entity_cls, Versioned):
                stmt = stmt.options(
                    with_loader_criteria(
                        entity_cls,
                        lambda cls: cls.valid_range.contains(current_time),
                        include_aliases=True,
                        propagate_to_loaders=True,
                    )
                )
        
        execute_state.statement = stmt

效果验证 #

修改后生成的SQL会同时包含A和B的过滤条件：

SELECT
    a.id AS a_id,
    a.valid_range AS a_valid_range,
    b.id AS b_id,
    b.a_id AS b_a_id,
    b.valid_range AS b_valid_range
FROM
    a
JOIN
    b
ON
    a.id = b.a_id
WHERE
    a.valid_range @> :current_datetime_1
    AND b.valid_range @> :current_datetime_2

补充说明 #

• include_aliases=True确保模型使用别名（如aliased(A)）时也会应用过滤条件
• propagate_to_loaders=True会让关联加载的实体（如通过relationship加载的子模型）也自动应用该条件
• 提前计算current_time避免多次调用生成不同的时间戳

内容的提问来源于stack exchange，提问作者Kris Harper