Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac: async db

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/utils/oauth.py
+44 -44
View File
@@ -700,7 +700,7 @@ class OAuthClientManager:
"""
try:
# Get the OAuth session
session = OAuthSessions.get_session_by_provider_and_user_id(client_id, user_id)
session = await OAuthSessions.get_session_by_provider_and_user_id(client_id, user_id)
if not session:
log.warning(f'No OAuth session found for user {user_id}, client_id {client_id}')
return None
@@ -714,7 +714,7 @@ class OAuthClientManager:
log.warning(
f'Token refresh failed for user {user_id}, client_id {session.provider}, deleting session {session.id}'
)
OAuthSessions.delete_session_by_id(session.id)
await OAuthSessions.delete_session_by_id(session.id)
return None
return session.token
@@ -738,7 +738,7 @@ class OAuthClientManager:
if refreshed_token:
# Update the session with new token data
session = OAuthSessions.update_session_by_id(session.id, refreshed_token)
session = await OAuthSessions.update_session_by_id(session.id, refreshed_token)
log.info(f'Successfully refreshed token for session {session.id}')
return session.token
else:
@@ -884,12 +884,12 @@ class OAuthClientManager:
token['expires_at'] = datetime.now().timestamp() + token['expires_in']
# Clean up any existing sessions for this user/client_id first
sessions = OAuthSessions.get_sessions_by_user_id(user_id)
sessions = await OAuthSessions.get_sessions_by_user_id(user_id)
for session in sessions:
if session.provider == client_id:
OAuthSessions.delete_session_by_id(session.id)
await OAuthSessions.delete_session_by_id(session.id)
session = OAuthSessions.create_session(
session = await OAuthSessions.create_session(
user_id=user_id,
provider=client_id,
token=token,
@@ -963,7 +963,7 @@ class OAuthManager:
"""
try:
# Get the OAuth session
session = OAuthSessions.get_session_by_id_and_user_id(session_id, user_id)
session = await OAuthSessions.get_session_by_id_and_user_id(session_id, user_id)
if not session:
log.warning(f'No OAuth session found for user {user_id}, session {session_id}')
return None
@@ -977,7 +977,7 @@ class OAuthManager:
log.warning(
f'Token refresh failed for user {user_id}, provider {session.provider}, deleting session {session.id}'
)
OAuthSessions.delete_session_by_id(session.id)
await OAuthSessions.delete_session_by_id(session.id)
return None
return session.token
@@ -1002,7 +1002,7 @@ class OAuthManager:
if refreshed_token:
# Update the session with new token data
session = OAuthSessions.update_session_by_id(session.id, refreshed_token)
session = await OAuthSessions.update_session_by_id(session.id, refreshed_token)
log.info(f'Successfully refreshed token for session {session.id}')
return session.token
else:
@@ -1102,8 +1102,8 @@ class OAuthManager:
log.error(f'Exception during token refresh for provider {provider}: {e}')
return None
def get_user_role(self, user, user_data):
user_count = Users.get_num_users()
async def get_user_role(self, user, user_data):
user_count = await Users.get_num_users()
if user and user_count == 1:
# If the user is the only user, assign the role "admin" - actually repairs role for single user on login
log.debug('Assigning the only user the admin role')
@@ -1188,7 +1188,7 @@ class OAuthManager:
return role
def update_user_groups(self, user, user_data, default_permissions, db=None):
async def update_user_groups(self, user, user_data, default_permissions, db=None):
log.debug('Running OAUTH Group management')
oauth_claim = auth_manager_config.OAUTH_GROUPS_CLAIM
@@ -1217,8 +1217,8 @@ class OAuthManager:
else:
user_oauth_groups = []
user_current_groups: list[GroupModel] = Groups.get_groups_by_member_id(user.id, db=db)
all_available_groups: list[GroupModel] = Groups.get_all_groups(db=db)
user_current_groups: list[GroupModel] = await Groups.get_groups_by_member_id(user.id, db=db)
all_available_groups: list[GroupModel] = await Groups.get_all_groups(db=db)
# Create groups if they don't exist and creation is enabled
if auth_manager_config.ENABLE_OAUTH_GROUP_CREATION:
@@ -1226,7 +1226,7 @@ class OAuthManager:
all_group_names = {g.name for g in all_available_groups}
groups_created = False
# Determine creator ID: Prefer admin, fallback to current user if no admin exists
admin_user = Users.get_super_admin_user()
admin_user = await Users.get_super_admin_user()
creator_id = admin_user.id if admin_user else user.id
log.debug(f'Using creator ID {creator_id} for potential group creation.')
@@ -1241,7 +1241,7 @@ class OAuthManager:
data={'config': {'share': auth_manager_config.OAUTH_GROUP_DEFAULT_SHARE}},
)
# Use determined creator ID (admin or fallback to current user)
created_group = Groups.insert_new_group(creator_id, new_group_form, db=db)
created_group = await Groups.insert_new_group(creator_id, new_group_form, db=db)
if created_group:
log.info(
f"Successfully created group '{group_name}' with ID {created_group.id} using creator ID {creator_id}"
@@ -1256,7 +1256,7 @@ class OAuthManager:
# Refresh the list of all available groups if any were created
if groups_created:
all_available_groups = Groups.get_all_groups(db=db)
all_available_groups = await Groups.get_all_groups(db=db)
log.debug('Refreshed list of all available groups after creation.')
log.debug(f'Oauth Groups claim: {oauth_claim}')
@@ -1273,14 +1273,14 @@ class OAuthManager:
):
# Remove group from user
log.debug(f'Removing user from group {group_model.name} as it is no longer in their oauth groups')
Groups.remove_users_from_group(group_model.id, [user.id], db=db)
await Groups.remove_users_from_group(group_model.id, [user.id], db=db)
# In case a group is created, but perms are never assigned to the group by hitting "save"
group_permissions = group_model.permissions
if not group_permissions:
group_permissions = default_permissions
Groups.update_group_by_id(
await Groups.update_group_by_id(
id=group_model.id,
form_data=GroupUpdateForm(
name=group_model.name,
@@ -1302,14 +1302,14 @@ class OAuthManager:
# Add user to group
log.debug(f'Adding user to group {group_model.name} as it was found in their oauth groups')
Groups.add_users_to_group(group_model.id, [user.id], db=db)
await Groups.add_users_to_group(group_model.id, [user.id], db=db)
# In case a group is created, but perms are never assigned to the group by hitting "save"
group_permissions = group_model.permissions
if not group_permissions:
group_permissions = default_permissions
Groups.update_group_by_id(
await Groups.update_group_by_id(
id=group_model.id,
form_data=GroupUpdateForm(
name=group_model.name,
@@ -1490,20 +1490,20 @@ class OAuthManager:
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
# Check if the user exists
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:
# If the user does not exist, check if merging is enabled
if auth_manager_config.OAUTH_MERGE_ACCOUNTS_BY_EMAIL:
# Check if the user exists by email
user = Users.get_user_by_email(email, db=db)
user = await Users.get_user_by_email(email, db=db)
if user:
# Update the user with the new oauth sub
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 user:
determined_role = self.get_user_role(user, user_data)
determined_role = await self.get_user_role(user, user_data)
if user.role != determined_role:
Users.update_user_role_by_id(user.id, determined_role, db=db)
await Users.update_user_role_by_id(user.id, determined_role, db=db)
# Update the user object in memory as well,
# to avoid problems with the ENABLE_OAUTH_GROUP_MANAGEMENT check below
user.role = determined_role
@@ -1513,7 +1513,7 @@ class OAuthManager:
if username_claim:
new_name = user_data.get(username_claim)
if new_name and new_name != user.name:
Users.update_user_by_id(user.id, {'name': new_name}, db=db)
await Users.update_user_by_id(user.id, {'name': new_name}, db=db)
user.name = new_name
log.debug(f'Updated name for user {user.email}')
@@ -1522,13 +1522,13 @@ class OAuthManager:
if email_claim:
new_email = user_data.get(email_claim)
if new_email and new_email.lower() != user.email.lower():
existing_user = Users.get_user_by_email(new_email, db=db)
existing_user = await Users.get_user_by_email(new_email, db=db)
if existing_user:
log.error(
f'Cannot update email to {new_email} for user {user.id} because it is already taken.'
)
else:
Auths.update_email_by_id(user.id, new_email.lower(), db=db)
await Auths.update_email_by_id(user.id, new_email.lower(), db=db)
user.email = new_email.lower()
log.debug(f'Updated email for user {user.id}')
@@ -1544,13 +1544,13 @@ class OAuthManager:
new_picture_url, token.get('access_token')
)
if processed_picture_url != user.profile_image_url:
Users.update_user_profile_image_url_by_id(user.id, processed_picture_url, db=db)
await Users.update_user_profile_image_url_by_id(user.id, processed_picture_url, db=db)
log.debug(f'Updated profile picture for user {user.email}')
else:
# If the user does not exist, check if signups are enabled
if auth_manager_config.ENABLE_OAUTH_SIGNUP:
# Check if an existing user with the same email already exists
existing_user = Users.get_user_by_email(email, db=db)
existing_user = await Users.get_user_by_email(email, db=db)
if existing_user:
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
@@ -1570,12 +1570,12 @@ class OAuthManager:
log.warning('Username claim is missing, using email as name')
name = email
user = Auths.insert_new_auth(
user = await Auths.insert_new_auth(
email=email,
password=get_password_hash(str(uuid.uuid4())), # Random password, not used
name=name,
profile_image_url=picture_url,
role=self.get_user_role(None, user_data),
role=await self.get_user_role(None, user_data),
oauth=oauth_data,
db=db,
)
@@ -1586,9 +1586,9 @@ class OAuthManager:
# Atomically check if this is the only user *after* the
# insert to avoid TOCTOU race on first-user registration.
# Matches signup_handler pattern.
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)
if auth_manager_config.WEBHOOK_URL:
await post_webhook(
@@ -1602,7 +1602,7 @@ class OAuthManager:
},
)
apply_default_group_assignment(request.app.state.config.DEFAULT_GROUP_ID, user.id, db=db)
await apply_default_group_assignment(request.app.state.config.DEFAULT_GROUP_ID, user.id, db=db)
else:
raise HTTPException(
@@ -1615,7 +1615,7 @@ class OAuthManager:
expires_delta=parse_duration(auth_manager_config.JWT_EXPIRES_IN),
)
if auth_manager_config.ENABLE_OAUTH_GROUP_MANAGEMENT:
self.update_user_groups(
await self.update_user_groups(
user=user,
user_data=user_data,
default_permissions=request.app.state.config.USER_PERMISSIONS,
@@ -1675,7 +1675,7 @@ class OAuthManager:
# Enforce max concurrent sessions per user/provider to prevent
# unbounded growth while allowing multi-device usage
sessions = OAuthSessions.get_sessions_by_user_id(user.id, db=db)
sessions = await OAuthSessions.get_sessions_by_user_id(user.id, db=db)
provider_sessions = sorted(
[session for session in sessions if session.provider == provider],
key=lambda session: session.created_at,
@@ -1684,9 +1684,9 @@ class OAuthManager:
# Keep the newest sessions up to the limit, prune the rest
if len(provider_sessions) >= OAUTH_MAX_SESSIONS_PER_USER:
for old_session in provider_sessions[OAUTH_MAX_SESSIONS_PER_USER - 1 :]:
OAuthSessions.delete_session_by_id(old_session.id, db=db)
await OAuthSessions.delete_session_by_id(old_session.id, db=db)
session = OAuthSessions.create_session(
session = await OAuthSessions.create_session(
user_id=user.id,
provider=provider,
token=token,
@@ -1847,7 +1847,7 @@ class OAuthManager:
# 8. Identify users to log out
users_to_logout = []
if sub:
user = Users.get_user_by_oauth_sub(matched_provider, sub, db=db)
user = await Users.get_user_by_oauth_sub(matched_provider, sub, db=db)
if user:
users_to_logout.append(user)
@@ -1868,9 +1868,9 @@ class OAuthManager:
revoked_count = 0
for user in users_to_logout:
sessions = OAuthSessions.get_sessions_by_user_id(user.id, db=db)
sessions = await OAuthSessions.get_sessions_by_user_id(user.id, db=db)
for oauth_session in sessions:
OAuthSessions.delete_session_by_id(oauth_session.id, db=db)
await OAuthSessions.delete_session_by_id(oauth_session.id, db=db)
if redis:
revocation_key = f'{REDIS_KEY_PREFIX}:auth:user:{user.id}:revoked_at'