From 3b4d6a2a3e3b38231b78fccd1af73b3d990b79ba Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sun, 7 Dec 2025 17:46:00 +0800 Subject: [PATCH] =?UTF-8?q?refactor(auth):=20=E5=B0=86=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E5=BA=93=E6=9F=A5=E8=AF=A2=E4=BB=8E=E5=90=8C=E6=AD=A5=E6=94=B9?= =?UTF-8?q?=E4=B8=BA=E5=BC=82=E6=AD=A5=E5=AE=9E=E7=8E=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将auth_router中的所有数据库查询操作从同步SQLAlchemy改为异步实现 添加async_check_first_run方法以支持异步检查首次运行 --- server/routers/auth_router.py | 108 ++++++++++++++++++++++++---------- src/storage/db/manager.py | 12 +++- 2 files changed, 86 insertions(+), 34 deletions(-) diff --git a/server/routers/auth_router.py b/server/routers/auth_router.py index 082ee2c7..96eba926 100644 --- a/server/routers/auth_router.py +++ b/server/routers/auth_router.py @@ -3,6 +3,8 @@ import re from fastapi import APIRouter, Depends, HTTPException, Request, status, UploadFile, File from fastapi.security import OAuth2PasswordRequestForm from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session from src.storage.db.manager import db_manager @@ -89,16 +91,22 @@ class UserIdGeneration(BaseModel): @auth.post("/token", response_model=Token) -async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)): +async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: AsyncSession = Depends(get_db)): # 查找用户 - 支持user_id和phone_number登录 login_identifier = form_data.username # OAuth2表单中的username字段作为登录标识符 # 尝试通过user_id查找 - user = db.query(User).filter(User.user_id == login_identifier).first() + result = await db.execute( + select(User).filter(User.user_id == login_identifier) + ) + user = result.scalar_one_or_none() # 如果通过user_id没找到，尝试通过phone_number查找 if not user: - user = db.query(User).filter(User.phone_number == login_identifier).first() + result = await db.execute( + select(User).filter(User.phone_number == login_identifier) + ) + user = result.scalar_one_or_none() # 如果用户不存在，为防止用户名枚举攻击，返回通用错误信息 if not user: @@ -176,15 +184,15 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends( # 路由：校验是否需要初始化管理员 @auth.get("/check-first-run") async def check_first_run(): - is_first_run = db_manager.check_first_run() + is_first_run = await db_manager.async_check_first_run() return {"first_run": is_first_run} # 路由：初始化管理员账户 @auth.post("/initialize", response_model=Token) -async def initialize_admin(admin_data: InitializeAdmin, db: Session = Depends(get_db)): +async def initialize_admin(admin_data: InitializeAdmin, db: AsyncSession = Depends(get_db)): # 检查是否是首次运行 - if not db_manager.check_first_run(): + if not await db_manager.async_check_first_run(): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="系统已经初始化，无法再次创建初始管理员", @@ -263,7 +271,7 @@ async def update_profile( profile_data: UserProfileUpdate, request: Request, current_user: User = Depends(get_required_user), - db: Session = Depends(get_db), + db: AsyncSession = Depends(get_db), ): """更新当前用户的个人资料""" update_details = [] @@ -279,9 +287,10 @@ async def update_profile( ) # 检查用户名是否已被其他用户使用 - existing_user = ( - db.query(User).filter(User.username == profile_data.username, User.id != current_user.id).first() + result = await db.execute( + select(User).filter(User.username == profile_data.username, User.id != current_user.id) ) + existing_user = result.scalar_one_or_none() if existing_user: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -299,11 +308,11 @@ async def update_profile( # 检查手机号是否已被其他用户使用 if profile_data.phone_number: - existing_phone = ( - db.query(User) + result = await db.execute( + select(User) .filter(User.phone_number == profile_data.phone_number, User.id != current_user.id) - .first() ) + existing_phone = result.scalar_one_or_none() if existing_phone: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="手机号已被其他用户使用") @@ -327,7 +336,7 @@ async def update_profile( @auth.post("/users", response_model=UserResponse) async def create_user( - user_data: UserCreate, request: Request, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db) + user_data: UserCreate, request: Request, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db) ): # 验证用户名 is_valid, error_msg = validate_username(user_data.username) @@ -338,7 +347,10 @@ async def create_user( ) # 检查用户名是否已存在 - existing_user = db.query(User).filter(User.username == user_data.username).first() + result = await db.execute( + select(User).filter(User.username == user_data.username) + ) + existing_user = result.scalar_one_or_none() if existing_user: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -347,7 +359,10 @@ async def create_user( # 检查手机号是否已存在（如果提供了） if user_data.phone_number: - existing_phone = db.query(User).filter(User.phone_number == user_data.phone_number).first() + result = await db.execute( + select(User).filter(User.phone_number == user_data.phone_number) + ) + existing_phone = result.scalar_one_or_none() if existing_phone: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -355,7 +370,8 @@ async def create_user( ) # 生成唯一的user_id - existing_user_ids = [user.user_id for user in db.query(User.user_id).all()] + result = await db.execute(select(User.user_id)) + existing_user_ids = [user_id for (user_id,) in result.all()] user_id = generate_unique_user_id(user_data.username, existing_user_ids) # 创建新用户 @@ -397,16 +413,25 @@ async def create_user( # 路由：获取所有用户（管理员权限） @auth.get("/users", response_model=list[UserResponse]) async def read_users( - skip: int = 0, limit: int = 100, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db) + skip: int = 0, limit: int = 100, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db) ): - users = db.query(User).filter(User.is_deleted == 0).offset(skip).limit(limit).all() + result = await db.execute( + select(User) + .filter(User.is_deleted == 0) + .offset(skip) + .limit(limit) + ) + users = result.scalars().all() return [user.to_dict() for user in users] # 路由：获取特定用户信息（管理员权限） @auth.get("/users/{user_id}", response_model=UserResponse) -async def read_user(user_id: int, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db)): - user = db.query(User).filter(User.id == user_id, User.is_deleted == 0).first() +async def read_user(user_id: int, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db)): + result = await db.execute( + select(User).filter(User.id == user_id, User.is_deleted == 0) + ) + user = result.scalar_one_or_none() if user is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -422,9 +447,12 @@ async def update_user( user_data: UserUpdate, request: Request, current_user: User = Depends(get_admin_user), - db: Session = Depends(get_db), + db: AsyncSession = Depends(get_db), ): - user = db.query(User).filter(User.id == user_id, User.is_deleted == 0).first() + result = await db.execute( + select(User).filter(User.id == user_id, User.is_deleted == 0) + ) + user = result.scalar_one_or_none() if user is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -450,7 +478,10 @@ async def update_user( if user_data.username is not None: # 检查用户名是否已被其他用户使用 - existing_user = db.query(User).filter(User.username == user_data.username, User.id != user_id).first() + result = await db.execute( + select(User).filter(User.username == user_data.username, User.id != user_id) + ) + existing_user = result.scalar_one_or_none() if existing_user: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -478,9 +509,12 @@ async def update_user( # 路由：删除用户（管理员权限） @auth.delete("/users/{user_id}", response_model=dict) async def delete_user( - user_id: int, request: Request, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db) + user_id: int, request: Request, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db) ): - user = db.query(User).filter(User.id == user_id, User.is_deleted == 0).first() + result = await db.execute( + select(User).filter(User.id == user_id, User.is_deleted == 0) + ) + user = result.scalar_one_or_none() if user is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -497,7 +531,10 @@ async def delete_user( ) # 检查是否是最后一个超级管理员 - superadmin_count = db.query(User).filter(User.role == "superadmin", User.is_deleted == 0).count() + result = await db.execute( + select(db.func.count(User.id)).filter(User.role == "superadmin", User.is_deleted == 0) + ) + superadmin_count = result.scalar() if superadmin_count <= 1: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -544,7 +581,7 @@ async def delete_user( # 路由：验证用户名并生成user_id @auth.post("/validate-username", response_model=UserIdGeneration) async def validate_username_and_generate_user_id( - validation_data: UsernameValidation, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db) + validation_data: UsernameValidation, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db) ): """验证用户名格式并生成可用的user_id""" # 验证用户名格式 @@ -556,7 +593,10 @@ async def validate_username_and_generate_user_id( ) # 检查用户名是否已存在 - existing_user = db.query(User).filter(User.username == validation_data.username).first() + result = await db.execute( + select(User).filter(User.username == validation_data.username) + ) + existing_user = result.scalar_one_or_none() if existing_user: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -564,7 +604,8 @@ async def validate_username_and_generate_user_id( ) # 生成唯一的user_id - existing_user_ids = [user.user_id for user in db.query(User.user_id).all()] + result = await db.execute(select(User.user_id)) + existing_user_ids = [user_id for (user_id,) in result.all()] user_id = generate_unique_user_id(validation_data.username, existing_user_ids) return UserIdGeneration(username=validation_data.username, user_id=user_id, is_available=True) @@ -573,17 +614,20 @@ async def validate_username_and_generate_user_id( # 路由：检查user_id是否可用 @auth.get("/check-user-id/{user_id}") async def check_user_id_availability( - user_id: str, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db) + user_id: str, current_user: User = Depends(get_admin_user), db: AsyncSession = Depends(get_db) ): """检查user_id是否可用""" - existing_user = db.query(User).filter(User.user_id == user_id).first() + result = await db.execute( + select(User).filter(User.user_id == user_id) + ) + existing_user = result.scalar_one_or_none() return {"user_id": user_id, "is_available": existing_user is None} # 路由：上传用户头像 @auth.post("/upload-avatar") async def upload_user_avatar( - file: UploadFile = File(...), current_user: User = Depends(get_required_user), db: Session = Depends(get_db) + file: UploadFile = File(...), current_user: User = Depends(get_required_user), db: AsyncSession = Depends(get_db) ): """上传用户头像""" # 检查文件类型 diff --git a/src/storage/db/manager.py b/src/storage/db/manager.py index a546b282..d1df8b8a 100644 --- a/src/storage/db/manager.py +++ b/src/storage/db/manager.py @@ -3,7 +3,7 @@ import os import pathlib from contextlib import asynccontextmanager, contextmanager -from sqlalchemy import create_engine +from sqlalchemy import create_engine, select, func from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import sessionmaker @@ -131,7 +131,7 @@ class DBManager(metaclass=SingletonMeta): await session.close() def check_first_run(self): - """检查是否首次运行""" + """检查是否首次运行（同步版本）""" session = self.get_session() try: # 检查是否有任何用户存在 @@ -139,6 +139,14 @@ class DBManager(metaclass=SingletonMeta): finally: session.close() + async def async_check_first_run(self): + """检查是否首次运行（异步版本）""" + async with self.get_async_session_context() as session: + # 检查是否有任何用户存在 + result = await session.execute(select(func.count(User.id))) + count = result.scalar() + return count == 0 + # 创建全局数据库管理器实例 db_manager = DBManager()