Trace sqla get_session

This commit is contained in:
Kevin Shi 2024-02-29 15:45:53 -05:00 committed by Chris Weaver
parent c7d228e292
commit d07345c533

View File

@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator
from collections.abc import Generator from collections.abc import Generator
from datetime import datetime from datetime import datetime
from ddtrace import tracer
from sqlalchemy import text from sqlalchemy import text
from sqlalchemy.engine import create_engine from sqlalchemy.engine import create_engine
from sqlalchemy.engine import Engine from sqlalchemy.engine import Engine
@ -68,6 +69,7 @@ def get_sqlalchemy_async_engine() -> AsyncEngine:
return _ASYNC_ENGINE return _ASYNC_ENGINE
@tracer.wrap("db.get_session")
def get_session() -> Generator[Session, None, None]: def get_session() -> Generator[Session, None, None]:
with Session(get_sqlalchemy_engine(), expire_on_commit=False) as session: with Session(get_sqlalchemy_engine(), expire_on_commit=False) as session:
yield session yield session