52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.resource_capacity import RESOURCE_CAPACITY_EXCEEDED_MESSAGE, ResourceCapacityErrorCodeEnum
|
|
from app.models.user import User
|
|
from app.services.resource_capacity_service import get_user_resource_capacity_usage
|
|
from app.services.upload_resource.accounting_service import get_or_create_total_stat
|
|
|
|
|
|
def is_admin_user(user: User | object | None) -> bool:
|
|
if user is None:
|
|
return False
|
|
try:
|
|
if getattr(user, "user_type", None) == "admin":
|
|
return True
|
|
if bool(getattr(user, "is_admin", False)):
|
|
return True
|
|
except Exception:
|
|
return False
|
|
return False
|
|
|
|
|
|
async def assert_upload_capacity_available(
|
|
db: AsyncSession,
|
|
*,
|
|
user: User,
|
|
file_size_bytes: int,
|
|
) -> None:
|
|
"""上传容量拦截。
|
|
|
|
admin 用户不拦截，但统计仍会入账。普通用户锁定 total_stat 行后判断：
|
|
active_size_bytes + 本次上传大小 <= 容量上限。
|
|
"""
|
|
if is_admin_user(user):
|
|
return
|
|
|
|
total_stat = await get_or_create_total_stat(db, user.id, for_update=True)
|
|
usage = await get_user_resource_capacity_usage(db, user.id)
|
|
if not usage.enabled or usage.total_bytes is None:
|
|
return
|
|
|
|
used = int(total_stat.active_size_bytes or 0)
|
|
size = max(int(file_size_bytes or 0), 0)
|
|
if used + size > int(usage.total_bytes or 0):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"{RESOURCE_CAPACITY_EXCEEDED_MESSAGE}，本次上传 {size} 字节，当前已用 {used} 字节，总容量 {usage.total_bytes} 字节",
|
|
headers={"X-Error-Code": ResourceCapacityErrorCodeEnum.RESOURCE_CAPACITY_EXCEEDED.value},
|
|
)
|