75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
"""Budget storage and period helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .shared import Budget, Optional, and_, datetime, desc, select
|
|
|
|
|
|
class BudgetEnforcementStorageMixin:
|
|
"""Budget lookup and period-window management helpers."""
|
|
|
|
async def _get_active_budget(self, tenant_id: int, project_id: Optional[int]) -> Optional[Budget]:
|
|
"""优先读取项目预算，否则回退租户预算。"""
|
|
if project_id is not None:
|
|
result = await self.db.execute(
|
|
select(Budget)
|
|
.where(
|
|
and_(
|
|
Budget.tenant_id == tenant_id,
|
|
Budget.budget_type == "project",
|
|
Budget.project_id == project_id,
|
|
)
|
|
)
|
|
.order_by(desc(Budget.created_at))
|
|
)
|
|
project_budget = result.scalar_one_or_none()
|
|
if project_budget:
|
|
return project_budget
|
|
|
|
result = await self.db.execute(
|
|
select(Budget)
|
|
.where(
|
|
and_(
|
|
Budget.tenant_id == tenant_id,
|
|
Budget.budget_type == "tenant",
|
|
)
|
|
)
|
|
.order_by(desc(Budget.created_at))
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def _ensure_current_period(self, budget: Budget, now: datetime) -> None:
|
|
"""若预算周期已过，自动切换到新周期并清零使用量。"""
|
|
start = budget.current_period_start
|
|
end = budget.current_period_end
|
|
if start is None or end is None or now >= end:
|
|
next_start, next_end = self._compute_period_window(now, budget.period)
|
|
budget.current_period_start = next_start
|
|
budget.current_period_end = next_end
|
|
budget.current_usage_usd = 0.0
|
|
budget.status = "active"
|
|
await self.db.commit()
|
|
|
|
def _compute_period_window(self, now: datetime, period: str) -> tuple[datetime, datetime]:
|
|
period = (period or "monthly").lower()
|
|
if period == "yearly":
|
|
start = datetime(now.year, 1, 1)
|
|
end = datetime(now.year + 1, 1, 1)
|
|
return start, end
|
|
if period == "quarterly":
|
|
quarter = (now.month - 1) // 3
|
|
start_month = quarter * 3 + 1
|
|
start = datetime(now.year, start_month, 1)
|
|
if start_month == 10:
|
|
end = datetime(now.year + 1, 1, 1)
|
|
else:
|
|
end = datetime(now.year, start_month + 3, 1)
|
|
return start, end
|
|
|
|
start = datetime(now.year, now.month, 1)
|
|
if now.month == 12:
|
|
end = datetime(now.year + 1, 1, 1)
|
|
else:
|
|
end = datetime(now.year, now.month + 1, 1)
|
|
return start, end
|