263 lines
10 KiB
Python
263 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import date, datetime, timezone
|
|
from typing import Any
|
|
|
|
from sqlalchemy import func, select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.upload_resource import UploadResourceTypeEnum
|
|
from app.models.generated_resource import GeneratedResource
|
|
from app.models.upload_resource import UploadResource
|
|
from app.models.user_resource_month_stat import UserResourceMonthStat
|
|
from app.models.user_resource_total_stat import UserResourceTotalStat
|
|
from app.utils.id_gen import generate_id
|
|
|
|
|
|
def resource_month_from_datetime(value: datetime | None = None) -> date:
|
|
value = value or datetime.now(timezone.utc)
|
|
return date(value.year, value.month, 1)
|
|
|
|
|
|
def _int(value: Any) -> int:
|
|
return int(value or 0)
|
|
|
|
|
|
def _add_non_negative(obj: Any, field: str, delta: int) -> None:
|
|
current = _int(getattr(obj, field, 0))
|
|
setattr(obj, field, max(0, current + int(delta or 0)))
|
|
|
|
|
|
def _add_raw(obj: Any, field: str, delta: int) -> None:
|
|
current = _int(getattr(obj, field, 0))
|
|
setattr(obj, field, current + int(delta or 0))
|
|
|
|
|
|
async def get_or_create_month_stat(db: AsyncSession, user_id: str, stat_month: date) -> UserResourceMonthStat:
|
|
result = await db.execute(
|
|
select(UserResourceMonthStat).where(
|
|
UserResourceMonthStat.user_id == user_id,
|
|
UserResourceMonthStat.stat_month == stat_month,
|
|
).limit(1)
|
|
)
|
|
stat = result.scalar_one_or_none()
|
|
if stat:
|
|
return stat
|
|
stat = UserResourceMonthStat(id=generate_id(), user_id=user_id, stat_month=stat_month)
|
|
db.add(stat)
|
|
await db.flush()
|
|
return stat
|
|
|
|
|
|
async def get_or_create_total_stat(db: AsyncSession, user_id: str, *, for_update: bool = False) -> UserResourceTotalStat:
|
|
stmt = select(UserResourceTotalStat).where(UserResourceTotalStat.user_id == user_id).limit(1)
|
|
if for_update:
|
|
stmt = stmt.with_for_update()
|
|
result = await db.execute(stmt)
|
|
stat = result.scalar_one_or_none()
|
|
if stat:
|
|
return stat
|
|
stat = UserResourceTotalStat(id=generate_id(), user_id=user_id)
|
|
db.add(stat)
|
|
await db.flush()
|
|
if for_update:
|
|
result = await db.execute(
|
|
select(UserResourceTotalStat)
|
|
.where(UserResourceTotalStat.user_id == user_id)
|
|
.with_for_update()
|
|
.limit(1)
|
|
)
|
|
locked = result.scalar_one_or_none()
|
|
if locked:
|
|
return locked
|
|
return stat
|
|
|
|
|
|
async def apply_upload_resource_stat_delta(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
stat_month: date,
|
|
resource_type: str,
|
|
active_size_delta: int = 0,
|
|
active_count_delta: int = 0,
|
|
deleted_size_delta: int = 0,
|
|
deleted_count_delta: int = 0,
|
|
upload_size_delta: int = 0,
|
|
upload_count_delta: int = 0,
|
|
) -> None:
|
|
month_stat = await get_or_create_month_stat(db, user_id, stat_month)
|
|
total_stat = await get_or_create_total_stat(db, user_id)
|
|
now = datetime.now(timezone.utc)
|
|
|
|
for stat in (month_stat, total_stat):
|
|
_add_non_negative(stat, "active_size_bytes", active_size_delta)
|
|
_add_non_negative(stat, "active_count", active_count_delta)
|
|
_add_non_negative(stat, "deleted_size_bytes", deleted_size_delta)
|
|
_add_non_negative(stat, "deleted_count", deleted_count_delta)
|
|
_add_raw(stat, "upload_size_bytes", upload_size_delta)
|
|
_add_raw(stat, "upload_count", upload_count_delta)
|
|
|
|
if resource_type == UploadResourceTypeEnum.IMAGE.value:
|
|
_add_non_negative(stat, "image_size_bytes", active_size_delta)
|
|
_add_non_negative(stat, "image_count", active_count_delta)
|
|
elif resource_type == UploadResourceTypeEnum.VIDEO.value:
|
|
_add_non_negative(stat, "video_size_bytes", active_size_delta)
|
|
_add_non_negative(stat, "video_count", active_count_delta)
|
|
elif resource_type == UploadResourceTypeEnum.AUDIO.value:
|
|
_add_non_negative(stat, "audio_size_bytes", active_size_delta)
|
|
_add_non_negative(stat, "audio_count", active_count_delta)
|
|
elif resource_type == UploadResourceTypeEnum.SHOT_SEGMENT.value:
|
|
_add_non_negative(stat, "shot_segment_size_bytes", active_size_delta)
|
|
_add_non_negative(stat, "shot_segment_count", active_count_delta)
|
|
|
|
stat.last_recalculated_at = now
|
|
|
|
|
|
async def release_upload_resource_capacity(db: AsyncSession, resource: UploadResource, *, released_at: datetime | None = None) -> bool:
|
|
if resource.capacity_released_at is not None:
|
|
return False
|
|
released_at = released_at or datetime.now(timezone.utc)
|
|
stat_month = resource_month_from_datetime(resource.created_at or released_at)
|
|
size = _int(resource.file_size_bytes)
|
|
resource.capacity_released_at = released_at
|
|
await apply_upload_resource_stat_delta(
|
|
db,
|
|
user_id=resource.user_id,
|
|
stat_month=stat_month,
|
|
resource_type=resource.resource_type,
|
|
active_size_delta=-size,
|
|
active_count_delta=-1,
|
|
deleted_size_delta=size,
|
|
deleted_count_delta=1,
|
|
)
|
|
return True
|
|
|
|
|
|
async def rebuild_user_resource_stats(db: AsyncSession, *, user_ids: list[str] | None = None) -> dict[str, int]:
|
|
"""按数据库真实资源账本重算统计。
|
|
|
|
说明：这里只重置并重算 user_ids 范围内的统计。未传 user_ids 时重算所有在资源表中出现过的用户。
|
|
"""
|
|
if user_ids is None:
|
|
ids: set[str] = set()
|
|
for model in (GeneratedResource, UploadResource):
|
|
result = await db.execute(select(model.user_id).distinct())
|
|
ids.update(v for v in result.scalars().all() if v)
|
|
user_ids = sorted(ids)
|
|
else:
|
|
user_ids = sorted({v for v in user_ids if v})
|
|
if not user_ids:
|
|
return {"users": 0, "month_rows": 0, "total_rows": 0}
|
|
|
|
await db.execute(update(UserResourceTotalStat).where(UserResourceTotalStat.user_id.in_(user_ids)).values(
|
|
active_size_bytes=0,
|
|
deleted_size_bytes=0,
|
|
total_generated_size_bytes=0,
|
|
upload_size_bytes=0,
|
|
image_size_bytes=0,
|
|
video_size_bytes=0,
|
|
audio_size_bytes=0,
|
|
shot_segment_size_bytes=0,
|
|
active_count=0,
|
|
deleted_count=0,
|
|
upload_count=0,
|
|
image_count=0,
|
|
video_count=0,
|
|
audio_count=0,
|
|
shot_segment_count=0,
|
|
last_recalculated_at=datetime.now(timezone.utc),
|
|
))
|
|
await db.execute(update(UserResourceMonthStat).where(UserResourceMonthStat.user_id.in_(user_ids)).values(
|
|
active_size_bytes=0,
|
|
deleted_size_bytes=0,
|
|
total_generated_size_bytes=0,
|
|
upload_size_bytes=0,
|
|
image_size_bytes=0,
|
|
video_size_bytes=0,
|
|
audio_size_bytes=0,
|
|
shot_segment_size_bytes=0,
|
|
active_count=0,
|
|
deleted_count=0,
|
|
upload_count=0,
|
|
image_count=0,
|
|
video_count=0,
|
|
audio_count=0,
|
|
shot_segment_count=0,
|
|
last_recalculated_at=datetime.now(timezone.utc),
|
|
))
|
|
|
|
month_rows = 0
|
|
total_rows = 0
|
|
|
|
gen_rows = await db.execute(
|
|
select(
|
|
GeneratedResource.user_id,
|
|
GeneratedResource.resource_month,
|
|
GeneratedResource.resource_type,
|
|
GeneratedResource.deleted_at,
|
|
func.count(GeneratedResource.id),
|
|
func.coalesce(func.sum(GeneratedResource.file_size_bytes), 0),
|
|
).where(GeneratedResource.user_id.in_(user_ids)).group_by(
|
|
GeneratedResource.user_id,
|
|
GeneratedResource.resource_month,
|
|
GeneratedResource.resource_type,
|
|
GeneratedResource.deleted_at,
|
|
)
|
|
)
|
|
for user_id, month, rtype, deleted_at, count, size in gen_rows.all():
|
|
active = deleted_at is None
|
|
await apply_upload_resource_stat_delta(
|
|
db,
|
|
user_id=user_id,
|
|
stat_month=month,
|
|
resource_type=rtype,
|
|
active_size_delta=int(size or 0) if active else 0,
|
|
active_count_delta=int(count or 0) if active else 0,
|
|
deleted_size_delta=0 if active else int(size or 0),
|
|
deleted_count_delta=0 if active else int(count or 0),
|
|
)
|
|
# 生成资源字段单独累加
|
|
month_stat = await get_or_create_month_stat(db, user_id, month)
|
|
total_stat = await get_or_create_total_stat(db, user_id)
|
|
if active:
|
|
_add_raw(month_stat, "total_generated_size_bytes", int(size or 0))
|
|
_add_raw(total_stat, "total_generated_size_bytes", int(size or 0))
|
|
|
|
upload_rows = await db.execute(
|
|
select(
|
|
UploadResource.user_id,
|
|
func.date_trunc("month", UploadResource.created_at).label("month"),
|
|
UploadResource.resource_type,
|
|
UploadResource.deleted_at,
|
|
func.count(UploadResource.id),
|
|
func.coalesce(func.sum(UploadResource.file_size_bytes), 0),
|
|
).where(UploadResource.user_id.in_(user_ids)).group_by(
|
|
UploadResource.user_id,
|
|
"month",
|
|
UploadResource.resource_type,
|
|
UploadResource.deleted_at,
|
|
)
|
|
)
|
|
for user_id, month_dt, rtype, deleted_at, count, size in upload_rows.all():
|
|
month = resource_month_from_datetime(month_dt or datetime.now(timezone.utc))
|
|
active = deleted_at is None
|
|
await apply_upload_resource_stat_delta(
|
|
db,
|
|
user_id=user_id,
|
|
stat_month=month,
|
|
resource_type=rtype,
|
|
active_size_delta=int(size or 0) if active else 0,
|
|
active_count_delta=int(count or 0) if active else 0,
|
|
deleted_size_delta=0 if active else int(size or 0),
|
|
deleted_count_delta=0 if active else int(count or 0),
|
|
upload_size_delta=int(size or 0) if active else 0,
|
|
upload_count_delta=int(count or 0) if active else 0,
|
|
)
|
|
|
|
await db.flush()
|
|
total_rows = len(user_ids)
|
|
month_count = await db.execute(select(func.count(UserResourceMonthStat.id)).where(UserResourceMonthStat.user_id.in_(user_ids)))
|
|
month_rows = int(month_count.scalar_one() or 0)
|
|
return {"users": len(user_ids), "month_rows": month_rows, "total_rows": total_rows}
|