refac: async db
This commit is contained in:
@@ -70,8 +70,8 @@ from open_webui.utils.auth import (
|
||||
get_password_hash,
|
||||
get_http_authorization_cred,
|
||||
)
|
||||
from open_webui.internal.db import get_session
|
||||
from sqlalchemy.orm import Session
|
||||
from open_webui.internal.db import get_async_session
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from open_webui.utils.webhook import post_webhook
|
||||
from open_webui.utils.access_control import get_permissions, has_permission
|
||||
from open_webui.utils.groups import apply_default_group_assignment
|
||||
@@ -96,7 +96,7 @@ log = logging.getLogger(__name__)
|
||||
signin_rate_limiter = RateLimiter(redis_client=get_redis_client(), limit=5 * 3, window=60 * 3)
|
||||
|
||||
|
||||
def create_session_response(request: Request, user, db, response: Response = None, set_cookie: bool = False) -> dict:
|
||||
async def create_session_response(request: Request, user, db, response: Response = None, set_cookie: bool = False) -> dict:
|
||||
"""
|
||||
Create JWT token and build session response for a user.
|
||||
Shared helper for signin, signup, ldap_auth, add_user, and token_exchange endpoints.
|
||||
@@ -131,7 +131,7 @@ def create_session_response(request: Request, user, db, response: Response = Non
|
||||
**({'max_age': max_age} if max_age is not None else {}),
|
||||
)
|
||||
|
||||
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 {
|
||||
'token': token,
|
||||
@@ -167,7 +167,7 @@ async def get_session_user(
|
||||
request: Request,
|
||||
response: Response,
|
||||
user=Depends(get_current_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
auth_header = request.headers.get('Authorization')
|
||||
auth_token = get_http_authorization_cred(auth_header)
|
||||
@@ -197,7 +197,7 @@ async def get_session_user(
|
||||
**({'max_age': max_age} if max_age is not None else {}),
|
||||
)
|
||||
|
||||
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 {
|
||||
'token': token,
|
||||
@@ -227,10 +227,10 @@ async def get_session_user(
|
||||
async def update_profile(
|
||||
form_data: UpdateProfileForm,
|
||||
session_user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if session_user:
|
||||
user = Users.update_user_by_id(
|
||||
user = await Users.update_user_by_id(
|
||||
session_user.id,
|
||||
form_data.model_dump(),
|
||||
db=db,
|
||||
@@ -256,10 +256,10 @@ class UpdateTimezoneForm(BaseModel):
|
||||
async def update_timezone(
|
||||
form_data: UpdateTimezoneForm,
|
||||
session_user=Depends(get_current_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if session_user:
|
||||
Users.update_user_by_id(
|
||||
await Users.update_user_by_id(
|
||||
session_user.id,
|
||||
{'timezone': form_data.timezone},
|
||||
db=db,
|
||||
@@ -278,12 +278,12 @@ async def update_timezone(
|
||||
async def update_password(
|
||||
form_data: UpdatePasswordForm,
|
||||
session_user=Depends(get_current_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
|
||||
if session_user:
|
||||
user = Auths.authenticate_user(
|
||||
user = await Auths.authenticate_user(
|
||||
session_user.email,
|
||||
lambda pw: verify_password(form_data.password, pw),
|
||||
db=db,
|
||||
@@ -295,7 +295,7 @@ async def update_password(
|
||||
except Exception as e:
|
||||
raise HTTPException(400, detail=str(e))
|
||||
hashed = get_password_hash(form_data.new_password)
|
||||
return Auths.update_user_password_by_id(user.id, hashed, db=db)
|
||||
return await Auths.update_user_password_by_id(user.id, hashed, db=db)
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INCORRECT_PASSWORD)
|
||||
else:
|
||||
@@ -310,7 +310,7 @@ async def ldap_auth(
|
||||
request: Request,
|
||||
response: Response,
|
||||
form_data: LdapForm,
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
# Security checks FIRST - before loading any config
|
||||
if not request.app.state.config.ENABLE_LDAP:
|
||||
@@ -476,12 +476,12 @@ async def ldap_auth(
|
||||
if not await asyncio.to_thread(connection_user.bind):
|
||||
raise HTTPException(400, 'Authentication failed.')
|
||||
|
||||
user = Users.get_user_by_email(email, db=db)
|
||||
user = await Users.get_user_by_email(email, db=db)
|
||||
if not user:
|
||||
try:
|
||||
# Insert with default role first to avoid TOCTOU race on
|
||||
# first-user registration. Matches signup_handler pattern.
|
||||
user = Auths.insert_new_auth(
|
||||
user = await Auths.insert_new_auth(
|
||||
email=email,
|
||||
password=str(uuid.uuid4()),
|
||||
name=cn,
|
||||
@@ -494,11 +494,11 @@ async def ldap_auth(
|
||||
|
||||
# Atomically check if this is the only user *after* the
|
||||
# insert. Only the single user present should become admin.
|
||||
if Users.get_num_users(db=db) == 1:
|
||||
Users.update_user_role_by_id(user.id, 'admin', db=db)
|
||||
user = Users.get_user_by_id(user.id, db=db)
|
||||
if await Users.get_num_users(db=db) == 1:
|
||||
await Users.update_user_role_by_id(user.id, 'admin', db=db)
|
||||
user = await Users.get_user_by_id(user.id, db=db)
|
||||
|
||||
apply_default_group_assignment(
|
||||
await apply_default_group_assignment(
|
||||
request.app.state.config.DEFAULT_GROUP_ID,
|
||||
user.id,
|
||||
db=db,
|
||||
@@ -510,19 +510,19 @@ async def ldap_auth(
|
||||
log.error(f'LDAP user creation error: {str(err)}')
|
||||
raise HTTPException(500, detail='Internal error occurred during LDAP user creation.')
|
||||
|
||||
user = Auths.authenticate_user_by_email(email, db=db)
|
||||
user = await Auths.authenticate_user_by_email(email, db=db)
|
||||
|
||||
if user:
|
||||
if ENABLE_LDAP_GROUP_MANAGEMENT and user_groups:
|
||||
if ENABLE_LDAP_GROUP_CREATION:
|
||||
Groups.create_groups_by_group_names(user.id, user_groups, db=db)
|
||||
await Groups.create_groups_by_group_names(user.id, user_groups, db=db)
|
||||
try:
|
||||
Groups.sync_groups_by_group_names(user.id, user_groups, db=db)
|
||||
await Groups.sync_groups_by_group_names(user.id, user_groups, db=db)
|
||||
log.info(f'Successfully synced groups for user {user.id}: {user_groups}')
|
||||
except Exception as e:
|
||||
log.error(f'Failed to sync groups for user {user.id}: {e}')
|
||||
|
||||
return create_session_response(request, user, db, response, set_cookie=True)
|
||||
return await create_session_response(request, user, db, response, set_cookie=True)
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
else:
|
||||
@@ -542,7 +542,7 @@ async def signin(
|
||||
request: Request,
|
||||
response: Response,
|
||||
form_data: SigninForm,
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if not ENABLE_PASSWORD_AUTH:
|
||||
raise HTTPException(
|
||||
@@ -564,7 +564,7 @@ async def signin(
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
if not Users.get_user_by_email(email.lower(), db=db):
|
||||
if not await Users.get_user_by_email(email.lower(), db=db):
|
||||
await signup_handler(
|
||||
request,
|
||||
email,
|
||||
@@ -573,20 +573,20 @@ async def signin(
|
||||
db=db,
|
||||
)
|
||||
|
||||
user = Auths.authenticate_user_by_email(email, db=db)
|
||||
user = await Auths.authenticate_user_by_email(email, db=db)
|
||||
if user:
|
||||
if WEBUI_AUTH_TRUSTED_GROUPS_HEADER:
|
||||
group_names = request.headers.get(WEBUI_AUTH_TRUSTED_GROUPS_HEADER, '').split(',')
|
||||
group_names = [name.strip() for name in group_names if name.strip()]
|
||||
|
||||
if group_names:
|
||||
Groups.sync_groups_by_group_names(user.id, group_names, db=db)
|
||||
await Groups.sync_groups_by_group_names(user.id, group_names, db=db)
|
||||
|
||||
if WEBUI_AUTH_TRUSTED_ROLE_HEADER:
|
||||
trusted_role = request.headers.get(WEBUI_AUTH_TRUSTED_ROLE_HEADER, '').lower().strip()
|
||||
if trusted_role in {'admin', 'user', 'pending'}:
|
||||
if user.role != trusted_role:
|
||||
Users.update_user_role_by_id(user.id, trusted_role, db=db)
|
||||
await Users.update_user_role_by_id(user.id, trusted_role, db=db)
|
||||
elif trusted_role:
|
||||
log.warning(f'Ignoring invalid trusted role header value: {trusted_role}')
|
||||
|
||||
@@ -594,14 +594,14 @@ async def signin(
|
||||
admin_email = 'admin@localhost'
|
||||
admin_password = 'admin'
|
||||
|
||||
if Users.get_user_by_email(admin_email.lower(), db=db):
|
||||
user = Auths.authenticate_user(
|
||||
if await Users.get_user_by_email(admin_email.lower(), db=db):
|
||||
user = await Auths.authenticate_user(
|
||||
admin_email.lower(),
|
||||
lambda pw: verify_password(admin_password, pw),
|
||||
db=db,
|
||||
)
|
||||
else:
|
||||
if Users.has_users(db=db):
|
||||
if await Users.has_users(db=db):
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)
|
||||
|
||||
await signup_handler(
|
||||
@@ -612,7 +612,7 @@ async def signin(
|
||||
db=db,
|
||||
)
|
||||
|
||||
user = Auths.authenticate_user(
|
||||
user = await Auths.authenticate_user(
|
||||
admin_email.lower(),
|
||||
lambda pw: verify_password(admin_password, pw),
|
||||
db=db,
|
||||
@@ -633,14 +633,14 @@ async def signin(
|
||||
# decode safely — ignore incomplete UTF-8 sequences
|
||||
form_data.password = password_bytes.decode('utf-8', errors='ignore')
|
||||
|
||||
user = Auths.authenticate_user(
|
||||
user = await Auths.authenticate_user(
|
||||
form_data.email.lower(),
|
||||
lambda pw: verify_password(form_data.password, pw),
|
||||
db=db,
|
||||
)
|
||||
|
||||
if user:
|
||||
return create_session_response(request, user, db, response, set_cookie=True)
|
||||
return await create_session_response(request, user, db, response, set_cookie=True)
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
|
||||
@@ -657,7 +657,7 @@ async def signup_handler(
|
||||
name: str,
|
||||
profile_image_url: str = '/user.png',
|
||||
*,
|
||||
db: Session,
|
||||
db: AsyncSession,
|
||||
) -> UserModel:
|
||||
"""
|
||||
Core user-creation logic shared by the signup endpoint and
|
||||
@@ -671,7 +671,7 @@ async def signup_handler(
|
||||
# first-user registration can all see an empty table and each get admin.
|
||||
hashed = get_password_hash(password)
|
||||
|
||||
user = Auths.insert_new_auth(
|
||||
user = await Auths.insert_new_auth(
|
||||
email=email.lower(),
|
||||
password=hashed,
|
||||
name=name,
|
||||
@@ -684,9 +684,9 @@ async def signup_handler(
|
||||
|
||||
# Atomically check if this is the only user *after* the insert.
|
||||
# Only the single user present at this point should become admin.
|
||||
if Users.get_num_users(db=db) == 1:
|
||||
Users.update_user_role_by_id(user.id, 'admin', db=db)
|
||||
user = Users.get_user_by_id(user.id, db=db)
|
||||
if await Users.get_num_users(db=db) == 1:
|
||||
await Users.update_user_role_by_id(user.id, 'admin', db=db)
|
||||
user = await Users.get_user_by_id(user.id, db=db)
|
||||
request.app.state.config.ENABLE_SIGNUP = False
|
||||
|
||||
if request.app.state.config.WEBHOOK_URL:
|
||||
@@ -701,7 +701,7 @@ async def signup_handler(
|
||||
},
|
||||
)
|
||||
|
||||
apply_default_group_assignment(
|
||||
await apply_default_group_assignment(
|
||||
request.app.state.config.DEFAULT_GROUP_ID,
|
||||
user.id,
|
||||
db=db,
|
||||
@@ -715,9 +715,9 @@ async def signup(
|
||||
request: Request,
|
||||
response: Response,
|
||||
form_data: SignupForm,
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
has_users = Users.has_users(db=db)
|
||||
has_users = await Users.has_users(db=db)
|
||||
|
||||
if WEBUI_AUTH:
|
||||
if not request.app.state.config.ENABLE_SIGNUP or not request.app.state.config.ENABLE_LOGIN_FORM:
|
||||
@@ -730,7 +730,7 @@ async def signup(
|
||||
if not validate_email_format(form_data.email.lower()):
|
||||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT)
|
||||
|
||||
if Users.get_user_by_email(form_data.email.lower(), db=db):
|
||||
if await Users.get_user_by_email(form_data.email.lower(), db=db):
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
|
||||
|
||||
try:
|
||||
@@ -747,7 +747,7 @@ async def signup(
|
||||
form_data.profile_image_url,
|
||||
db=db,
|
||||
)
|
||||
return create_session_response(request, user, db, response, set_cookie=True)
|
||||
return await create_session_response(request, user, db, response, set_cookie=True)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as err:
|
||||
@@ -756,7 +756,7 @@ async def signup(
|
||||
|
||||
|
||||
@router.get('/signout')
|
||||
async def signout(request: Request, response: Response, db: Session = Depends(get_session)):
|
||||
async def signout(request: Request, response: Response, db: AsyncSession = Depends(get_async_session)):
|
||||
# get auth token from headers or cookies
|
||||
token = None
|
||||
auth_header = request.headers.get('Authorization')
|
||||
@@ -777,7 +777,7 @@ async def signout(request: Request, response: Response, db: Session = Depends(ge
|
||||
if oauth_session_id:
|
||||
response.delete_cookie('oauth_session_id')
|
||||
|
||||
session = OAuthSessions.get_session_by_id(oauth_session_id, db=db)
|
||||
session = await OAuthSessions.get_session_by_id(oauth_session_id, db=db)
|
||||
|
||||
# If a custom end_session_endpoint is configured (e.g. AWS Cognito), redirect
|
||||
# there directly instead of attempting OIDC discovery.
|
||||
@@ -852,12 +852,12 @@ async def add_user(
|
||||
request: Request,
|
||||
form_data: AddUserForm,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if not validate_email_format(form_data.email.lower()):
|
||||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT)
|
||||
|
||||
if Users.get_user_by_email(form_data.email.lower(), db=db):
|
||||
if await Users.get_user_by_email(form_data.email.lower(), db=db):
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
|
||||
|
||||
try:
|
||||
@@ -867,7 +867,7 @@ async def add_user(
|
||||
raise HTTPException(400, detail=str(e))
|
||||
|
||||
hashed = get_password_hash(form_data.password)
|
||||
user = Auths.insert_new_auth(
|
||||
user = await Auths.insert_new_auth(
|
||||
form_data.email.lower(),
|
||||
hashed,
|
||||
form_data.name,
|
||||
@@ -877,7 +877,7 @@ async def add_user(
|
||||
)
|
||||
|
||||
if user:
|
||||
apply_default_group_assignment(
|
||||
await apply_default_group_assignment(
|
||||
request.app.state.config.DEFAULT_GROUP_ID,
|
||||
user.id,
|
||||
db=db,
|
||||
@@ -909,7 +909,7 @@ async def add_user(
|
||||
|
||||
|
||||
@router.get('/admin/details')
|
||||
async def get_admin_details(request: Request, user=Depends(get_current_user), db: Session = Depends(get_session)):
|
||||
async def get_admin_details(request: Request, user=Depends(get_current_user), db: AsyncSession = Depends(get_async_session)):
|
||||
if request.app.state.config.SHOW_ADMIN_DETAILS:
|
||||
admin_email = request.app.state.config.ADMIN_EMAIL
|
||||
admin_name = None
|
||||
@@ -917,11 +917,11 @@ async def get_admin_details(request: Request, user=Depends(get_current_user), db
|
||||
log.info(f'Admin details - Email: {admin_email}, Name: {admin_name}')
|
||||
|
||||
if admin_email:
|
||||
admin = Users.get_user_by_email(admin_email, db=db)
|
||||
admin = await Users.get_user_by_email(admin_email, db=db)
|
||||
if admin:
|
||||
admin_name = admin.name
|
||||
else:
|
||||
admin = Users.get_first_user(db=db)
|
||||
admin = await Users.get_first_user(db=db)
|
||||
if admin:
|
||||
admin_email = admin.email
|
||||
admin_name = admin.name
|
||||
@@ -1173,10 +1173,10 @@ async def update_ldap_config(request: Request, form_data: LdapConfigForm, user=D
|
||||
|
||||
# create api key
|
||||
@router.post('/api_key', response_model=ApiKey)
|
||||
async def generate_api_key(request: Request, user=Depends(get_current_user), db: Session = Depends(get_session)):
|
||||
async def generate_api_key(request: Request, user=Depends(get_current_user), db: AsyncSession = Depends(get_async_session)):
|
||||
if not request.app.state.config.ENABLE_API_KEYS or (
|
||||
user.role != 'admin'
|
||||
and not has_permission(user.id, 'features.api_keys', request.app.state.config.USER_PERMISSIONS)
|
||||
and not await has_permission(user.id, 'features.api_keys', request.app.state.config.USER_PERMISSIONS)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
@@ -1184,7 +1184,7 @@ async def generate_api_key(request: Request, user=Depends(get_current_user), db:
|
||||
)
|
||||
|
||||
api_key = create_api_key()
|
||||
success = Users.update_user_api_key_by_id(user.id, api_key, db=db)
|
||||
success = await Users.update_user_api_key_by_id(user.id, api_key, db=db)
|
||||
|
||||
if success:
|
||||
return {
|
||||
@@ -1196,14 +1196,14 @@ async def generate_api_key(request: Request, user=Depends(get_current_user), db:
|
||||
|
||||
# delete api key
|
||||
@router.delete('/api_key', response_model=bool)
|
||||
async def delete_api_key(user=Depends(get_current_user), db: Session = Depends(get_session)):
|
||||
return Users.delete_user_api_key_by_id(user.id, db=db)
|
||||
async def delete_api_key(user=Depends(get_current_user), db: AsyncSession = Depends(get_async_session)):
|
||||
return await Users.delete_user_api_key_by_id(user.id, db=db)
|
||||
|
||||
|
||||
# get api key
|
||||
@router.get('/api_key', response_model=ApiKey)
|
||||
async def get_api_key(user=Depends(get_current_user), db: Session = Depends(get_session)):
|
||||
api_key = Users.get_user_api_key_by_id(user.id, db=db)
|
||||
async def get_api_key(user=Depends(get_current_user), db: AsyncSession = Depends(get_async_session)):
|
||||
api_key = await Users.get_user_api_key_by_id(user.id, db=db)
|
||||
if api_key:
|
||||
return {
|
||||
'api_key': api_key,
|
||||
@@ -1227,7 +1227,7 @@ async def token_exchange(
|
||||
response: Response,
|
||||
provider: str,
|
||||
form_data: TokenExchangeForm,
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
"""
|
||||
Exchange an external OAuth provider token for an OpenWebUI JWT.
|
||||
@@ -1296,14 +1296,14 @@ async def token_exchange(
|
||||
email = email.lower()
|
||||
|
||||
# Try to find the user by OAuth sub
|
||||
user = Users.get_user_by_oauth_sub(provider, sub, db=db)
|
||||
user = await Users.get_user_by_oauth_sub(provider, sub, db=db)
|
||||
|
||||
if not user and OAUTH_MERGE_ACCOUNTS_BY_EMAIL.value:
|
||||
# Try to find by email if merge is enabled
|
||||
user = Users.get_user_by_email(email, db=db)
|
||||
user = await Users.get_user_by_email(email, db=db)
|
||||
if user:
|
||||
# Link the OAuth sub to this user
|
||||
Users.update_user_oauth_by_id(user.id, provider, sub, db=db)
|
||||
await Users.update_user_oauth_by_id(user.id, provider, sub, db=db)
|
||||
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
@@ -1311,4 +1311,4 @@ async def token_exchange(
|
||||
detail='User not found. Please sign in via the web interface first.',
|
||||
)
|
||||
|
||||
return create_session_response(request, user, db)
|
||||
return await create_session_response(request, user, db)
|
||||
|
||||
Reference in New Issue
Block a user