From ec48142a2d388090a4cdc877ce2442042488d0ce Mon Sep 17 00:00:00 2001
From: Weves <chrisweaver101@gmail.com>
Date: Fri, 22 Mar 2024 15:20:12 -0700
Subject: [PATCH] Move some of the user re-work stuff to MIT repo

---
 backend/danswer/auth/users.py          | 33 ++++++++++++++++++--------
 backend/danswer/server/manage/users.py |  4 ++--
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py
index 31bdc41a208..975358b6cd0 100644
--- a/backend/danswer/auth/users.py
+++ b/backend/danswer/auth/users.py
@@ -279,13 +279,32 @@ fastapi_users = FastAPIUserWithLogoutRouter[User, uuid.UUID](
 # take care of that in `double_check_user` ourself. This is needed, since
 # we want the /me endpoint to still return a user even if they are not
 # yet verified, so that the frontend knows they exist
-optional_valid_user = fastapi_users.current_user(active=True, optional=True)
+optional_fastapi_current_user = fastapi_users.current_user(active=True, optional=True)
 
 
-async def double_check_user(
+async def optional_user_(
     request: Request,
     user: User | None,
     db_session: Session,
+) -> User | None:
+    """NOTE: `request` and `db_session` are not used here, but are included
+    for the EE version of this function."""
+    return user
+
+
+async def optional_user(
+    request: Request,
+    user: User | None = Depends(optional_fastapi_current_user),
+    db_session: Session = Depends(get_session),
+) -> User | None:
+    versioned_fetch_user = fetch_versioned_implementation(
+        "danswer.auth.users", "optional_user_"
+    )
+    return await versioned_fetch_user(request, user, db_session)
+
+
+async def double_check_user(
+    user: User | None,
     optional: bool = DISABLE_AUTH,
 ) -> User | None:
     if optional:
@@ -307,15 +326,9 @@ async def double_check_user(
 
 
 async def current_user(
-    request: Request,
-    user: User | None = Depends(optional_valid_user),
-    db_session: Session = Depends(get_session),
+    user: User | None = Depends(optional_user),
 ) -> User | None:
-    double_check_user = fetch_versioned_implementation(
-        "danswer.auth.users", "double_check_user"
-    )
-    user = await double_check_user(request, user, db_session)
-    return user
+    return await double_check_user(user)
 
 
 async def current_admin_user(user: User | None = Depends(current_user)) -> User | None:
diff --git a/backend/danswer/server/manage/users.py b/backend/danswer/server/manage/users.py
index 539d7212f2d..635d6dbe560 100644
--- a/backend/danswer/server/manage/users.py
+++ b/backend/danswer/server/manage/users.py
@@ -11,7 +11,7 @@ from danswer.auth.schemas import UserRead
 from danswer.auth.schemas import UserRole
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
-from danswer.auth.users import optional_valid_user
+from danswer.auth.users import optional_user
 from danswer.db.engine import get_session
 from danswer.db.engine import get_sqlalchemy_async_engine
 from danswer.db.models import User
@@ -57,7 +57,7 @@ async def get_user_role(user: User = Depends(current_user)) -> UserRoleResponse:
 
 
 @router.get("/me")
-def verify_user_logged_in(user: User | None = Depends(optional_valid_user)) -> UserInfo:
+def verify_user_logged_in(user: User | None = Depends(optional_user)) -> UserInfo:
     # NOTE: this does not use `current_user` / `current_admin_user` because we don't want
     # to enforce user verification here - the frontend always wants to get the info about
     # the current user regardless of if they are currently verified