codekingpro/portable-devtools

1from __future__ import annotations2 3import typing as t4 5import sqlalchemy as sa6import sqlalchemy.event as sa_event7import sqlalchemy.orm as sa_orm8from flask import current_app9from flask import has_app_context10from flask.signals import Namespace  # type: ignore[attr-defined]11 12if t.TYPE_CHECKING:13    from .session import Session14 15_signals = Namespace()16 17models_committed = _signals.signal("models-committed")18"""This Blinker signal is sent after the session is committed if there were changed19models in the session.20 21The sender is the application that emitted the changes. The receiver is passed the22``changes`` argument with a list of tuples in the form ``(instance, operation)``.23The operations are ``"insert"``, ``"update"``, and ``"delete"``.24"""25 26before_models_committed = _signals.signal("before-models-committed")27"""This signal works exactly like :data:`models_committed` but is emitted before the28commit takes place.29"""30 31 32def _listen(session: sa_orm.scoped_session[Session]) -> None:33    sa_event.listen(session, "before_flush", _record_ops, named=True)34    sa_event.listen(session, "before_commit", _record_ops, named=True)35    sa_event.listen(session, "before_commit", _before_commit)36    sa_event.listen(session, "after_commit", _after_commit)37    sa_event.listen(session, "after_rollback", _after_rollback)38 39 40def _record_ops(session: Session, **kwargs: t.Any) -> None:41    if not has_app_context():42        return43 44    if not current_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"]:45        return46 47    for targets, operation in (48        (session.new, "insert"),49        (session.dirty, "update"),50        (session.deleted, "delete"),51    ):52        for target in targets:53            state = sa.inspect(target)54            key = state.identity_key if state.has_identity else id(target)55            session._model_changes[key] = (target, operation)56 57 58def _before_commit(session: Session) -> None:59    if not has_app_context():60        return61 62    app = current_app._get_current_object()  # type: ignore[attr-defined]63 64    if not app.config["SQLALCHEMY_TRACK_MODIFICATIONS"]:65        return66 67    if session._model_changes:68        changes = list(session._model_changes.values())69        before_models_committed.send(app, changes=changes)70 71 72def _after_commit(session: Session) -> None:73    if not has_app_context():74        return75 76    app = current_app._get_current_object()  # type: ignore[attr-defined]77 78    if not app.config["SQLALCHEMY_TRACK_MODIFICATIONS"]:79        return80 81    if session._model_changes:82        changes = list(session._model_changes.values())83        models_committed.send(app, changes=changes)84        session._model_changes.clear()85 86 87def _after_rollback(session: Session) -> None:88    session._model_changes.clear()89