Fix OAuth w/ ext_perm_user for multi-tenant (onyx-dot-app#4723)

* Fix OAuth w/ ext_perm_user for multi-tenant

* Improve comment

Author: Weves

backend/onyx/auth/users.py

@@ -463,15 +463,8 @@ async def oauth_callback(
                     }
 
                     user = await self.user_db.create(user_dict)
-
-                    # Add OAuth account only if user creation was successful
-                    if user is not None:
-                        await self.user_db.add_oauth_account(user, oauth_account_dict)
-                        await self.on_after_register(user, request)
-                    else:
-                        raise HTTPException(
-                            status_code=500, detail="Failed to create user account"
-                        )
+                    await self.user_db.add_oauth_account(user, oauth_account_dict)
+                    await self.on_after_register(user, request)
 
             else:
                 # User exists, update OAuth account if needed
@@ -489,12 +482,6 @@ async def oauth_callback(
                                 oauth_account_dict,
                             )
 
-            # Ensure user is not None before proceeding
-            if user is None:
-                raise HTTPException(
-                    status_code=500, detail="Failed to authenticate or create user"
-                )
-
             # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
             # re-authenticate that frequently, so by default this is disabled
             if expires_at and TRACK_EXTERNAL_IDP_EXPIRY:
@@ -505,14 +492,20 @@ async def oauth_callback(
 
             # Handle case where user has used product outside of web and is now creating an account through web
             if not user.role.is_web_login():
+                # We must use the existing user in the session if it matches
+                # the user we just got by email/oauth
+                if user.id:
+                    user_by_session = await db_session.get(User, user.id)
+                    if user_by_session:
+                        user = user_by_session
+
                 await self.user_db.update(
                     user,
                     {
                         "is_verified": is_verified_by_default,
                         "role": UserRole.BASIC,
                     },
                 )
-                user.is_verified = is_verified_by_default
 
             # this is needed if an organization goes from `TRACK_EXTERNAL_IDP_EXPIRY=true` to `false`
             # otherwise, the oidc expiry will always be old, and the user will never be able to login

backend/onyx/db/engine.py

@@ -497,7 +497,9 @@ async def get_async_session(
     engine = get_sqlalchemy_async_engine()
 
     async with AsyncSession(engine, expire_on_commit=False) as async_session:
-        # set the search path on sync session as well to be extra safe
+        # IMPORTANT: do NOT remove. The search_path seems to get reset on every `.commit()`
+        # without this. Do not fully understand why atm
+        async_session.info["tenant_id"] = tenant_id
         event.listen(
             async_session.sync_session,
             "after_begin",