codekingpro/portable-devtools

1from __future__ import annotations2 3import typing as t4 5import sqlalchemy as sa6import sqlalchemy.exc as sa_exc7import sqlalchemy.orm as sa_orm8from flask.globals import app_ctx9 10if t.TYPE_CHECKING:11    from .extension import SQLAlchemy12 13 14class Session(sa_orm.Session):15    """A SQLAlchemy :class:`~sqlalchemy.orm.Session` class that chooses what engine to16    use based on the bind key associated with the metadata associated with the thing17    being queried.18 19    To customize ``db.session``, subclass this and pass it as the ``class_`` key in the20    ``session_options`` to :class:`.SQLAlchemy`.21 22    .. versionchanged:: 3.023        Renamed from ``SignallingSession``.24    """25 26    def __init__(self, db: SQLAlchemy, **kwargs: t.Any) -> None:27        super().__init__(**kwargs)28        self._db = db29        self._model_changes: dict[object, tuple[t.Any, str]] = {}30 31    def get_bind(32        self,33        mapper: t.Any | None = None,34        clause: t.Any | None = None,35        bind: sa.engine.Engine | sa.engine.Connection | None = None,36        **kwargs: t.Any,37    ) -> sa.engine.Engine | sa.engine.Connection:38        """Select an engine based on the ``bind_key`` of the metadata associated with39        the model or table being queried. If no bind key is set, uses the default bind.40 41        .. versionchanged:: 3.0.342            Fix finding the bind for a joined inheritance model.43 44        .. versionchanged:: 3.045            The implementation more closely matches the base SQLAlchemy implementation.46 47        .. versionchanged:: 2.148            Support joining an external transaction.49        """50        if bind is not None:51            return bind52 53        engines = self._db.engines54 55        if mapper is not None:56            try:57                mapper = sa.inspect(mapper)58            except sa_exc.NoInspectionAvailable as e:59                if isinstance(mapper, type):60                    raise sa_orm.exc.UnmappedClassError(mapper) from e61 62                raise63 64            engine = _clause_to_engine(mapper.local_table, engines)65 66            if engine is not None:67                return engine68 69        if clause is not None:70            engine = _clause_to_engine(clause, engines)71 72            if engine is not None:73                return engine74 75        if None in engines:76            return engines[None]77 78        return super().get_bind(mapper=mapper, clause=clause, bind=bind, **kwargs)79 80 81def _clause_to_engine(82    clause: sa.ClauseElement | None,83    engines: t.Mapping[str | None, sa.engine.Engine],84) -> sa.engine.Engine | None:85    """If the clause is a table, return the engine associated with the table's86    metadata's bind key.87    """88    table = None89 90    if clause is not None:91        if isinstance(clause, sa.Table):92            table = clause93        elif isinstance(clause, sa.UpdateBase) and isinstance(clause.table, sa.Table):94            table = clause.table95 96    if table is not None and "bind_key" in table.metadata.info:97        key = table.metadata.info["bind_key"]98 99        if key not in engines:100            raise sa_exc.UnboundExecutionError(101                f"Bind key '{key}' is not in 'SQLALCHEMY_BINDS' config."102            )103 104        return engines[key]105 106    return None107 108 109def _app_ctx_id() -> int:110    """Get the id of the current Flask application context for the session scope."""111    return id(app_ctx._get_current_object())  # type: ignore[attr-defined]112