From e5af4681d307b6a9ac1c4f96233e00b86b1f64b4 Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Mon, 28 Oct 2024 12:44:45 -0700 Subject: [PATCH] Fix nagging double auth issue (#2960) * fix nagging double auth issue * ports --- backend/danswer/auth/users.py | 5 +++++ backend/danswer/db/auth.py | 5 ++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py index 13a0902b8..aa5054983 100644 --- a/backend/danswer/auth/users.py +++ b/backend/danswer/auth/users.py @@ -49,6 +49,7 @@ from httpx_oauth.oauth2 import BaseOAuth2 from httpx_oauth.oauth2 import OAuth2Token from pydantic import BaseModel from sqlalchemy import select +from sqlalchemy import text from sqlalchemy.orm import attributes from sqlalchemy.orm import Session @@ -366,6 +367,10 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): } user = await self.user_db.create(user_dict) + + # Explicitly set the Postgres schema for this session to ensure + # OAuth account creation happens in the correct tenant schema + await db_session.execute(text(f'SET search_path = "{tenant_id}"')) user = await self.user_db.add_oauth_account( user, oauth_account_dict ) diff --git a/backend/danswer/db/auth.py b/backend/danswer/db/auth.py index 9eba3806d..53fba1e52 100644 --- a/backend/danswer/db/auth.py +++ b/backend/danswer/db/auth.py @@ -57,7 +57,10 @@ async def get_user_count() -> int: # Need to override this because FastAPI Users doesn't give flexibility for backend field creation logic in OAuth flow class SQLAlchemyUserAdminDB(SQLAlchemyUserDatabase): - async def create(self, create_dict: Dict[str, Any]) -> UP: + async def create( + self, + create_dict: Dict[str, Any], + ) -> UP: user_count = await get_user_count() if user_count == 0 or create_dict["email"] in get_default_admin_user_emails(): create_dict["role"] = UserRole.ADMIN