fastapi-traffic/examples/00_basic_usage.py

"""Basic usage examples for fastapi-traffic."""

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated, TypeAlias

from fastapi import Depends, FastAPI, Request
from fastapi.responses import JSONResponse

from fastapi_traffic import (
    Algorithm,
    RateLimiter,
    RateLimitExceeded,
    SQLiteBackend,
    rate_limit,
)
from fastapi_traffic.core.decorator import RateLimitDependency
from fastapi_traffic.core.limiter import set_limiter
from fastapi_traffic.core.models import RateLimitInfo

DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000

# Configure global rate limiter with SQLite backend for persistence
backend = SQLiteBackend("rate_limits.db")
limiter = RateLimiter(backend)
set_limiter(limiter)

basic_ratelimiter = RateLimitDependency(limit=20, window_size=60)
RateLimitDep: TypeAlias = Annotated[RateLimitInfo, Depends(basic_ratelimiter)]


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    """Manage application lifespan - startup and shutdown."""
    # Startup: Initialize the rate limiter
    await limiter.initialize()
    yield
    # Shutdown: Cleanup
    await limiter.close()


app = FastAPI(title="FastAPI Traffic Example", lifespan=lifespan)


# Exception handler for rate limit exceeded
@app.exception_handler(RateLimitExceeded)
async def rate_limit_handler(_: Request, exc: RateLimitExceeded) -> JSONResponse:
    """Handle rate limit exceeded exceptions."""
    headers = exc.limit_info.to_headers() if exc.limit_info else {}
    return JSONResponse(
        status_code=429,
        content={
            "error": "rate_limit_exceeded",
            "message": exc.message,
            "retry_after": exc.retry_after,
        },
        headers=headers,
    )


# Example 1: Basic decorator usage
@app.get("/api/basic")
@rate_limit(100, 60)  # 100 requests per minute
async def basic_endpoint(_: Request) -> dict[str, str]:
    """Basic rate-limited endpoint."""
    return {"message": "Hello, World!"}


# Example 2: Custom algorithm
@app.get("/api/token-bucket")
@rate_limit(
    limit=50,
    window_size=60,
    algorithm=Algorithm.TOKEN_BUCKET,
    burst_size=10,  # Allow bursts of up to 10 requests
)
async def token_bucket_endpoint(_: Request) -> dict[str, str]:
    """Endpoint using token bucket algorithm."""
    return {"message": "Token bucket rate limiting"}


# Example 3: Sliding window for precise rate limiting
@app.get("/api/sliding-window")
@rate_limit(
    limit=30,
    window_size=60,
    algorithm=Algorithm.SLIDING_WINDOW,
)
async def sliding_window_endpoint(_: Request) -> dict[str, str]:
    """Endpoint using sliding window algorithm."""
    return {"message": "Sliding window rate limiting"}


# Example 4: Custom key extractor (rate limit by API key)
def api_key_extractor(request: Request) -> str:
    """Extract API key from header for rate limiting."""
    api_key = request.headers.get("X-API-Key", "anonymous")
    return f"api_key:{api_key}"


@app.get("/api/by-api-key")
@rate_limit(
    limit=1000,
    window_size=3600,  # 1000 requests per hour
    key_extractor=api_key_extractor,
)
async def api_key_endpoint(_: Request) -> dict[str, str]:
    """Endpoint rate limited by API key."""
    return {"message": "Rate limited by API key"}


# Example 5: Using dependency injection
# Note: This dependency injection seems to be tripping pydantic, needs to be looked into.
"""@app.get("/api/dependency")
async def dependency_endpoint(
    _: Request,
    rate_info: RateLimitDep,
) -> dict[str, object]:
    '''Endpoint using rate limit as dependency.'''
    return {
        "message": "Rate limit info available",
        "rate_limit": rate_info,
    }
"""


# Example 6: Exempt certain requests
def is_admin(request: Request) -> bool:
    """Check if request is from admin."""
    return request.headers.get("X-Admin-Token") == "secret-admin-token"


@app.get("/api/admin-exempt")
@rate_limit(
    limit=10,
    window_size=60,
    exempt_when=is_admin,
)
async def admin_exempt_endpoint(_: Request) -> dict[str, str]:
    """Endpoint with admin exemption."""
    return {"message": "Admins are exempt from rate limiting"}


# Example 7: Different costs for different operations
@app.post("/api/expensive")
@rate_limit(
    limit=100,
    window_size=60,
    cost=10,  # This endpoint costs 10 tokens per request
)
async def expensive_endpoint(_: Request) -> dict[str, str]:
    """Expensive operation that costs more tokens."""
    return {"message": "Expensive operation completed"}


# Example 8: Global middleware rate limiting
# Uncomment to enable global rate limiting
# app.add_middleware(
#     RateLimitMiddleware,
#     limit=1000,
#     window_size=60,
#     exempt_paths={"/health", "/docs", "/openapi.json"},
# )


@app.get("/health")
async def health_check(_: Request) -> dict[str, str]:
    """Health check endpoint (typically exempt from rate limiting)."""
    return {"status": "healthy"}


if __name__ == "__main__":
    import argparse

    import uvicorn

    parser = argparse.ArgumentParser(
        description="Basic usage example for fastapi-traffic"
    )
    parser.add_argument(
        "--host",
        default=DEFAULT_HOST,
        help=f"Host to bind to (default: {DEFAULT_HOST})",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=DEFAULT_PORT,
        help=f"Port to bind to (default: {DEFAULT_PORT})",
    )
    args = parser.parse_args()

    uvicorn.run(app, host=args.host, port=args.port)