210 lines
7.5 KiB
Python
210 lines
7.5 KiB
Python
"""配额校验（批次 C1）：`check_and_count(tenant, kind)` 统一入口。
|
||
|
||
kind 取值与 `Plan.limits` 的键对应：
|
||
- users / products / bills_monthly → 数量型（达到即拒）
|
||
- ai_parse_order / ai_ask → 复用给 apps.ai.usage.quota_for()
|
||
- batch_managed / finance_ledger → 功能开关（False 即拒）
|
||
- storefront → 功能开关
|
||
|
||
调用方拿到 QuotaExceeded 后翻译成 403 + 升级引导（见 views / 业务拦截点）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from django.db.models import Count, Sum
|
||
|
||
from .models import Plan, Subscription, get_subscription
|
||
|
||
|
||
class QuotaExceeded(Exception):
|
||
"""超出套餐限制。payload 供前端弹"升级引导"（列差距 + 跳价格页）。"""
|
||
|
||
def __init__(self, kind: str, used, limit, plan_code: str, plan_name: str,
|
||
upgrade_hint: str = ""):
|
||
self.kind = kind
|
||
self.used = used
|
||
self.limit = limit
|
||
self.plan_code = plan_code
|
||
self.plan_name = plan_name
|
||
self.upgrade_hint = upgrade_hint or "升级套餐即可解除限制"
|
||
super().__init__(f"quota exceeded for {kind}: {used}/{limit} @ {plan_code}")
|
||
|
||
def as_dict(self) -> dict:
|
||
return {
|
||
"code": "quota_exceeded",
|
||
"kind": self.kind,
|
||
"detail": f"{self.kind} 已达 {self.used}/{self.limit} 上限（当前套餐：{self.plan_name}），{self.upgrade_hint}",
|
||
"used": float(self.used) if isinstance(self.used, (int, float)) else self.used,
|
||
"limit": float(self.limit) if isinstance(self.limit, (int, float)) else self.limit,
|
||
"plan_code": self.plan_code,
|
||
"plan_name": self.plan_name,
|
||
"upgrade_url": "/#/pricing",
|
||
}
|
||
|
||
|
||
def limits_for(tenant) -> dict:
|
||
return get_subscription(tenant).effective_limits()
|
||
|
||
|
||
def plan_of(tenant) -> Plan:
|
||
return get_subscription(tenant).plan
|
||
|
||
|
||
def _count_users(tenant) -> int:
|
||
"""该租户下的活跃用户数（经审计日志里出现过的操作者去重）。
|
||
|
||
注：当前用户表尚未挂 tenant 外键（阶段 1 只做了 Org/租户字段预留），
|
||
故用"在该租户留下过审计记录的用户数"作为保守口径；
|
||
接入成员表后替换此函数即可，调用方无需改动。
|
||
"""
|
||
from apps.core.models import AuditLog
|
||
|
||
return (
|
||
AuditLog.objects.filter(tenant=tenant, user__isnull=False)
|
||
.values("user_id")
|
||
.distinct()
|
||
.count()
|
||
)
|
||
|
||
|
||
def _count_products(tenant) -> int:
|
||
from apps.catalog.models import Product
|
||
|
||
return Product.objects.filter(tenant=tenant, is_deleted=False).count()
|
||
|
||
|
||
def _count_bills_monthly(tenant) -> int:
|
||
"""本月单据量 = 销售单 + 进货单（草稿与已过账都算，is_deleted 排除）。"""
|
||
from datetime import date
|
||
|
||
from apps.purchase.models import PurchaseBill
|
||
from apps.sales.models import SalesBill
|
||
|
||
today = date.today()
|
||
month_start = today.replace(day=1)
|
||
sales = SalesBill.objects.filter(
|
||
tenant=tenant, is_deleted=False, bill_date__gte=month_start
|
||
).count()
|
||
purchase = PurchaseBill.objects.filter(
|
||
tenant=tenant, is_deleted=False, bill_date__gte=month_start
|
||
).count()
|
||
return sales + purchase
|
||
|
||
|
||
COUNTERS = {
|
||
"users": _count_users,
|
||
"products": _count_products,
|
||
"bills_monthly": _count_bills_monthly,
|
||
}
|
||
|
||
FEATURE_FLAGS = {"batch_managed", "finance_ledger", "storefront"}
|
||
|
||
|
||
def usage_snapshot(tenant) -> dict:
|
||
"""当前用量 vs 限额（前端"套餐用量"页 + 升级引导弹窗的数据源）。"""
|
||
sub = get_subscription(tenant)
|
||
limits = sub.effective_limits()
|
||
snap = {
|
||
"plan": {
|
||
"code": sub.plan.code, "name": sub.plan.name,
|
||
"price_monthly": float(sub.plan.price_monthly),
|
||
},
|
||
"status": sub.status,
|
||
"trial_ends_at": sub.trial_ends_at.isoformat() if sub.trial_ends_at else None,
|
||
"trial_expired": sub.is_trial_expired,
|
||
"period_end": sub.period_end.isoformat() if sub.period_end else None,
|
||
"items": [],
|
||
}
|
||
for kind, fn in COUNTERS.items():
|
||
limit = int(limits.get(kind) or 0)
|
||
used = fn(tenant)
|
||
snap["items"].append({
|
||
"kind": kind,
|
||
"used": used,
|
||
"limit": limit,
|
||
"unlimited": limit == 0,
|
||
"percent": None if limit == 0 else round(used / limit * 100, 1),
|
||
})
|
||
# AI 次数（复用 ai.usage）
|
||
from apps.ai import usage as ai_usage
|
||
|
||
for kind in (ai_usage.KIND_PARSE_ORDER, ai_usage.KIND_ASK):
|
||
limit = int(limits.get(f"ai_{kind}") or 0)
|
||
used = ai_usage.monthly_count(tenant, kind)
|
||
snap["items"].append({
|
||
"kind": f"ai_{kind}",
|
||
"used": used,
|
||
"limit": limit,
|
||
"unlimited": limit == 0,
|
||
"percent": None if limit == 0 else round(used / limit * 100, 1),
|
||
})
|
||
snap["features"] = {k: bool(limits.get(k)) for k in FEATURE_FLAGS}
|
||
return snap
|
||
|
||
|
||
def check_and_count(tenant, kind: str, *, delta: int = 1) -> dict:
|
||
"""配额校验统一入口。
|
||
|
||
数量型：used + delta > limit → 拒（limit=0 视为不限）
|
||
功能开关：limits[kind] 为假 → 拒
|
||
AI 次数：交给 apps.ai.usage.check_quota（保持单一实现）
|
||
|
||
返回 {used, limit, remaining} 或抛 QuotaExceeded。
|
||
"""
|
||
sub = get_subscription(tenant)
|
||
limits = sub.effective_limits()
|
||
plan = sub.plan
|
||
|
||
if kind in COUNTERS:
|
||
limit = int(limits.get(kind) or 0)
|
||
used = COUNTERS[kind](tenant)
|
||
if limit > 0 and used + delta > limit:
|
||
raise QuotaExceeded(kind, used, limit, plan.code, plan.name,
|
||
upgrade_hint=_hint_for(kind))
|
||
return {
|
||
"kind": kind, "used": used, "limit": limit,
|
||
"remaining": (limit - used) if limit > 0 else None,
|
||
}
|
||
if kind in FEATURE_FLAGS:
|
||
if not limits.get(kind):
|
||
raise QuotaExceeded(kind, 0, 0, plan.code, plan.name,
|
||
upgrade_hint=f"{_label(kind)}需专业版")
|
||
return {"kind": kind, "enabled": True}
|
||
|
||
if kind.startswith("ai_"):
|
||
from apps.ai import usage as ai_usage
|
||
|
||
out = ai_usage.check_quota(tenant, kind[3:]) # ai_usage 内部已抛自己的 QuotaExceeded
|
||
return out
|
||
|
||
raise ValueError(f"unknown quota kind: {kind}")
|
||
|
||
|
||
def _label(kind: str) -> str:
|
||
return {
|
||
"users": "用户数", "products": "商品数", "bills_monthly": "月度单据量",
|
||
"ai_parse_order": "AI 录单次数", "ai_ask": "AI 问答次数",
|
||
"batch_managed": "批次效期管理", "finance_ledger": "总账核算",
|
||
"storefront": "订货商城",
|
||
}.get(kind, kind)
|
||
|
||
|
||
def _hint_for(kind: str) -> str:
|
||
return {
|
||
"users": "升级到基础版可加到 5 个用户，专业版 50 个",
|
||
"products": "升级到基础版可管 3000 个商品",
|
||
"bills_monthly": "升级到专业版不限制单据量",
|
||
"ai_parse_order": "升级套餐可获得更多 AI 录单次数",
|
||
"ai_ask": "升级套餐可获得更多 AI 问答次数",
|
||
}.get(kind, "升级套餐即可解除限制")
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# 供 apps.ai 反向调用（billing 未安装时 ai 会回落到 settings）
|
||
# ------------------------------------------------------------
|
||
|
||
def quota_for(tenant, kind: str) -> int:
|
||
"""返回某配额项的月度/总量上限（0 = 不限）。apps.ai.usage 优先调用本函数。"""
|
||
limits = limits_for(tenant)
|
||
return int(limits.get(kind) or 0)
|