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



根据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

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

任何想法?

原来这个问题是由两个问题引起的:

  1. sqlalchemy的_generative修饰符,它将在获取属性时自动创建该属性的副本。

  2. 提示需要添加到子查询中,而不是原始语句

总而言之,我提出了以下解决方案:通过__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)

最新更新