refac: async db
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user