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
+40 -38
View File
@@ -2,8 +2,9 @@ import logging
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.users import User, UserModel, UserProfileImageResponse, Users
from open_webui.utils.validate import validate_profile_image_url
from pydantic import BaseModel, field_validator
@@ -88,7 +89,7 @@ class AddUserForm(SignupForm):
class AuthsTable:
def insert_new_auth(
async def insert_new_auth(
self,
email: str,
password: str,
@@ -96,9 +97,9 @@ class AuthsTable:
profile_image_url: str = '/user.png',
role: str = 'pending',
oauth: Optional[dict] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[UserModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
log.info('insert_new_auth')
id = str(uuid.uuid4())
@@ -107,28 +108,29 @@ class AuthsTable:
result = Auth(**auth.model_dump())
db.add(result)
user = Users.insert_new_user(id, name, email, profile_image_url, role, oauth=oauth, db=db)
user = await Users.insert_new_user(id, name, email, profile_image_url, role, oauth=oauth, db=db)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result and user:
return user
else:
return None
def authenticate_user(
self, email: str, verify_password: callable, db: Optional[Session] = None
async def authenticate_user(
self, email: str, verify_password: callable, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
log.info(f'authenticate_user: {email}')
user = Users.get_user_by_email(email, db=db)
user = await Users.get_user_by_email(email, db=db)
if not user:
return None
try:
with get_db_context(db) as db:
auth = db.query(Auth).filter_by(id=user.id, active=True).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Auth).filter_by(id=user.id, active=True))
auth = result.scalars().first()
if auth:
if verify_password(auth.password):
return user
@@ -139,66 +141,66 @@ class AuthsTable:
except Exception:
return None
def authenticate_user_by_api_key(self, api_key: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def authenticate_user_by_api_key(self, api_key: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
log.info(f'authenticate_user_by_api_key')
# if no api_key, return None
if not api_key:
return None
try:
user = Users.get_user_by_api_key(api_key, db=db)
user = await Users.get_user_by_api_key(api_key, db=db)
return user if user else None
except Exception:
return False
def authenticate_user_by_email(self, email: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def authenticate_user_by_email(self, email: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
log.info(f'authenticate_user_by_email: {email}')
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Single JOIN query instead of two separate queries
result = (
db.query(Auth, User)
result = await db.execute(
select(Auth, User)
.join(User, Auth.id == User.id)
.filter(Auth.email == email, Auth.active == True)
.first()
)
if result:
_, user = result
row = result.first()
if row:
_, user = row
return UserModel.model_validate(user)
return None
except Exception:
return None
def update_user_password_by_id(self, id: str, new_password: str, db: Optional[Session] = None) -> bool:
async def update_user_password_by_id(self, id: str, new_password: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
result = db.query(Auth).filter_by(id=id).update({'password': new_password})
db.commit()
return True if result == 1 else False
async with get_async_db_context(db) as db:
result = await db.execute(update(Auth).filter_by(id=id).values(password=new_password))
await db.commit()
return True if result.rowcount == 1 else False
except Exception:
return False
def update_email_by_id(self, id: str, email: str, db: Optional[Session] = None) -> bool:
async def update_email_by_id(self, id: str, email: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
result = db.query(Auth).filter_by(id=id).update({'email': email})
db.commit()
if result == 1:
Users.update_user_by_id(id, {'email': email}, db=db)
async with get_async_db_context(db) as db:
result = await db.execute(update(Auth).filter_by(id=id).values(email=email))
await db.commit()
if result.rowcount == 1:
await Users.update_user_by_id(id, {'email': email}, db=db)
return True
return False
except Exception:
return False
def delete_auth_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_auth_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete User
result = Users.delete_user_by_id(id, db=db)
result = await Users.delete_user_by_id(id, db=db)
if result:
db.query(Auth).filter_by(id=id).delete()
db.commit()
await db.execute(delete(Auth).filter_by(id=id))
await db.commit()
return True
else: