Reisacher-Forderung/app/db.py

from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session, declarative_base, sessionmaker

from .config import settings
from .logging_setup import db_operation, db_query, log_query_error, query_origin


def install_engine_logging(engine):
    if getattr(engine, "_gc_logging", False):
        return
    engine._gc_logging = True

    @event.listens_for(engine, "before_cursor_execute")
    def executing(conn, cursor, statement, parameters, context, executemany):
        context._gc_operation = "execute"

    @event.listens_for(engine, "after_cursor_execute")
    def executed(conn, cursor, statement, parameters, context, executemany):
        context._gc_operation = "fetch"

    @event.listens_for(engine, "handle_error")
    def query_failed(context):
        # Observational only: never replace exceptions or affect connection invalidation.
        error = context.sqlalchemy_exception or context.original_exception
        options = getattr(context.execution_context, "execution_options", {}) or {}
        active = db_query.get() or {}
        statement = context.statement if context.statement is not None else active.get("statement")
        log_query_error(
            error,
            statement,
            options.get("query_name") or active.get("query_name"),
            db_operation.get() or getattr(context.execution_context, "_gc_operation", "connect_or_transaction"),
            parameters=context.parameters,
        )


class LoggingSession(Session):
    def execute(self, statement, params=None, *, execution_options=None, **kwargs):
        options = dict(execution_options or {})
        origin = options.setdefault(
            "query_name", getattr(statement, "_execution_options", {}).get("query_name") or query_origin()
        )
        token = db_operation.set("execute")
        query_token = db_query.set({"statement": statement, "query_name": origin})
        try:
            return super().execute(statement, params, execution_options=options, **kwargs)
        except Exception as error:
            log_query_error(error, statement, origin, "execute", parameters=params)
            raise
        finally:
            db_query.reset(query_token)
            db_operation.reset(token)

    def _transaction_operation(self, operation):
        token = db_operation.set(operation)
        try:
            return getattr(super(), operation)()
        except Exception as error:
            log_query_error(error, operation=operation)
            raise
        finally:
            db_operation.reset(token)

    def commit(self):
        return self._transaction_operation("commit")

    def rollback(self):
        return self._transaction_operation("rollback")

    def flush(self, objects=None):
        token = db_operation.set("flush")
        try:
            return super().flush(objects)
        except Exception as error:
            log_query_error(error, operation="flush")
            raise
        finally:
            db_operation.reset(token)


engine = create_engine(
    settings.DATABASE_URL,
    connect_args={"check_same_thread": False} if settings.DATABASE_URL.startswith("sqlite") else {},
    hide_parameters=True,
)
install_engine_logging(engine)
SessionLocal = sessionmaker(bind=engine, class_=LoggingSession, autoflush=False, autocommit=False)
Base = declarative_base()


def get_session():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()