From 7c4487585d67396b8aca0aba740adb64e82f2cc7 Mon Sep 17 00:00:00 2001 From: rkuo-danswer Date: Mon, 14 Apr 2025 14:48:35 -0700 Subject: [PATCH] rollback properly on exception (#4073) * rollback properly on exception * rollback on exception * don't continue if we can't set the search path * cleaner handling via context manager --------- Co-authored-by: Richard Kuo (Danswer) --- backend/onyx/db/engine.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/backend/onyx/db/engine.py b/backend/onyx/db/engine.py index 04419b5a33..41bde426e3 100644 --- a/backend/onyx/db/engine.py +++ b/backend/onyx/db/engine.py @@ -487,19 +487,23 @@ def get_session_with_tenant(*, tenant_id: str) -> Generator[Session, None, None] f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}" ) ) + except Exception: + raise RuntimeError(f"search_path not set for {tenant_id}") finally: cursor.close() - with Session(bind=connection, expire_on_commit=False) as session: - try: + try: + # automatically rollback or close + with Session(bind=connection, expire_on_commit=False) as session: yield session - finally: - if MULTI_TENANT: - cursor = dbapi_connection.cursor() - try: - cursor.execute('SET search_path TO "$user", public') - finally: - cursor.close() + finally: + # always reset the search path on exit + if MULTI_TENANT: + cursor = dbapi_connection.cursor() + try: + cursor.execute('SET search_path TO "$user", public') + finally: + cursor.close() def set_search_path_on_checkout(