query hints on all sqlalchemy query tables in do_orm_execute

user_5678230 0
Viewed 251

Based on information in the session.info, I want to dynamically add a type hint to all the selected tables of a sqlalchmey query. For that, I am using the do_orm_execute session event and alter the statement with hints such as shown here:

@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

This works fine on top level selectables, but when I have more complex queries which also involves subqueries, this recursion does not seem to work. I mean, the recursion itself works, but the output queries do not have the type hint on the subqueries or joins.

any idea?

recursion sqlalchemy python
1 Answers

it turns out the problem is due to two problems:

  1. the _generative decorator of sqlalchemy which will automatically create a copy of the attribute when being get.

  2. the hints need to be added to the subquery instead of the original statement

all in all, i came up with the following solution which bypasses the generative decorator by accessing the original function through the __wrapped__ attribute and passing the statements recursively if required

@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)
user_56782300
Related