From c1f07da4c86ca15db361c6a5c08cab833e9c31e0 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Wed, 15 Jan 2025 10:10:11 -0800
Subject: [PATCH] functional

---
 backend/onyx/auth/users.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py
index 2deb4defd93..ea13ce63b90 100644
--- a/backend/onyx/auth/users.py
+++ b/backend/onyx/auth/users.py
@@ -46,6 +46,7 @@
 from httpx_oauth.oauth2 import BaseOAuth2
 from httpx_oauth.oauth2 import OAuth2Token
 from pydantic import BaseModel
+from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from onyx.auth.api_key import get_hashed_api_key_from_request
@@ -371,6 +372,8 @@ async def oauth_callback(
 
             except exceptions.UserNotExists:
                 try:
+                    # Set the search path for the current session
+                    await db_session.execute(text(f'SET search_path = "{tenant_id}"'))
                     # Attempt to get user by email
                     user = await self.get_by_email(account_email)
                     if not associate_by_email:
@@ -395,7 +398,9 @@ async def oauth_callback(
                     # OAuth account creation happens in the correct tenant schema
 
                     # Add OAuth account
+                    await db_session.execute(text(f'SET search_path = "{tenant_id}"'))
                     await self.user_db.add_oauth_account(user, oauth_account_dict)
+
                     await self.on_after_register(user, request)
 
             else: