38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
"""User quota service for plan-based monthly limits."""
|
|
from __future__ import annotations
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.repositories.usage_repository import UsageRepository
|
|
|
|
PLAN_MONTHLY_LIMITS = {
|
|
"free": 10.0,
|
|
"basic": 50.0,
|
|
"pro": 200.0,
|
|
"enterprise": 1000.0,
|
|
}
|
|
|
|
|
|
class UserQuotaService:
|
|
"""Service for checking user quota based on subscription plan."""
|
|
|
|
def __init__(self, session: AsyncSession | None = None):
|
|
self._session = session
|
|
self._repository = UsageRepository(session) if session else None
|
|
|
|
def get_monthly_limit(self, user_plan: str) -> float:
|
|
"""Get monthly cost limit based on user plan."""
|
|
return PLAN_MONTHLY_LIMITS.get(user_plan, PLAN_MONTHLY_LIMITS["free"])
|
|
|
|
async def check_quota_with_plan(
|
|
self,
|
|
user_id: str,
|
|
user_plan: str,
|
|
) -> dict:
|
|
"""Check quota using the limit from user plan."""
|
|
if not self._repository:
|
|
raise RuntimeError("UserQuotaService not initialized with AsyncSession")
|
|
|
|
monthly_limit = self.get_monthly_limit(user_plan)
|
|
return await self._repository.check_quota(user_id, monthly_limit=monthly_limit)
|