do_orm_execute 中所有 sqlalchemy 查询表的查询提示

Question

根据 session.info 中的信息，我想为 sqlalchmey 查询的所有选定表动态添加类型提示。为此，我正在使用 do_orm_execute 会话事件并使用提示更改语句，如下所示：

@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):
        for from_ in [*getattr(statement, "froms", []), getattr(statement, "left", None), getattr(statement, "right", None), getattr(statement, "original", None)]:
            if isinstance(from_, sa.Table):
                orm_execute_state.statement = orm_execute_state.statement.with_hint(from_, hint)
            elif from_ is not None:
                _recursive_helper(from_)

    _recursive_helper(orm_execute_state.statement)

    return None

这在顶级可选项上工作得很好，但是当我有更复杂的查询也涉及子查询时，这个递归似乎不起作用。我的意思是，递归本身有效，但输出查询在子查询或连接上没有类型提示。

有什么想法吗？

Answer 1

原来是两个问题造成的：

sqlalchemy 的 _generative 装饰器，在获取属性时会自动创建一个副本。
需要在子查询中添加提示而不是原来的语句

总而言之，我提出了以下解决方案，它通过 __wrapped__ 属性访问原始函数并在需要时递归传递语句来绕过生成装饰器

@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(iter_, statement=None):
        stmt = iter_ if statement is None else statement
        for from_ in [*getattr(iter_, "froms", []), getattr(iter_, "left", None),
                      getattr(iter_, "right", None), getattr(iter_, "original", None)]:
            if hasattr(stmt, "with_hint") and isinstance(from_, sa.Table):
                stmt.with_hint.__wrapped__(stmt, from_, hint)  # use __wrapped__ to bypass attribute clone/generator
            elif from_ is not None:
                _recursive_helper(from_, statement=None if hasattr(from_, "with_hint") else stmt)

    _recursive_helper(orm_execute_state.statement)

do_orm_execute 中所有 sqlalchemy 查询表的查询提示

query hints on all sqlalchemy query tables in do_orm_execute

python

recursion

sqlalchemy