task_management_api/app/utils/rate_limiter.py - master - tinypot/integr

task_management_api/app/utils/rate_limiter.py
105 строк4 KB
Tim Polus
lab2
16 фев 2026, 19:22
16 фев 2026, 19:22ec900ad
"""
rate limiting
"""

import time
from typing import Optional
from fastapi import Request, HTTPException, status
from fastapi.responses import JSONResponse

from app.core.config import settings
from app.utils.redis_client import get_redis_client


class RateLimiter:

    def __init__(self, requests_per_minute: int = None):
        self.requests_per_minute = requests_per_minute or settings.RATE_LIMIT_REQUESTS_PER_MINUTE

    async def is_rate_limited(self, identifier: str) -> tuple[bool, Optional[int]]:
        try:
            redis_client = await get_redis_client()

            # Create key for this identifier (e.g., IP address or user ID)
            key = f"rate_limit:{identifier}"

            # Get current minute timestamp
            current_minute = int(time.time() // 60)

            # Clean up old entries and count requests in current minute
            # Use Redis pipeline for atomic operations
            async with redis_client.pipeline() as pipe:
                # Remove entries older than current minute
                await pipe.zremrangebyscore(key, 0, current_minute - 1)
                # Count remaining entries
                await pipe.zcard(key)
                # Add current request
                await pipe.zadd(key, {str(time.time()): current_minute})
                # Set expiration for the key (keep data for 2 minutes)
                await pipe.expire(key, 120)

                results = await pipe.execute()

            request_count = results[1]

            if request_count >= self.requests_per_minute:
                # Calculate retry after (when next minute starts)
                next_minute = (current_minute + 1) * 60
                retry_after = int(next_minute - time.time())
                return True, retry_after

            return False, None
        except Exception:
            # If Redis is not available, allow all requests (fail open)
            return False, None


async def rate_limit_middleware(request: Request, call_next):
    # Get client identifier (IP address)
    client_ip = request.client.host

    # Skip rate limiting for health checks
    if request.url.path == "/health":
        return await call_next(request)

    rate_limiter = RateLimiter()
    is_limited, retry_after = await rate_limiter.is_rate_limited(client_ip)

    if is_limited:
        return JSONResponse(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            content={
                "detail": "Rate limit exceeded",
                "retry_after": retry_after
            },
            headers={"Retry-After": str(retry_after)}
        )

    # Add rate limit headers to response
    response = await call_next(request)

    # Calculate remaining requests (only if Redis is available)
    try:
        redis_client = await get_redis_client()
        key = f"rate_limit:{client_ip}"
        current_minute = int(time.time() // 60)

        async with redis_client.pipeline() as pipe:
            await pipe.zremrangebyscore(key, 0, current_minute - 1)
            await pipe.zcard(key)
            results = await pipe.execute()

        request_count = results[1]
        remaining = max(0, settings.RATE_LIMIT_REQUESTS_PER_MINUTE - request_count)

        response.headers["X-RateLimit-Limit"] = str(settings.RATE_LIMIT_REQUESTS_PER_MINUTE)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Reset"] = str((current_minute + 1) * 60)

        # Additional headers as per requirements
        response.headers["X-Limit-Remaining"] = str(remaining)
    except Exception:
        # If Redis is not available, skip rate limit headers
        pass

    return response