如何在do_orm_execute中为SQLAlchemy所有查询表动态添加hint

问题原因 #

你的原代码问题出在所有with_hint调用都作用在顶层查询对象上，而SQLAlchemy的with_hint只会为当前查询语句自身的FROM子句中关联的表添加hint，子查询内部的表必须在子查询自身的语句对象上调用with_hint才会生效，给顶层语句加子查询内部表的hint属于无效操作。

修复后代码 #

@sa.event.listens_for(sa.orm.Session, 'do_orm_execute')
def _valid_as_of(orm_execute_state):
    valid_as_of = orm_execute_state.session.info.get('valid_as_of', None)
    if valid_as_of is None:
        return None
    hint = f"FOR SYSTEM_TIME {valid_as_of}"

    def _recursive_helper(statement):
        # 处理子查询、别名等包裹对象的内部语句
        if hasattr(statement, "element"):
            updated_inner = _recursive_helper(statement.element)
            if updated_inner is not statement.element:
                return type(statement)(updated_inner)
        # 处理JOIN结构的左右节点
        if hasattr(statement, "left") and statement.left is not None:
            updated_right = _recursive_helper(statement.right)
            if updated_right is not statement.right:
                statement = statement.join(
                    updated_right,
                    onclause=statement.onclause,
                    isouter=statement.isouter
                )
        # 处理当前语句的所有FROM对象
        if hasattr(statement, "froms"):
            for from_obj in statement.froms:
                if isinstance(from_obj, sa.Table):
                    # 给当前层级的语句加hint，而非顶层语句
                    statement = statement.with_hint(from_obj, hint)
                else:
                    updated_from = _recursive_helper(from_obj)
                    if updated_from is not from_obj:
                        statement = statement.select_from(updated_from)
        return statement

    # 替换原始查询为处理后的完整语句树
    orm_execute_state.statement = _recursive_helper(orm_execute_state.statement)
    return None

核心修改点 #

• 递归函数改为返回修改后的语句对象，所有修改都作用在当前遍历到的语句层级，而非统一修改顶层语句
• 新增对Subquery、Alias这类包裹类对象的处理，先递归修改其内部的element再重新构造包裹对象
• 遍历到当前层级的表时，直接给当前层级的语句调用with_hint，保证hint会被绑定到对应层级的SQL片段上
• 处理JOIN结构时，会递归处理JOIN的左右节点再重新构造JOIN对象
• 最后将完全处理后的整棵语句树赋值回orm_execute_state.statement，替换原始查询

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