"""Redis-based rate limiter для проверки лимитов по тарифному плану.""" import json from datetime import datetime, timezone import redis from app.config import settings # Лимиты по тарифным планам PLAN_LIMITS: dict[str, dict[str, int | None]] = { "free": { "search_per_day": 10, "summarize_per_month": 3, "plagiarism_per_month": 1, "concurrent": 1, }, "student": { "search_per_day": None, # None = безлимит "summarize_per_month": 30, "plagiarism_per_month": 10, "concurrent": 2, }, "premium": { "search_per_day": None, "summarize_per_month": None, "plagiarism_per_month": 50, "concurrent": 5, }, "science": { "search_per_day": None, "summarize_per_month": None, "plagiarism_per_month": None, "concurrent": 10, }, } # Маппинг действий на ключи лимитов ACTION_TO_LIMIT: dict[str, tuple[str, str]] = { "search": ("search_per_day", "day"), "summarize": ("summarize_per_month", "month"), "plagiarism": ("plagiarism_per_month", "month"), "gost": ("gost_per_month", "month"), # ГОСТ всегда разрешён } def get_redis_client() -> redis.Redis: """Создать синхронный Redis клиент.""" return redis.from_url(settings.REDIS_URL, decode_responses=True) def _get_period_key(period: str) -> str: """Получить строку периода для Redis ключа.""" now = datetime.now(timezone.utc) if period == "day": return now.strftime("%Y-%m-%d") elif period == "month": return now.strftime("%Y-%m") return now.strftime("%Y-%m-%d") def check_and_increment_limit(user_id: int, action: str, plan: str) -> dict: """ Проверить лимит и инкрементировать счётчик. Args: user_id: ID пользователя action: Действие (search, plagiarism, summarize, gost) plan: Тарифный план пользователя Returns: dict с полями: - allowed: bool — разрешено ли действие - current: int — текущее количество использований - limit: int | None — лимит (None = безлимит) - remaining: int | None — осталось использований - reset_at: str — когда сбрасывается счётчик """ limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) if action not in ACTION_TO_LIMIT: # Неизвестное действие — разрешаем return {"allowed": True, "current": 0, "limit": None, "remaining": None} limit_key, period = ACTION_TO_LIMIT[action] limit_value = limits.get(limit_key) # Безлимитный план if limit_value is None: return { "allowed": True, "current": 0, "limit": None, "remaining": None, "reset_at": None, } period_str = _get_period_key(period) redis_key = f"rate:{user_id}:{action}:{period_str}" r = get_redis_client() # Атомарно инкрементировать pipe = r.pipeline() pipe.incr(redis_key) # Устанавливаем TTL: для дня — 86400 сек, для месяца — 32 дня ttl = 86400 if period == "day" else 86400 * 32 pipe.expire(redis_key, ttl) results = pipe.execute() current = results[0] if current > limit_value: # Декрементировать обратно (не считать запрещённые) r.decr(redis_key) current -= 1 return { "allowed": False, "current": current, "limit": limit_value, "remaining": 0, "reset_at": period_str, } return { "allowed": True, "current": current, "limit": limit_value, "remaining": limit_value - current, "reset_at": period_str, } def check_concurrent_limit(user_id: int, plan: str) -> bool: """ Проверить лимит одновременных задач. Returns: True если можно создать новую задачу, False если превышен лимит. """ limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) max_concurrent = limits.get("concurrent", 1) r = get_redis_client() key = f"concurrent:{user_id}" current = r.get(key) return (current is None) or (int(current) < max_concurrent) def increment_concurrent(user_id: int) -> None: """Увеличить счётчик одновременных задач (при создании задачи).""" r = get_redis_client() key = f"concurrent:{user_id}" pipe = r.pipeline() pipe.incr(key) pipe.expire(key, 3600) # Автосброс через 1 час pipe.execute() def decrement_concurrent(user_id: int) -> None: """Уменьшить счётчик одновременных задач (при завершении задачи).""" r = get_redis_client() key = f"concurrent:{user_id}" current = r.get(key) if current and int(current) > 0: r.decr(key)