From c3ead772f1cf50b0e33d1ac084e269291d818fa4 Mon Sep 17 00:00:00 2001 From: PVBLIC Foundation Date: Mon, 11 May 2026 10:54:58 -0700 Subject: [PATCH] fix: migrate local Gmail/Drive/spend-limit paths to async after v0.9.5 Upstream v0.9.5 made Users.*, OAuthSessions.*, Files.*, and Knowledges.get_knowledge_by_id async. Local code that wasn't part of the merge still called these synchronously, which caused both an import-time crash (knowledge.py used `db: Session = Depends(get_session)` without importing the symbols) and runtime breakage in Gmail/Drive sync. Changes: - routers/knowledge.py: 6 Drive endpoints now use AsyncSession; awaited Knowledges.get_knowledge_by_id and has_access calls - routers/users.py: 6 spend-limit endpoints now use AsyncSession; awaited Users.get_user_by_id and Users.update_user_by_id; also awaited the force_gmail_sync admin endpoint's Users + OAuthSessions calls - routers/gmail.py: awaited Users.get_user_by_id (5x), Users.update_user_by_id (5x), OAuthSessions .get_session_by_provider_and_user_id (3x) - utils/spend_limit.py: made check_user_spend_limit async; updated enforce_spend_limit to await it; awaited Users.get_user_by_id - utils/gmail_auto_sync.py: awaited Users.get_user_by_id (3x), OAuthSessions.get_session_by_provider_and_user_id (2x), OAuthSessions.get_sessions_by_user_id - utils/knowledge_drive_sync.py: awaited Files.insert_new_file, Files.get_file_by_id (2x), Files.update_file_data_by_id (2x), Users.get_user_by_id UserUsages.* methods stay sync (they use the local sync get_db_context fallback restored in 14321d1b3) and continue to work when called with either sync or async sessions thanks to the isinstance(Session) check. Made-with: Cursor --- backend/open_webui/routers/gmail.py | 26 +++++++------- backend/open_webui/routers/knowledge.py | 36 +++++++++---------- backend/open_webui/routers/users.py | 26 +++++++------- backend/open_webui/utils/gmail_auto_sync.py | 12 +++---- .../open_webui/utils/knowledge_drive_sync.py | 12 +++---- backend/open_webui/utils/spend_limit.py | 6 ++-- 6 files changed, 59 insertions(+), 59 deletions(-) diff --git a/backend/open_webui/routers/gmail.py b/backend/open_webui/routers/gmail.py index 564756360f..33fdc872ed 100644 --- a/backend/open_webui/routers/gmail.py +++ b/backend/open_webui/routers/gmail.py @@ -34,7 +34,7 @@ async def get_gmail_status( Returns sync status, last sync time, email counts, etc. """ - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") @@ -45,7 +45,7 @@ async def get_gmail_status( gmail_settings = settings_dict.get("gmail", {}) if isinstance(settings_dict, dict) else {} # Check if user has Google OAuth session - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) has_gmail_oauth = oauth_session is not None # Check if OAuth token has Gmail scopes @@ -78,12 +78,12 @@ async def enable_gmail_sync( Does NOT trigger sync - user must click "Sync Now" separately. """ - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") # Check if user has Google OAuth with Gmail scopes - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) if not oauth_session: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="User must log in with Google OAuth first") @@ -109,7 +109,7 @@ async def enable_gmail_sync( "total_vectors": existing_gmail.get("total_vectors", 0), } - Users.update_user_by_id(user_id, {"settings": user_settings}) + await Users.update_user_by_id(user_id, {"settings": user_settings}) logger.info(f"✅ Gmail sync enabled for user {user_id}") @@ -134,7 +134,7 @@ async def trigger_gmail_sync( - User has valid Google OAuth session with Gmail scopes """ - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") @@ -149,7 +149,7 @@ async def trigger_gmail_sync( raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Gmail sync is not enabled for this user") # Get OAuth session and token (with automatic refresh if expired) - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) if not oauth_session: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -186,7 +186,7 @@ async def trigger_gmail_sync( **gmail_settings, "sync_status": "syncing", } - Users.update_user_by_id(user_id, {"settings": user_settings}) + await Users.update_user_by_id(user_id, {"settings": user_settings}) # Trigger background sync task with refreshed token try: @@ -209,7 +209,7 @@ async def trigger_gmail_sync( # Reset status to ready on error user_settings["gmail"]["sync_status"] = "error" - Users.update_user_by_id(user_id, {"settings": user_settings}) + await Users.update_user_by_id(user_id, {"settings": user_settings}) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to start Gmail sync: {str(e)}" @@ -228,7 +228,7 @@ async def disable_gmail_sync( Sets sync_enabled to False. Does not delete existing data. """ - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") @@ -246,7 +246,7 @@ async def disable_gmail_sync( "sync_status": "disabled", } - Users.update_user_by_id(user_id, {"settings": user_settings}) + await Users.update_user_by_id(user_id, {"settings": user_settings}) logger.info(f"🛑 Gmail sync disabled for user {user_id}") @@ -265,7 +265,7 @@ async def delete_gmail_data( Removes all vectors and disables sync. """ - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") @@ -291,7 +291,7 @@ async def delete_gmail_data( "total_vectors": 0, } - Users.update_user_by_id(user_id, {"settings": user_settings}) + await Users.update_user_by_id(user_id, {"settings": user_settings}) return {"status": "deleted", "message": "All Gmail data has been deleted from Pinecone"} diff --git a/backend/open_webui/routers/knowledge.py b/backend/open_webui/routers/knowledge.py index 702ccc712f..e5f0b5ee69 100644 --- a/backend/open_webui/routers/knowledge.py +++ b/backend/open_webui/routers/knowledge.py @@ -1257,12 +1257,12 @@ class DriveSyncResponse(BaseModel): async def get_drive_sources( id: str, user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get all Google Drive sources connected to a Knowledge base. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1272,7 +1272,7 @@ async def get_drive_sources( # Check access if ( knowledge.user_id != user.id - and not has_access(user.id, "read", knowledge.access_control, db=db) + and not await has_access(user.id, "read", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( @@ -1316,7 +1316,7 @@ async def connect_drive_folder( id: str, form_data: KnowledgeDriveSourceForm, user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Connect a Google Drive folder to a Knowledge base. @@ -1324,7 +1324,7 @@ async def connect_drive_folder( The folder will be synced automatically based on the configured interval. Requires user to have Google OAuth with Drive scope. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1334,7 +1334,7 @@ async def connect_drive_folder( # Check write access if ( knowledge.user_id != user.id - and not has_access(user.id, "write", knowledge.access_control, db=db) + and not await has_access(user.id, "write", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( @@ -1429,14 +1429,14 @@ async def disconnect_drive_folder( id: str, source_id: str, user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Disconnect a Google Drive folder from a Knowledge base. This removes the sync connection but does not delete files already synced. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1446,7 +1446,7 @@ async def disconnect_drive_folder( # Check write access if ( knowledge.user_id != user.id - and not has_access(user.id, "write", knowledge.access_control, db=db) + and not await has_access(user.id, "write", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( @@ -1482,7 +1482,7 @@ async def sync_drive_source( source_id: str, force_full: bool = Query(False, description="Force full sync instead of incremental"), user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Manually trigger a sync for a Drive source. @@ -1490,7 +1490,7 @@ async def sync_drive_source( By default, performs incremental sync (only changed files). Set force_full=true to resync all files. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1500,7 +1500,7 @@ async def sync_drive_source( # Check write access if ( knowledge.user_id != user.id - and not has_access(user.id, "write", knowledge.access_control, db=db) + and not await has_access(user.id, "write", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( @@ -1566,12 +1566,12 @@ async def sync_all_drive_sources( id: str, force_full: bool = Query(False, description="Force full sync instead of incremental"), user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Trigger sync for all Drive sources connected to a Knowledge base. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1581,7 +1581,7 @@ async def sync_all_drive_sources( # Check write access if ( knowledge.user_id != user.id - and not has_access(user.id, "write", knowledge.access_control, db=db) + and not await has_access(user.id, "write", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( @@ -1644,12 +1644,12 @@ async def update_drive_source( source_id: str, form_data: DriveSourceUpdateForm, user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Update settings for a Drive source. """ - knowledge = Knowledges.get_knowledge_by_id(id=id, db=db) + knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db) if not knowledge: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1659,7 +1659,7 @@ async def update_drive_source( # Check write access if ( knowledge.user_id != user.id - and not has_access(user.id, "write", knowledge.access_control, db=db) + and not await has_access(user.id, "write", knowledge.access_control, db=db) and user.role != "admin" ): raise HTTPException( diff --git a/backend/open_webui/routers/users.py b/backend/open_webui/routers/users.py index bff1ef4644..73da15cddb 100644 --- a/backend/open_webui/routers/users.py +++ b/backend/open_webui/routers/users.py @@ -412,7 +412,7 @@ async def update_user_info_by_session_user( @router.get("/me/spend", response_model=UserSpendSummary) async def get_my_spend( user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get current user's spend summary. @@ -425,14 +425,14 @@ async def get_my_spend( @router.get("/me/spend-limits", response_model=dict) async def get_my_spend_limits( user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get current user's spend limits and current usage. Returns whether limits are enabled, the limit values, and current usage. """ - current_user = Users.get_user_by_id(user.id, db=db) + current_user = await Users.get_user_by_id(user.id, db=db) spend_summary = UserUsages.get_user_spend_summary(user.id, db=db) return { @@ -454,7 +454,7 @@ async def get_my_usage_history( skip: int = 0, limit: int = 30, user=Depends(get_verified_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get current user's usage history. @@ -889,7 +889,7 @@ async def force_gmail_sync(user_id: str, request: Request, user=Depends(get_admi from open_webui.utils.gmail_auto_sync import trigger_gmail_sync_if_needed # Validate user exists - target_user = Users.get_user_by_id(user_id) + target_user = await Users.get_user_by_id(user_id) if not target_user: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found") @@ -898,7 +898,7 @@ async def force_gmail_sync(user_id: str, request: Request, user=Depends(get_admi raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Gmail sync is not enabled for this user") # Check if user has Google OAuth session - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) if not oauth_session: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="User has not connected their Google account" @@ -964,14 +964,14 @@ async def force_gmail_sync(user_id: str, request: Request, user=Depends(get_admi async def get_user_spend_limits( user_id: str, user=Depends(get_admin_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get spend limit configuration for a user (admin only). Returns the user's spend limits and current usage. """ - target_user = Users.get_user_by_id(user_id, db=db) + target_user = await Users.get_user_by_id(user_id, db=db) if not target_user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1000,7 +1000,7 @@ async def update_user_spend_limits( user_id: str, form_data: UserSpendLimitForm, user=Depends(get_admin_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Update spend limit configuration for a user (admin only). @@ -1008,7 +1008,7 @@ async def update_user_spend_limits( Set daily and/or monthly spend limits in USD. Set limits to null to remove that limit. """ - target_user = Users.get_user_by_id(user_id, db=db) + target_user = await Users.get_user_by_id(user_id, db=db) if not target_user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1022,7 +1022,7 @@ async def update_user_spend_limits( "spend_limit_monthly": form_data.spend_limit_monthly, } - updated_user = Users.update_user_by_id(user_id, update_data, db=db) + updated_user = await Users.update_user_by_id(user_id, update_data, db=db) if not updated_user: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -1048,14 +1048,14 @@ async def get_user_usage_history( skip: int = 0, limit: int = 30, user=Depends(get_admin_user), - db: Session = Depends(get_session), + db: AsyncSession = Depends(get_async_session), ): """ Get usage history for a user (admin only). Returns paginated list of usage records with totals. """ - target_user = Users.get_user_by_id(user_id, db=db) + target_user = await Users.get_user_by_id(user_id, db=db) if not target_user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, diff --git a/backend/open_webui/utils/gmail_auto_sync.py b/backend/open_webui/utils/gmail_auto_sync.py index 9f981ef4af..00e3be64b9 100644 --- a/backend/open_webui/utils/gmail_auto_sync.py +++ b/backend/open_webui/utils/gmail_auto_sync.py @@ -691,7 +691,7 @@ async def trigger_gmail_sync_if_needed( return # Check if admin has enabled Gmail sync for this user first - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: logger.info(f" ⏭️ SKIP: User {user_id} not found") return @@ -780,7 +780,7 @@ async def _background_gmail_sync(request, user_id: str, oauth_token: dict, force try: # Validate user - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: logger.error(f"User {user_id} not found") return @@ -806,7 +806,7 @@ async def _background_gmail_sync(request, user_id: str, oauth_token: dict, force """Refresh OAuth token using oauth_manager""" try: # Get OAuth session for this user - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) if not oauth_session: logger.error(f"No OAuth session found for user {user_id}") return None @@ -1432,7 +1432,7 @@ async def _sync_user_periodic(user_id: str) -> bool: logger.info(f"🔄 Periodic sync starting for user: {user_id}") # Validation: Check user exists - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: logger.warning(f"⚠️ User {user_id} not found in database, skipping") return False @@ -1448,7 +1448,7 @@ async def _sync_user_periodic(user_id: str) -> bool: # Validation: Check OAuth session exists # Debug: List all OAuth sessions for this user to see what providers exist - all_sessions = OAuthSessions.get_sessions_by_user_id(user_id) + all_sessions = await OAuthSessions.get_sessions_by_user_id(user_id) if all_sessions: logger.info(f" Found {len(all_sessions)} OAuth session(s) for user:") for s in all_sessions: @@ -1456,7 +1456,7 @@ async def _sync_user_periodic(user_id: str) -> bool: else: logger.info(f" No OAuth sessions found for user {user_id}") - oauth_session = OAuthSessions.get_session_by_provider_and_user_id("google", user_id) + oauth_session = await OAuthSessions.get_session_by_provider_and_user_id("google", user_id) if not oauth_session: logger.info(f"⏭️ No Google OAuth session for user {user_id}, skipping") return False diff --git a/backend/open_webui/utils/knowledge_drive_sync.py b/backend/open_webui/utils/knowledge_drive_sync.py index ce9fcf8da7..03327c9b8f 100644 --- a/backend/open_webui/utils/knowledge_drive_sync.py +++ b/backend/open_webui/utils/knowledge_drive_sync.py @@ -484,7 +484,7 @@ class KnowledgeDriveSyncService: # Create file record log.info(f" 📝 Creating file record...") - file_record = Files.insert_new_file( + file_record = await Files.insert_new_file( user_id, FileForm( id=file_id, @@ -540,12 +540,12 @@ class KnowledgeDriveSyncService: from open_webui.models.users import Users # Get user for auth - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: raise ValueError(f"User {user_id} not found") # Get file record - file_record = Files.get_file_by_id(file_id) + file_record = await Files.get_file_by_id(file_id) if not file_record: raise ValueError(f"File {file_id} not found") @@ -585,7 +585,7 @@ class KnowledgeDriveSyncService: ) # Verify content was extracted - file_check = Files.get_file_by_id(file_id) + file_check = await Files.get_file_by_id(file_id) if not file_check or not file_check.data or not file_check.data.get("content"): # Check if file status indicates failure if file_check and file_check.data and file_check.data.get("status") == "failed": @@ -618,7 +618,7 @@ class KnowledgeDriveSyncService: if not result2 or not result2.get("status"): log.error(f" ❌ Step 2 failed for {file_id[:8]}: {result2}") - Files.update_file_data_by_id( + await Files.update_file_data_by_id( file_id, {"status": "error", "error": "Failed to add to knowledge base"} ) return @@ -634,7 +634,7 @@ class KnowledgeDriveSyncService: log.error(f" Traceback: {traceback.format_exc()}") # Don't raise - file was synced, just RAG failed # Mark file with error status - Files.update_file_data_by_id(file_id, {"status": "error", "error": str(e)}) + await Files.update_file_data_by_id(file_id, {"status": "error", "error": str(e)}) def get_sync_status(self, source_id: str) -> Optional[Dict[str, Any]]: """Get current sync status for a source""" diff --git a/backend/open_webui/utils/spend_limit.py b/backend/open_webui/utils/spend_limit.py index ada33c936d..fedd7f7846 100644 --- a/backend/open_webui/utils/spend_limit.py +++ b/backend/open_webui/utils/spend_limit.py @@ -34,7 +34,7 @@ class SpendLimitExceeded(HTTPException): ) -def check_user_spend_limit(user_id: str) -> Tuple[bool, Optional[dict]]: +async def check_user_spend_limit(user_id: str) -> Tuple[bool, Optional[dict]]: """ Check if a user has exceeded their spend limits. @@ -50,7 +50,7 @@ def check_user_spend_limit(user_id: str) -> Tuple[bool, Optional[dict]]: SpendLimitExceeded: If user has exceeded their limit """ try: - user = Users.get_user_by_id(user_id) + user = await Users.get_user_by_id(user_id) if not user: return True, None @@ -117,4 +117,4 @@ async def enforce_spend_limit(user) -> None: # if user.role == "admin": # return - check_user_spend_limit(user.id) + await check_user_spend_limit(user.id)