refactor(auth): 将数据库查询从同步改为异步实现
将auth_router中的所有数据库查询操作从同步SQLAlchemy改为异步实现 添加async_check_first_run方法以支持异步检查首次运行
This commit is contained in:
Wenjie Zhang
parent
704954ba96
commit
3b4d6a2a3e
@ -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)
|
||||
):
|
||||
"""上传用户头像"""
|
||||
# 检查文件类型
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user