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) <rkuo@onyx.app>

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:
+        finally:
-                if MULTI_TENANT:
+            # always reset the search path on exit
-                    cursor = dbapi_connection.cursor()
+            if MULTI_TENANT:
-                    try:
+                cursor = dbapi_connection.cursor()
-                        cursor.execute('SET search_path TO "$user", public')
+                try:
-                    finally:
+                    cursor.execute('SET search_path TO "$user", public')
-                        cursor.close()
+                finally:
+                    cursor.close()
 
 
 def set_search_path_on_checkout(