From 2331bf9b367ae63ee51716ee34ee9c164e08fc08 Mon Sep 17 00:00:00 2001 From: Weves Date: Thu, 29 Feb 2024 13:49:16 -0800 Subject: [PATCH] Add trace for db session creation --- backend/danswer/db/engine.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/backend/danswer/db/engine.py b/backend/danswer/db/engine.py index b28647f6d..e355084d5 100644 --- a/backend/danswer/db/engine.py +++ b/backend/danswer/db/engine.py @@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator 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 @@ -69,8 +70,9 @@ def get_sqlalchemy_async_engine() -> AsyncEngine: def get_session() -> Generator[Session, None, None]: - with Session(get_sqlalchemy_engine(), expire_on_commit=False) as session: - yield session + with tracer.trace("db.get_session"): + with Session(get_sqlalchemy_engine(), expire_on_commit=False) as session: + yield session async def get_async_session() -> AsyncGenerator[AsyncSession, None]: