root
f5d7cfadb4
2、图片返回的 URL 是本地路径 3、视频生成中媒体文件重复下载 4、幂等性检查无数据库唯一约束 5、幂等键冲突返回 409 改为返回已有任务信息 6、虚拟素材库配额校验 TOCTOU 7、项目级联删除与独立素材删除任务并发冲突
184 lines
6.6 KiB
Python
184 lines
6.6 KiB
Python
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
from fastapi import HTTPException
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.enums.private_portrait import (
|
||
PRIVATE_PORTRAIT_REMOTE_PROJECT_NAME,
|
||
PRIVATE_PORTRAIT_VIRTUAL_GROUP_TYPE,
|
||
PrivatePortraitAssetType,
|
||
PrivatePortraitAssetStatus,
|
||
PrivatePortraitLibraryType,
|
||
PrivatePortraitProjectStatus,
|
||
)
|
||
from app.models.virtual_portrait_v3 import (
|
||
VpV3ApiKeyQuota,
|
||
VpV3Asset,
|
||
VpV3Project,
|
||
)
|
||
from app.utils.id_gen import generate_id
|
||
|
||
MB_BYTES = 1024 * 1024
|
||
_SAFE_SLUG = re.compile(r"[^a-zA-Z0-9_-]")
|
||
|
||
|
||
def _slug(name: str) -> str:
|
||
if not name:
|
||
return "unnamed"
|
||
return _SAFE_SLUG.sub("_", name.strip())[:80] or "unnamed"
|
||
|
||
|
||
def _bytes_to_mb(b: int | float | None) -> float:
|
||
if not b:
|
||
return 0.0
|
||
return round(b / MB_BYTES, 3)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 配额读写（确保 VpV3ApiKeyQuota 记录存在）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def _upsert_quota(db: AsyncSession, api_key_id: str) -> VpV3ApiKeyQuota:
|
||
"""获取配额记录；不存在则创建（默认全 0=不可用）。"""
|
||
from sqlalchemy.dialects.postgresql import insert
|
||
|
||
stmt = (
|
||
insert(VpV3ApiKeyQuota)
|
||
.values(
|
||
id=generate_id(),
|
||
api_key_id=api_key_id,
|
||
project_limit=0,
|
||
asset_limit=0,
|
||
storage_mb_limit=0,
|
||
project_used=0,
|
||
asset_used=0,
|
||
storage_mb_used=0,
|
||
)
|
||
.on_conflict_do_nothing(index_elements=["api_key_id"])
|
||
)
|
||
await db.execute(stmt)
|
||
row = (await db.execute(
|
||
select(VpV3ApiKeyQuota).where(VpV3ApiKeyQuota.api_key_id == api_key_id).limit(1)
|
||
)).scalar_one()
|
||
return row
|
||
|
||
|
||
async def _refresh_quota_used(db: AsyncSession, quota: VpV3ApiKeyQuota) -> None:
|
||
"""按真实数据重算已使用量（最终一致性）。"""
|
||
project_result = await db.execute(
|
||
select(func.count(VpV3Project.id)).where(
|
||
VpV3Project.api_key_id == quota.api_key_id,
|
||
VpV3Project.deleted_at.is_(None),
|
||
)
|
||
)
|
||
asset_result = await db.execute(
|
||
select(
|
||
func.count(VpV3Asset.id),
|
||
func.coalesce(func.sum(VpV3Asset.file_size_bytes), 0),
|
||
).where(
|
||
VpV3Asset.api_key_id == quota.api_key_id,
|
||
VpV3Asset.deleted_at.is_(None),
|
||
)
|
||
)
|
||
project_used = project_result.scalar() or 0
|
||
asset_row = asset_result.one()
|
||
asset_used = asset_row[0] or 0
|
||
storage_bytes = asset_row[1] or 0
|
||
quota.project_used = int(project_used)
|
||
quota.asset_used = int(asset_used)
|
||
quota.storage_mb_used = int(_bytes_to_mb(storage_bytes))
|
||
|
||
|
||
async def get_quota(db: AsyncSession, *, api_key_id: str, refresh: bool = True) -> VpV3ApiKeyQuota:
|
||
"""获取当前 API Key 的配额（含已使用量）。不存在则创建默认 0。"""
|
||
quota = await _upsert_quota(db, api_key_id)
|
||
if refresh:
|
||
await _refresh_quota_used(db, quota)
|
||
return quota
|
||
|
||
|
||
async def ensure_quota_enabled(db: AsyncSession, *, api_key_id: str) -> VpV3ApiKeyQuota:
|
||
"""校验是否已启用虚拟素材库功能，未启用直接 403。返回已刷新的配额。
|
||
|
||
判定口径（与后台设置保持一致）：只要「项目数上限」或「素材数上限」任一 > 0 即视为启用；
|
||
存储上限已从配置中移除（不再作为启用条件，也不做硬性限制，仅保留数据库字段做统计展示）。
|
||
"""
|
||
quota = await get_quota(db, api_key_id=api_key_id, refresh=True)
|
||
if (quota.project_limit or 0) <= 0 and (quota.asset_limit or 0) <= 0:
|
||
raise HTTPException(status_code=403, detail="当前 API Key 未开启虚拟素材库功能，请联系管理员配置配额")
|
||
return quota
|
||
|
||
|
||
def _check(limit: int | None, used: int | float | None, delta: int | float, field: str) -> None:
|
||
"""通用配额上限校验。
|
||
|
||
约定：limit <= 0 视为该维度「未配置 / 不做限制」，此时直接跳过不报错；
|
||
只有 limit > 0 时才按「已用 + 本次 <= 上限」判断，避免影响已移除的维度（如存储上限）。
|
||
"""
|
||
if (limit or 0) <= 0:
|
||
return # 不限制，直接通过
|
||
if (used or 0) + delta > limit:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=f"虚拟素材库配额不足：{field} 上限 {limit}，已使用 {used}，本次需要 {delta}，超出上限",
|
||
)
|
||
|
||
|
||
async def check_project_quota(db: AsyncSession, *, api_key_id: str, delta: int = 1) -> VpV3ApiKeyQuota:
|
||
"""创建项目前校验配额（带行锁，防止并发超配）。"""
|
||
quota = await ensure_quota_enabled(db, api_key_id=api_key_id)
|
||
# 用行锁重新读取，保证并发安全
|
||
quota = (
|
||
await db.execute(
|
||
select(VpV3ApiKeyQuota)
|
||
.where(VpV3ApiKeyQuota.api_key_id == api_key_id)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
).scalar_one()
|
||
await _refresh_quota_used(db, quota)
|
||
_check(quota.project_limit, quota.project_used, delta, "项目数")
|
||
return quota
|
||
|
||
|
||
async def check_asset_quota(
|
||
db: AsyncSession,
|
||
*,
|
||
api_key_id: str,
|
||
asset_count_delta: int = 1,
|
||
file_size_bytes: int | None = None,
|
||
) -> VpV3ApiKeyQuota:
|
||
"""上传素材前校验配额（带行锁，防止并发超配）。
|
||
|
||
注：「存储空间上限」已从业务约束中移除（不再做硬性配额限制），仅保留素材数量上限
|
||
与项目数量上限的校验；storage_mb_used 字段仍会在 get_quota 中刷新用于统计展示。
|
||
"""
|
||
del file_size_bytes # 不再用于配额校验（仅保留形参兼容现有调用点）
|
||
quota = await ensure_quota_enabled(db, api_key_id=api_key_id)
|
||
# 用行锁重新读取，保证并发安全
|
||
quota = (
|
||
await db.execute(
|
||
select(VpV3ApiKeyQuota)
|
||
.where(VpV3ApiKeyQuota.api_key_id == api_key_id)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
).scalar_one()
|
||
await _refresh_quota_used(db, quota)
|
||
_check(quota.asset_limit, quota.asset_used, asset_count_delta, "素材总数")
|
||
return quota
|
||
|
||
|
||
def remote_project_name() -> str:
|
||
"""火山 ProjectName（V3 中转统一共用这个 Project）。"""
|
||
return PRIVATE_PORTRAIT_REMOTE_PROJECT_NAME
|
||
|
||
|
||
def remote_group_name(*, api_key_id: str, project_slug: str, id: str) -> str:
|
||
"""火山 GroupName：vp-api-{api_key_id_short}-{id}-{slug} 最多 128 字符。"""
|
||
short_key = (api_key_id or "")
|
||
return f"vp-api-{short_key}-{id}-{project_slug}"[:128]
|