refac: async db

This commit is contained in:
Timothy Jaeryang Baek
2026-04-12 14:22:11 -05:00
parent b618d84065
commit 27169124f2
74 changed files with 4831 additions and 4479 deletions
+55 -55
View File
@@ -1,6 +1,6 @@
import logging
from typing import Optional
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
import base64
import io
@@ -30,7 +30,7 @@ from open_webui.models.users import (
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import STATIC_DIR
from open_webui.internal.db import get_session
from open_webui.internal.db import get_async_session
from open_webui.utils.auth import (
@@ -63,7 +63,7 @@ async def get_users(
direction: Optional[str] = None,
page: Optional[int] = 1,
user=Depends(get_admin_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
limit = PAGE_ITEM_COUNT
@@ -80,14 +80,14 @@ async def get_users(
filter['direction'] = direction
result = Users.get_users(filter=filter, skip=skip, limit=limit, db=db)
result = await Users.get_users(filter=filter, skip=skip, limit=limit, db=db)
users = result['users']
total = result['total']
# Fetch groups for all users in a single query to avoid N+1
user_ids = [user.id for user in users]
user_groups = Groups.get_groups_by_member_ids(user_ids, db=db)
user_groups = await Groups.get_groups_by_member_ids(user_ids, db=db)
return {
'users': [
@@ -106,9 +106,9 @@ async def get_users(
@router.get('/all', response_model=UserInfoListResponse)
async def get_all_users(
user=Depends(get_admin_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
return Users.get_users(db=db)
return await Users.get_users(db=db)
@router.get('/search', response_model=UserInfoListResponse)
@@ -118,7 +118,7 @@ async def search_users(
direction: Optional[str] = None,
page: Optional[int] = 1,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
limit = PAGE_ITEM_COUNT
@@ -133,7 +133,7 @@ async def search_users(
if direction:
filter['direction'] = direction
return Users.get_users(filter=filter, skip=skip, limit=limit, db=db)
return await Users.get_users(filter=filter, skip=skip, limit=limit, db=db)
############################
@@ -142,8 +142,8 @@ async def search_users(
@router.get('/groups')
async def get_user_groups(user=Depends(get_verified_user), db: Session = Depends(get_session)):
return Groups.get_groups_by_member_id(user.id, db=db)
async def get_user_groups(user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
return await Groups.get_groups_by_member_id(user.id, db=db)
############################
@@ -155,9 +155,9 @@ async def get_user_groups(user=Depends(get_verified_user), db: Session = Depends
async def get_user_permissisions(
request: Request,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
user_permissions = get_permissions(user.id, request.app.state.config.USER_PERMISSIONS, db=db)
user_permissions = await get_permissions(user.id, request.app.state.config.USER_PERMISSIONS, db=db)
return user_permissions
@@ -272,8 +272,8 @@ async def update_default_user_permissions(request: Request, form_data: UserPermi
@router.get('/user/settings', response_model=Optional[UserSettings])
async def get_user_settings_by_session_user(user=Depends(get_verified_user), db: Session = Depends(get_session)):
user = Users.get_user_by_id(user.id, db=db)
async def get_user_settings_by_session_user(user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
user = await Users.get_user_by_id(user.id, db=db)
if user:
return user.settings
else:
@@ -293,7 +293,7 @@ async def update_user_settings_by_session_user(
request: Request,
form_data: UserSettings,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
updated_user_settings = form_data.model_dump()
ui_settings = updated_user_settings.get('ui')
@@ -301,7 +301,7 @@ async def update_user_settings_by_session_user(
user.role != 'admin'
and ui_settings is not None
and 'toolServers' in ui_settings.keys()
and not has_permission(
and not await has_permission(
user.id,
'features.direct_tool_servers',
request.app.state.config.USER_PERMISSIONS,
@@ -310,7 +310,7 @@ async def update_user_settings_by_session_user(
# If the user is not an admin and does not have permission to use tool servers, remove the key
updated_user_settings['ui'].pop('toolServers', None)
user = Users.update_user_settings_by_id(user.id, updated_user_settings, db=db)
user = await Users.update_user_settings_by_id(user.id, updated_user_settings, db=db)
if user:
return user.settings
else:
@@ -329,14 +329,14 @@ async def update_user_settings_by_session_user(
async def get_user_status_by_session_user(
request: Request,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
if not request.app.state.config.ENABLE_USER_STATUS:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=ERROR_MESSAGES.ACTION_PROHIBITED,
)
user = Users.get_user_by_id(user.id, db=db)
user = await Users.get_user_by_id(user.id, db=db)
if user:
return user
else:
@@ -356,16 +356,16 @@ async def update_user_status_by_session_user(
request: Request,
form_data: UserStatus,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
if not request.app.state.config.ENABLE_USER_STATUS:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=ERROR_MESSAGES.ACTION_PROHIBITED,
)
user = Users.get_user_by_id(user.id, db=db)
user = await Users.get_user_by_id(user.id, db=db)
if user:
user = Users.update_user_status_by_id(user.id, form_data, db=db)
user = await Users.update_user_status_by_id(user.id, form_data, db=db)
return user
else:
raise HTTPException(
@@ -380,8 +380,8 @@ async def update_user_status_by_session_user(
@router.get('/user/info', response_model=Optional[dict])
async def get_user_info_by_session_user(user=Depends(get_verified_user), db: Session = Depends(get_session)):
user = Users.get_user_by_id(user.id, db=db)
async def get_user_info_by_session_user(user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
user = await Users.get_user_by_id(user.id, db=db)
if user:
return user.info
else:
@@ -398,14 +398,14 @@ async def get_user_info_by_session_user(user=Depends(get_verified_user), db: Ses
@router.post('/user/info/update', response_model=Optional[dict])
async def update_user_info_by_session_user(
form_data: dict, user=Depends(get_verified_user), db: Session = Depends(get_session)
form_data: dict, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)
):
user = Users.get_user_by_id(user.id, db=db)
user = await Users.get_user_by_id(user.id, db=db)
if user:
if user.info is None:
user.info = {}
user = Users.update_user_by_id(user.id, {'info': {**user.info, **form_data}}, db=db)
user = await Users.update_user_by_id(user.id, {'info': {**user.info, **form_data}}, db=db)
if user:
return user.info
else:
@@ -435,12 +435,12 @@ class UserActiveResponse(UserStatus):
@router.get('/{user_id}', response_model=UserActiveResponse)
async def get_user_by_id(user_id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
async def get_user_by_id(user_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
# Check if user_id is a shared chat
# If it is, get the user_id from the chat
if user_id.startswith('shared-'):
chat_id = user_id.replace('shared-', '')
chat = Chats.get_chat_by_id(chat_id)
chat = await Chats.get_chat_by_id(chat_id)
if chat:
user_id = chat.user_id
else:
@@ -449,14 +449,14 @@ async def get_user_by_id(user_id: str, user=Depends(get_admin_user), db: Session
detail=ERROR_MESSAGES.USER_NOT_FOUND,
)
user = Users.get_user_by_id(user_id, db=db)
user = await Users.get_user_by_id(user_id, db=db)
if user:
groups = Groups.get_groups_by_member_id(user_id, db=db)
groups = await Groups.get_groups_by_member_id(user_id, db=db)
return UserActiveResponse(
**{
**user.model_dump(),
'groups': [{'id': group.id, 'name': group.name} for group in groups],
'is_active': Users.is_user_active(user_id, db=db),
'is_active': await Users.is_user_active(user_id, db=db),
}
)
else:
@@ -467,15 +467,15 @@ async def get_user_by_id(user_id: str, user=Depends(get_admin_user), db: Session
@router.get('/{user_id}/info', response_model=UserInfoResponse)
async def get_user_info_by_id(user_id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)):
user = Users.get_user_by_id(user_id, db=db)
async def get_user_info_by_id(user_id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
user = await Users.get_user_by_id(user_id, db=db)
if user:
groups = Groups.get_groups_by_member_id(user_id, db=db)
groups = await Groups.get_groups_by_member_id(user_id, db=db)
return UserInfoResponse(
**{
**user.model_dump(),
'groups': [{'id': group.id, 'name': group.name} for group in groups],
'is_active': Users.is_user_active(user_id, db=db),
'is_active': await Users.is_user_active(user_id, db=db),
}
)
else:
@@ -486,8 +486,8 @@ async def get_user_info_by_id(user_id: str, user=Depends(get_verified_user), db:
@router.get('/{user_id}/oauth/sessions')
async def get_user_oauth_sessions_by_id(user_id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
sessions = OAuthSessions.get_sessions_by_user_id(user_id, db=db)
async def get_user_oauth_sessions_by_id(user_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
sessions = await OAuthSessions.get_sessions_by_user_id(user_id, db=db)
if sessions and len(sessions) > 0:
return sessions
else:
@@ -503,8 +503,8 @@ async def get_user_oauth_sessions_by_id(user_id: str, user=Depends(get_admin_use
@router.get('/{user_id}/profile/image')
def get_user_profile_image_by_id(user_id: str, user=Depends(get_verified_user)):
user = Users.get_user_by_id(user_id)
async def get_user_profile_image_by_id(user_id: str, user=Depends(get_verified_user)):
user = await Users.get_user_by_id(user_id)
if user:
if user.profile_image_url:
# check if it's url or base64
@@ -542,10 +542,10 @@ def get_user_profile_image_by_id(user_id: str, user=Depends(get_verified_user)):
@router.get('/{user_id}/active', response_model=dict)
async def get_user_active_status_by_id(
user_id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
user_id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)
):
return {
'active': Users.is_user_active(user_id, db=db),
'active': await Users.is_user_active(user_id, db=db),
}
@@ -559,11 +559,11 @@ async def update_user_by_id(
user_id: str,
form_data: UserUpdateForm,
session_user=Depends(get_admin_user),
db: Session = Depends(get_session),
db: AsyncSession = Depends(get_async_session),
):
# Prevent modification of the primary admin user by other admins
try:
first_user = Users.get_first_user(db=db)
first_user = await Users.get_first_user(db=db)
if first_user:
if user_id == first_user.id:
if session_user.id != user_id:
@@ -587,11 +587,11 @@ async def update_user_by_id(
detail='Could not verify primary admin status.',
)
user = Users.get_user_by_id(user_id, db=db)
user = await Users.get_user_by_id(user_id, db=db)
if user:
if form_data.email.lower() != user.email:
email_user = Users.get_user_by_email(form_data.email.lower(), db=db)
email_user = await Users.get_user_by_email(form_data.email.lower(), db=db)
if email_user:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
@@ -605,10 +605,10 @@ async def update_user_by_id(
raise HTTPException(400, detail=str(e))
hashed = get_password_hash(form_data.password)
Auths.update_user_password_by_id(user_id, hashed, db=db)
await Auths.update_user_password_by_id(user_id, hashed, db=db)
Auths.update_email_by_id(user_id, form_data.email.lower(), db=db)
updated_user = Users.update_user_by_id(
await Auths.update_email_by_id(user_id, form_data.email.lower(), db=db)
updated_user = await Users.update_user_by_id(
user_id,
{
'role': form_data.role,
@@ -639,10 +639,10 @@ async def update_user_by_id(
@router.delete('/{user_id}', response_model=bool)
async def delete_user_by_id(user_id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
async def delete_user_by_id(user_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
# Prevent deletion of the primary admin user
try:
first_user = Users.get_first_user(db=db)
first_user = await Users.get_first_user(db=db)
if first_user and user_id == first_user.id:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
@@ -656,7 +656,7 @@ async def delete_user_by_id(user_id: str, user=Depends(get_admin_user), db: Sess
)
if user.id != user_id:
result = Auths.delete_auth_by_id(user_id, db=db)
result = await Auths.delete_auth_by_id(user_id, db=db)
if result:
return True
@@ -679,5 +679,5 @@ async def delete_user_by_id(user_id: str, user=Depends(get_admin_user), db: Sess
@router.get('/{user_id}/groups')
async def get_user_groups_by_id(user_id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
return Groups.get_groups_by_member_id(user_id, db=db)
async def get_user_groups_by_id(user_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
return await Groups.get_groups_by_member_id(user_id, db=db)