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()