Trace sqla get_session

danswer-ai · Feb 29, 2024 · d07345c · d07345c
1 parent c7d228e
commit d07345c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/backend/danswer/db/engine.py b/backend/danswer/db/engine.py
@@ -2,6 +2,7 @@
 from collections.abc import Generator
 from datetime import datetime
 
+from ddtrace import tracer
 from sqlalchemy import text
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine import Engine
@@ -68,6 +69,7 @@ def get_sqlalchemy_async_engine() -> AsyncEngine:
     return _ASYNC_ENGINE
 
 
+@tracer.wrap("db.get_session")
 def get_session() -> Generator[Session, None, None]:
     with Session(get_sqlalchemy_engine(), expire_on_commit=False) as session:
         yield session