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