Spaces:
Running
Running
File size: 1,917 Bytes
5d6c194 91e0f22 5d6c194 91e0f22 5d6c194 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from contextlib import contextmanager
from datetime import datetime
import time
from sqlalchemy import event
from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker
from app.core.config import get_settings
class Base(DeclarativeBase):
pass
settings = get_settings()
engine = create_engine(settings.database_url, pool_pre_ping=True, future=True)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)
@event.listens_for(engine, "before_cursor_execute")
def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
conn.info.setdefault("blum_query_start_time", []).append(time.perf_counter())
conn.info.setdefault("blum_query_wall_start_time", []).append(datetime.utcnow())
@event.listens_for(engine, "after_cursor_execute")
def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
start_stack = conn.info.get("blum_query_start_time", [])
wall_stack = conn.info.get("blum_query_wall_start_time", [])
started = start_stack.pop() if start_stack else time.perf_counter()
wall_started = wall_stack.pop() if wall_stack else datetime.utcnow()
duration_ms = (time.perf_counter() - started) * 1000
try:
from app.services.performance import performance_recorder
performance_recorder.record_db_query(
statement=statement,
duration_ms=duration_ms,
rowcount=getattr(cursor, "rowcount", None),
started_at=wall_started,
parameters=parameters,
)
except Exception:
pass
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
@contextmanager
def session_scope():
db = SessionLocal()
try:
yield db
db.commit()
except Exception:
db.rollback()
raise
finally:
db.close()
|