xsl
12583e9739
Typed pydantic-settings config, structlog JSON logging with request/trace id propagation and redaction of phone numbers, tokens and face keys, the unified response envelope and error-code table from 架构文档 §9.2, async SQLAlchemy session management, ORM models for the eight foundation tables, and the OpenTelemetry/Prometheus wiring. packages/domain stays pure logic -- snowflake ids, the quota three-phase protocol and the tk_task state machine -- so it can be tested without a database and is held to mypy strict.
141 lines
4.6 KiB
Python
Executable File
141 lines
4.6 KiB
Python
Executable File
"""Quota arithmetic — pure logic, no IO (架构文档 §7.4 V1 配额三段式).
|
|
|
|
The three-phase protocol lives here as data transformations so it can be unit
|
|
tested without a database:
|
|
|
|
reserve → 冻结额度并写 bl_ledger(status=reserved)
|
|
commit → 确认扣除 (status=committed)
|
|
refund → 释放额度并写反向流水 (tx_type=refund)
|
|
|
|
V1 doesn't involve money, but 用户对「次数被白扣」的敏感度不比扣钱低，and V1
|
|
has no customer-service refund path — the mechanism is the only guarantee.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
from datetime import UTC, datetime, timedelta, timezone
|
|
from typing import Literal
|
|
|
|
from packages.domain.errors import QuotaExceededError
|
|
|
|
PeriodType = Literal["day", "month", "lifetime"]
|
|
LedgerStatus = Literal["reserved", "committed", "refunded"]
|
|
TxType = Literal["system_grant", "consume", "refund", "purchase", "gift"]
|
|
|
|
#: FR-BILL-01: 跨时区按 UTC+8 统一。
|
|
DEFAULT_TZ_OFFSET_HOURS = 8
|
|
|
|
|
|
def period_key(
|
|
now: datetime,
|
|
period_type: PeriodType,
|
|
tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS,
|
|
) -> str:
|
|
"""Bucket key that makes quota reset naturally, with no cron job.
|
|
|
|
``day`` → ``d:2026-08-13``; ``month`` → ``m:2026-08``; ``lifetime`` → ``life``.
|
|
"""
|
|
if period_type == "lifetime":
|
|
return "life"
|
|
local = now.astimezone(UTC) + timedelta(hours=tz_offset_hours)
|
|
if period_type == "day":
|
|
return f"d:{local:%Y-%m-%d}"
|
|
if period_type == "month":
|
|
return f"m:{local:%Y-%m}"
|
|
msg = f"unknown period_type: {period_type}"
|
|
raise ValueError(msg)
|
|
|
|
|
|
def period_reset_at(
|
|
now: datetime,
|
|
period_type: PeriodType,
|
|
tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS,
|
|
) -> datetime | None:
|
|
"""When the current period rolls over (used for the `reset_at` hint)."""
|
|
if period_type == "lifetime":
|
|
return None
|
|
tz = timezone(timedelta(hours=tz_offset_hours))
|
|
local = now.astimezone(tz)
|
|
if period_type == "day":
|
|
nxt = (local + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
|
|
else:
|
|
year, month = (local.year + 1, 1) if local.month == 12 else (local.year, local.month + 1)
|
|
nxt = local.replace(
|
|
year=year, month=month, day=1, hour=0, minute=0, second=0, microsecond=0
|
|
)
|
|
return nxt.astimezone(UTC)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class QuotaState:
|
|
"""Current period counters for one (user, quota_key)."""
|
|
|
|
used: int
|
|
base_limit: int
|
|
bonus: int = 0
|
|
|
|
@property
|
|
def effective_limit(self) -> int:
|
|
return self.base_limit + self.bonus
|
|
|
|
@property
|
|
def remaining(self) -> int:
|
|
return max(self.effective_limit - self.used, 0)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class QuotaDecision:
|
|
allowed: bool
|
|
required: int
|
|
remaining: int
|
|
effective_limit: int
|
|
|
|
|
|
def evaluate(state: QuotaState, amount: int) -> QuotaDecision:
|
|
"""Server-side hard check (FR-BILL-06). Client hints are UX only."""
|
|
if amount <= 0:
|
|
msg = "amount must be positive"
|
|
raise ValueError(msg)
|
|
return QuotaDecision(
|
|
allowed=state.used + amount <= state.effective_limit,
|
|
required=amount,
|
|
remaining=state.remaining,
|
|
effective_limit=state.effective_limit,
|
|
)
|
|
|
|
|
|
def reserve(state: QuotaState, amount: int) -> QuotaState:
|
|
"""Freeze ``amount`` units; raises :class:`QuotaExceededError` when over the limit."""
|
|
decision = evaluate(state, amount)
|
|
if not decision.allowed:
|
|
raise QuotaExceededError(required=amount, remaining=decision.remaining)
|
|
return replace(state, used=state.used + amount)
|
|
|
|
|
|
def refund(state: QuotaState, amount: int) -> QuotaState:
|
|
"""Release a previously reserved amount. Never drives ``used`` below zero."""
|
|
if amount <= 0:
|
|
msg = "amount must be positive"
|
|
raise ValueError(msg)
|
|
return replace(state, used=max(state.used - amount, 0))
|
|
|
|
|
|
def balance_after(state: QuotaState) -> int:
|
|
"""Value written to ``bl_ledger.balance_after``.
|
|
|
|
V1 has no wallet balance, so the ledger records *remaining quota* — that is
|
|
the number the user actually sees ("今日剩余 N 次"). When V2 introduces
|
|
``bl_wallet`` this becomes the wallet balance without a schema change.
|
|
"""
|
|
return state.remaining
|
|
|
|
|
|
def reserve_idempotency_key(*, user_id: int, quota_key: str, request_key: str) -> str:
|
|
return f"rsv:{user_id}:{quota_key}:{request_key}"
|
|
|
|
|
|
def refund_idempotency_key(reserve_key: str) -> str:
|
|
"""FR-BILL-05: 退还必须幂等。Derived, so a retry cannot double-refund."""
|
|
return f"rfd:{reserve_key}"
|