441 lines
15 KiB
Python
441 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from decimal import Decimal, ROUND_HALF_UP
|
|
from typing import Any, Iterable
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.resource_capacity import (
|
|
RESOURCE_CAPACITY_EXCEEDED_MESSAGE,
|
|
ResourceCapacityConfigKeyEnum,
|
|
ResourceCapacityErrorCodeEnum,
|
|
ResourceCapacitySourceEnum,
|
|
ResourceCapacityUnitEnum,
|
|
)
|
|
from app.models.system_config import SystemConfig
|
|
from app.models.user import User
|
|
from app.models.user_resource_capacity_config import UserResourceCapacityConfig
|
|
from app.models.user_resource_total_stat import UserResourceTotalStat
|
|
from app.schemas.resource_capacity import (
|
|
AdminUserResourceCapacityOut,
|
|
ResourceCapacityConfigOut,
|
|
ResourceCapacityConfigUpdate,
|
|
ResourceCapacityUsageOut,
|
|
)
|
|
from app.utils.id_gen import generate_id
|
|
|
|
DEFAULT_LIMIT_VALUE = Decimal("1.000")
|
|
DEFAULT_LIMIT_UNIT = ResourceCapacityUnitEnum.GB
|
|
DEFAULT_LIMIT_BYTES = DEFAULT_LIMIT_UNIT.bytes_multiplier
|
|
|
|
|
|
def _normalize_limit_value(value: Decimal | int | float | str | None) -> Decimal:
|
|
if value is None:
|
|
return DEFAULT_LIMIT_VALUE
|
|
decimal_value = Decimal(str(value))
|
|
return decimal_value.quantize(Decimal("0.001"), rounding=ROUND_HALF_UP)
|
|
|
|
|
|
def _limit_value_to_str(value: Decimal | int | float | str | None) -> str:
|
|
return format(_normalize_limit_value(value), "f")
|
|
|
|
|
|
def calculate_limit_bytes(
|
|
limit_value: Decimal | int | float | str | None,
|
|
limit_unit: ResourceCapacityUnitEnum | str | None,
|
|
) -> int:
|
|
unit = ResourceCapacityUnitEnum(limit_unit or DEFAULT_LIMIT_UNIT.value)
|
|
value = _normalize_limit_value(limit_value)
|
|
return int((value * Decimal(unit.bytes_multiplier)).to_integral_value(rounding=ROUND_HALF_UP))
|
|
|
|
|
|
def _default_config_out() -> ResourceCapacityConfigOut:
|
|
return ResourceCapacityConfigOut(
|
|
enabled=False,
|
|
limit_value=_limit_value_to_str(DEFAULT_LIMIT_VALUE),
|
|
limit_unit=DEFAULT_LIMIT_UNIT,
|
|
limit_bytes=DEFAULT_LIMIT_BYTES,
|
|
)
|
|
|
|
|
|
def _config_out(
|
|
*,
|
|
enabled: bool,
|
|
limit_value: Decimal | int | float | str | None,
|
|
limit_unit: ResourceCapacityUnitEnum | str | None,
|
|
limit_bytes: int | None = None,
|
|
) -> ResourceCapacityConfigOut:
|
|
unit = ResourceCapacityUnitEnum(limit_unit or DEFAULT_LIMIT_UNIT.value)
|
|
value = _normalize_limit_value(limit_value)
|
|
return ResourceCapacityConfigOut(
|
|
enabled=bool(enabled),
|
|
limit_value=_limit_value_to_str(value),
|
|
limit_unit=unit,
|
|
limit_bytes=int(limit_bytes if limit_bytes is not None else calculate_limit_bytes(value, unit)),
|
|
)
|
|
|
|
|
|
def _config_model_to_out(config: UserResourceCapacityConfig | None) -> ResourceCapacityConfigOut | None:
|
|
if config is None:
|
|
return None
|
|
return _config_out(
|
|
enabled=config.enabled,
|
|
limit_value=config.limit_value,
|
|
limit_unit=config.limit_unit,
|
|
limit_bytes=config.limit_bytes,
|
|
)
|
|
|
|
|
|
def _global_value_to_out(value: str | None) -> ResourceCapacityConfigOut:
|
|
if not value:
|
|
return _default_config_out()
|
|
try:
|
|
data = json.loads(value)
|
|
if not isinstance(data, dict):
|
|
return _default_config_out()
|
|
enabled = bool(data.get("enabled", False))
|
|
limit_unit = data.get("limit_unit") or DEFAULT_LIMIT_UNIT.value
|
|
limit_value = data.get("limit_value") or DEFAULT_LIMIT_VALUE
|
|
limit_bytes = data.get("limit_bytes")
|
|
return _config_out(
|
|
enabled=enabled,
|
|
limit_value=limit_value,
|
|
limit_unit=limit_unit,
|
|
limit_bytes=int(limit_bytes) if limit_bytes is not None else None,
|
|
)
|
|
except Exception:
|
|
return _default_config_out()
|
|
|
|
|
|
def _config_to_json(config: ResourceCapacityConfigOut) -> str:
|
|
return json.dumps(
|
|
{
|
|
"enabled": config.enabled,
|
|
"limit_value": config.limit_value,
|
|
"limit_unit": config.limit_unit.value,
|
|
"limit_bytes": config.limit_bytes,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
|
|
def _config_snapshot(config: ResourceCapacityConfigOut | None) -> dict[str, Any] | None:
|
|
if config is None:
|
|
return None
|
|
return {
|
|
"enabled": config.enabled,
|
|
"limit_value": config.limit_value,
|
|
"limit_unit": config.limit_unit.value,
|
|
"limit_bytes": config.limit_bytes,
|
|
}
|
|
|
|
|
|
def _usage_snapshot(usage: ResourceCapacityUsageOut | None) -> dict[str, Any] | None:
|
|
if usage is None:
|
|
return None
|
|
return {
|
|
"enabled": usage.enabled,
|
|
"source": usage.source.value,
|
|
"has_user_config": usage.has_user_config,
|
|
"used_bytes": usage.used_bytes,
|
|
"available_bytes": usage.available_bytes,
|
|
"total_bytes": usage.total_bytes,
|
|
"usage_percent": usage.usage_percent,
|
|
"exceeded": usage.exceeded,
|
|
"limit_value": usage.limit_value,
|
|
"limit_unit": usage.limit_unit.value if usage.limit_unit else None,
|
|
}
|
|
|
|
|
|
def build_global_resource_capacity_operation_detail(
|
|
before: ResourceCapacityConfigOut | None,
|
|
after: ResourceCapacityConfigOut | None,
|
|
) -> str:
|
|
"""构造全局容量配置操作日志详情。"""
|
|
|
|
return json.dumps(
|
|
{
|
|
"target": "global_resource_capacity",
|
|
"before": _config_snapshot(before),
|
|
"after": _config_snapshot(after),
|
|
},
|
|
ensure_ascii=False,
|
|
default=str,
|
|
)
|
|
|
|
|
|
def build_user_resource_capacity_operation_detail(
|
|
*,
|
|
target_user_id: str,
|
|
operation: str,
|
|
before: AdminUserResourceCapacityOut | None,
|
|
after: AdminUserResourceCapacityOut | None,
|
|
remark: str | None = None,
|
|
) -> str:
|
|
"""构造用户个人容量配置操作日志详情。"""
|
|
|
|
payload: dict[str, Any] = {
|
|
"target": "user_resource_capacity",
|
|
"target_user_id": target_user_id,
|
|
"operation": operation,
|
|
"before": {
|
|
"has_user_config": before.has_user_config if before else False,
|
|
"user_config": _config_snapshot(before.user_config) if before else None,
|
|
"effective": _usage_snapshot(before.effective) if before else None,
|
|
},
|
|
"after": {
|
|
"has_user_config": after.has_user_config if after else False,
|
|
"user_config": _config_snapshot(after.user_config) if after else None,
|
|
"effective": _usage_snapshot(after.effective) if after else None,
|
|
},
|
|
}
|
|
if remark:
|
|
payload["remark"] = remark
|
|
return json.dumps(payload, ensure_ascii=False, default=str)
|
|
|
|
|
|
def _build_usage_out(
|
|
*,
|
|
used_bytes: int,
|
|
has_user_config: bool,
|
|
source: ResourceCapacitySourceEnum,
|
|
config: ResourceCapacityConfigOut | None,
|
|
) -> ResourceCapacityUsageOut:
|
|
used = max(int(used_bytes or 0), 0)
|
|
if not config or not config.enabled:
|
|
return ResourceCapacityUsageOut(
|
|
enabled=False,
|
|
source=ResourceCapacitySourceEnum.DISABLED,
|
|
has_user_config=has_user_config,
|
|
used_bytes=used,
|
|
available_bytes=None,
|
|
total_bytes=None,
|
|
usage_percent=None,
|
|
exceeded=False,
|
|
limit_value=None,
|
|
limit_unit=None,
|
|
)
|
|
|
|
total = max(int(config.limit_bytes or 0), 0)
|
|
available = max(total - used, 0)
|
|
usage_percent = round((used / total) * 100, 2) if total > 0 else None
|
|
return ResourceCapacityUsageOut(
|
|
enabled=True,
|
|
source=source,
|
|
has_user_config=has_user_config,
|
|
used_bytes=used,
|
|
available_bytes=available,
|
|
total_bytes=total,
|
|
usage_percent=usage_percent,
|
|
exceeded=bool(total > 0 and used >= total),
|
|
limit_value=config.limit_value,
|
|
limit_unit=config.limit_unit,
|
|
)
|
|
|
|
|
|
async def get_global_resource_capacity_config(db: AsyncSession) -> ResourceCapacityConfigOut:
|
|
result = await db.execute(
|
|
select(SystemConfig.value)
|
|
.where(SystemConfig.key == ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value)
|
|
.limit(1)
|
|
)
|
|
return _global_value_to_out(result.scalar_one_or_none())
|
|
|
|
|
|
async def save_global_resource_capacity_config(
|
|
db: AsyncSession,
|
|
req: ResourceCapacityConfigUpdate,
|
|
) -> ResourceCapacityConfigOut:
|
|
limit_value = req.limit_value if req.limit_value is not None else DEFAULT_LIMIT_VALUE
|
|
limit_unit = req.limit_unit if req.limit_unit is not None else DEFAULT_LIMIT_UNIT
|
|
config_out = _config_out(
|
|
enabled=req.enabled,
|
|
limit_value=limit_value,
|
|
limit_unit=limit_unit,
|
|
)
|
|
result = await db.execute(
|
|
select(SystemConfig)
|
|
.where(SystemConfig.key == ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value)
|
|
.limit(1)
|
|
)
|
|
config = result.scalar_one_or_none()
|
|
if config:
|
|
config.value = _config_to_json(config_out)
|
|
config.description = "全局生成资源空间容量管控配置"
|
|
else:
|
|
db.add(
|
|
SystemConfig(
|
|
id=generate_id(),
|
|
key=ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value,
|
|
value=_config_to_json(config_out),
|
|
description="全局生成资源空间容量管控配置",
|
|
)
|
|
)
|
|
await db.flush()
|
|
return config_out
|
|
|
|
|
|
async def _get_user_config(db: AsyncSession, user_id: str) -> UserResourceCapacityConfig | None:
|
|
result = await db.execute(
|
|
select(UserResourceCapacityConfig)
|
|
.where(UserResourceCapacityConfig.user_id == user_id)
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def _get_used_bytes(db: AsyncSession, user_id: str) -> int:
|
|
result = await db.execute(
|
|
select(UserResourceTotalStat.active_size_bytes)
|
|
.where(UserResourceTotalStat.user_id == user_id)
|
|
.limit(1)
|
|
)
|
|
return int(result.scalar_one_or_none() or 0)
|
|
|
|
|
|
async def get_user_resource_capacity_usage(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
) -> ResourceCapacityUsageOut:
|
|
global_config = await get_global_resource_capacity_config(db)
|
|
user_config = await _get_user_config(db, user_id)
|
|
used_bytes = await _get_used_bytes(db, user_id)
|
|
|
|
user_config_out = _config_model_to_out(user_config)
|
|
if user_config_out is not None:
|
|
return _build_usage_out(
|
|
used_bytes=used_bytes,
|
|
has_user_config=True,
|
|
source=ResourceCapacitySourceEnum.USER if user_config_out.enabled else ResourceCapacitySourceEnum.DISABLED,
|
|
config=user_config_out,
|
|
)
|
|
|
|
return _build_usage_out(
|
|
used_bytes=used_bytes,
|
|
has_user_config=False,
|
|
source=ResourceCapacitySourceEnum.GLOBAL if global_config.enabled else ResourceCapacitySourceEnum.DISABLED,
|
|
config=global_config,
|
|
)
|
|
|
|
|
|
async def batch_get_user_resource_capacity_usage(
|
|
db: AsyncSession,
|
|
user_ids: Iterable[str],
|
|
) -> dict[str, ResourceCapacityUsageOut]:
|
|
ids = [user_id for user_id in dict.fromkeys(user_ids) if user_id]
|
|
if not ids:
|
|
return {}
|
|
|
|
global_config = await get_global_resource_capacity_config(db)
|
|
|
|
config_result = await db.execute(
|
|
select(UserResourceCapacityConfig)
|
|
.where(UserResourceCapacityConfig.user_id.in_(ids))
|
|
)
|
|
user_config_map = {item.user_id: item for item in config_result.scalars().all()}
|
|
|
|
stat_result = await db.execute(
|
|
select(UserResourceTotalStat.user_id, UserResourceTotalStat.active_size_bytes)
|
|
.where(UserResourceTotalStat.user_id.in_(ids))
|
|
)
|
|
used_map = {row.user_id: int(row.active_size_bytes or 0) for row in stat_result.all()}
|
|
|
|
usage_map: dict[str, ResourceCapacityUsageOut] = {}
|
|
for user_id in ids:
|
|
user_config_out = _config_model_to_out(user_config_map.get(user_id))
|
|
if user_config_out is not None:
|
|
usage_map[user_id] = _build_usage_out(
|
|
used_bytes=used_map.get(user_id, 0),
|
|
has_user_config=True,
|
|
source=ResourceCapacitySourceEnum.USER if user_config_out.enabled else ResourceCapacitySourceEnum.DISABLED,
|
|
config=user_config_out,
|
|
)
|
|
else:
|
|
usage_map[user_id] = _build_usage_out(
|
|
used_bytes=used_map.get(user_id, 0),
|
|
has_user_config=False,
|
|
source=ResourceCapacitySourceEnum.GLOBAL if global_config.enabled else ResourceCapacitySourceEnum.DISABLED,
|
|
config=global_config,
|
|
)
|
|
return usage_map
|
|
|
|
|
|
async def assert_user_resource_capacity_available(db: AsyncSession, user_id: str) -> None:
|
|
usage = await get_user_resource_capacity_usage(db, user_id)
|
|
if usage.enabled and usage.exceeded:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=RESOURCE_CAPACITY_EXCEEDED_MESSAGE,
|
|
headers={"X-Error-Code": ResourceCapacityErrorCodeEnum.RESOURCE_CAPACITY_EXCEEDED.value},
|
|
)
|
|
|
|
|
|
async def ensure_user_exists(db: AsyncSession, user_id: str) -> User:
|
|
result = await db.execute(select(User).where(User.id == user_id).limit(1))
|
|
user = result.scalar_one_or_none()
|
|
if not user:
|
|
raise HTTPException(status_code=404, detail="用户不存在")
|
|
return user
|
|
|
|
|
|
async def get_admin_user_resource_capacity(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
) -> AdminUserResourceCapacityOut:
|
|
await ensure_user_exists(db, user_id)
|
|
global_config = await get_global_resource_capacity_config(db)
|
|
user_config = await _get_user_config(db, user_id)
|
|
effective = await get_user_resource_capacity_usage(db, user_id)
|
|
return AdminUserResourceCapacityOut(
|
|
has_user_config=user_config is not None,
|
|
user_config=_config_model_to_out(user_config),
|
|
global_config=global_config,
|
|
effective=effective,
|
|
)
|
|
|
|
|
|
async def save_user_resource_capacity_config(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
req: ResourceCapacityConfigUpdate,
|
|
) -> AdminUserResourceCapacityOut:
|
|
await ensure_user_exists(db, user_id)
|
|
limit_value = req.limit_value if req.limit_value is not None else DEFAULT_LIMIT_VALUE
|
|
limit_unit = req.limit_unit if req.limit_unit is not None else DEFAULT_LIMIT_UNIT
|
|
limit_bytes = calculate_limit_bytes(limit_value, limit_unit)
|
|
|
|
config = await _get_user_config(db, user_id)
|
|
if config:
|
|
config.enabled = req.enabled
|
|
config.limit_value = _normalize_limit_value(limit_value)
|
|
config.limit_unit = ResourceCapacityUnitEnum(limit_unit).value
|
|
config.limit_bytes = limit_bytes
|
|
else:
|
|
db.add(
|
|
UserResourceCapacityConfig(
|
|
id=generate_id(),
|
|
user_id=user_id,
|
|
enabled=req.enabled,
|
|
limit_value=_normalize_limit_value(limit_value),
|
|
limit_unit=ResourceCapacityUnitEnum(limit_unit).value,
|
|
limit_bytes=limit_bytes,
|
|
)
|
|
)
|
|
await db.flush()
|
|
return await get_admin_user_resource_capacity(db, user_id)
|
|
|
|
|
|
async def delete_user_resource_capacity_config(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
) -> AdminUserResourceCapacityOut:
|
|
await ensure_user_exists(db, user_id)
|
|
await db.execute(
|
|
delete(UserResourceCapacityConfig).where(UserResourceCapacityConfig.user_id == user_id)
|
|
)
|
|
await db.flush()
|
|
return await get_admin_user_resource_capacity(db, user_id)
|