142 lines
5.5 KiB
Python
142 lines
5.5 KiB
Python
from __future__ import annotations
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.private_portrait import (
|
|
PRIVATE_PORTRAIT_DEFAULT_ASSET_LIMIT,
|
|
PRIVATE_PORTRAIT_ENABLED_ASSET_TYPES,
|
|
PrivatePortraitAssetStatus,
|
|
PrivatePortraitEventSource,
|
|
PrivatePortraitEventStatus,
|
|
PrivatePortraitEventType,
|
|
PrivatePortraitLibraryType,
|
|
)
|
|
from app.models.private_portrait import PrivatePortraitAsset
|
|
from app.models.user import User
|
|
from app.schemas.private_portrait import PrivatePortraitConfigOut
|
|
from app.services.operation_log_service import log_operation_event
|
|
|
|
DOMAIN = "private_portrait"
|
|
|
|
_COUNTING_STATUSES = {
|
|
PrivatePortraitAssetStatus.CREATING.value,
|
|
PrivatePortraitAssetStatus.PROCESSING.value,
|
|
PrivatePortraitAssetStatus.ACTIVE.value,
|
|
}
|
|
_COUNTING_LIBRARY_TYPES = {
|
|
PrivatePortraitLibraryType.REAL_PERSON.value,
|
|
PrivatePortraitLibraryType.AIGC_VIRTUAL.value,
|
|
}
|
|
|
|
|
|
async def get_user_or_404(db: AsyncSession, *, user_id: str, for_update: bool = False) -> User:
|
|
stmt = select(User).where(User.id == user_id).limit(1)
|
|
if for_update:
|
|
stmt = stmt.with_for_update()
|
|
user = (await db.execute(stmt)).scalar_one_or_none()
|
|
if not user:
|
|
raise HTTPException(status_code=404, detail="用户不存在")
|
|
return user
|
|
|
|
|
|
def get_user_asset_limit_value(user: User) -> int:
|
|
return int(getattr(user, "private_portrait_asset_limit", PRIVATE_PORTRAIT_DEFAULT_ASSET_LIMIT) or 0)
|
|
|
|
|
|
async def count_user_counting_assets(db: AsyncSession, *, user_id: str) -> int:
|
|
total = (
|
|
await db.execute(
|
|
select(func.count(PrivatePortraitAsset.id)).where(
|
|
PrivatePortraitAsset.user_id == user_id,
|
|
PrivatePortraitAsset.library_type.in_(_COUNTING_LIBRARY_TYPES),
|
|
PrivatePortraitAsset.asset_type.in_(PRIVATE_PORTRAIT_ENABLED_ASSET_TYPES),
|
|
PrivatePortraitAsset.deleted_at.is_(None),
|
|
PrivatePortraitAsset.status.in_(_COUNTING_STATUSES),
|
|
)
|
|
)
|
|
).scalar_one()
|
|
return int(total or 0)
|
|
|
|
|
|
async def get_user_private_portrait_config(db: AsyncSession, *, user_id: str) -> PrivatePortraitConfigOut:
|
|
user = await get_user_or_404(db, user_id=user_id)
|
|
limit = get_user_asset_limit_value(user)
|
|
used = await count_user_counting_assets(db, user_id=user_id)
|
|
remaining = max(0, limit - used) if limit > 0 else 0
|
|
return PrivatePortraitConfigOut(
|
|
enabled=limit > 0,
|
|
asset_limit=limit,
|
|
used_asset_count=used,
|
|
remaining_asset_count=remaining,
|
|
image_limit=limit,
|
|
used_image_count=used,
|
|
remaining_image_count=remaining,
|
|
)
|
|
|
|
|
|
async def set_user_private_portrait_limit(db: AsyncSession, *, user_id: str, limit: int) -> User:
|
|
user = await get_user_or_404(db, user_id=user_id)
|
|
user.private_portrait_asset_limit = max(0, int(limit))
|
|
await db.flush()
|
|
return user
|
|
|
|
|
|
async def ensure_private_portrait_asset_quota_available(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
project_id: str | None = None,
|
|
library_type: str | None = None,
|
|
asset_type: str | None = None,
|
|
) -> tuple[int, int]:
|
|
user = await get_user_or_404(db, user_id=user_id, for_update=True)
|
|
limit = get_user_asset_limit_value(user)
|
|
log_operation_event(
|
|
domain=DOMAIN,
|
|
event_type=PrivatePortraitEventType.QUOTA_CHECK_START.value,
|
|
event_status=PrivatePortraitEventStatus.PENDING.value,
|
|
source=PrivatePortraitEventSource.SERVICE.value,
|
|
user_id=user_id,
|
|
project_id=project_id,
|
|
detail={"asset_limit": limit, "library_type": library_type, "asset_type": asset_type},
|
|
)
|
|
if limit <= 0:
|
|
log_operation_event(
|
|
domain=DOMAIN,
|
|
event_type=PrivatePortraitEventType.QUOTA_CHECK_DENY.value,
|
|
event_status=PrivatePortraitEventStatus.FAILED.value,
|
|
source=PrivatePortraitEventSource.SERVICE.value,
|
|
user_id=user_id,
|
|
project_id=project_id,
|
|
detail={"asset_limit": limit, "library_type": library_type, "asset_type": asset_type, "reason": "disabled"},
|
|
message="用户私域人像素材库未启用",
|
|
)
|
|
raise HTTPException(status_code=403, detail="私域人像素材库未启用")
|
|
|
|
used = await count_user_counting_assets(db, user_id=user_id)
|
|
if used >= limit:
|
|
log_operation_event(
|
|
domain=DOMAIN,
|
|
event_type=PrivatePortraitEventType.QUOTA_CHECK_DENY.value,
|
|
event_status=PrivatePortraitEventStatus.FAILED.value,
|
|
source=PrivatePortraitEventSource.SERVICE.value,
|
|
user_id=user_id,
|
|
project_id=project_id,
|
|
detail={"asset_limit": limit, "used_asset_count": used, "library_type": library_type, "asset_type": asset_type, "reason": "max_limit"},
|
|
message="用户私域人像素材总量已达上限",
|
|
)
|
|
raise HTTPException(status_code=400, detail=f"你的私域人像素材库最多可上传 {limit} 个素材，请删除已有素材后再上传")
|
|
|
|
log_operation_event(
|
|
domain=DOMAIN,
|
|
event_type=PrivatePortraitEventType.QUOTA_CHECK_PASS.value,
|
|
event_status=PrivatePortraitEventStatus.SUCCESS.value,
|
|
source=PrivatePortraitEventSource.SERVICE.value,
|
|
user_id=user_id,
|
|
project_id=project_id,
|
|
detail={"asset_limit": limit, "used_asset_count": used, "remaining_asset_count": max(0, limit - used), "library_type": library_type, "asset_type": asset_type},
|
|
)
|
|
return limit, used
|