zanewalker
f3453cb0fc
- Refactor Redis backend connection handling and pool management - Update algorithm implementations with improved type annotations - Enhance config loader validation with stricter Pydantic schemas - Improve decorator and middleware error handling - Expand example scripts with better docstrings and usage patterns - Add new 00_basic_usage.py example for quick start - Reorganize examples directory structure - Fix type annotation inconsistencies across core modules - Update dependencies in pyproject.toml
82 lines
2.0 KiB
Python
82 lines
2.0 KiB
Python
"""Quickstart example - minimal setup to get rate limiting working."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from fastapi_traffic import (
|
|
MemoryBackend,
|
|
RateLimiter,
|
|
RateLimitExceeded,
|
|
rate_limit,
|
|
)
|
|
from fastapi_traffic.core.limiter import set_limiter
|
|
|
|
DEFAULT_HOST = "127.0.0.1"
|
|
DEFAULT_PORT = 8000
|
|
|
|
# Step 1: Create a backend and limiter
|
|
backend = MemoryBackend()
|
|
limiter = RateLimiter(backend)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
"""Lifespan context manager for startup/shutdown."""
|
|
await limiter.initialize()
|
|
set_limiter(limiter)
|
|
yield
|
|
await limiter.close()
|
|
|
|
|
|
app = FastAPI(title="Quickstart Example", lifespan=lifespan)
|
|
|
|
|
|
# Step 2: Add exception handler for rate limit errors
|
|
@app.exception_handler(RateLimitExceeded)
|
|
async def rate_limit_handler(_: Request, exc: RateLimitExceeded) -> JSONResponse:
|
|
return JSONResponse(
|
|
status_code=429,
|
|
content={"error": "Too many requests", "retry_after": exc.retry_after},
|
|
)
|
|
|
|
|
|
# Step 3: Apply rate limiting to endpoints
|
|
@app.get("/")
|
|
@rate_limit(10, 60) # 10 requests per minute
|
|
async def hello(_: Request) -> dict[str, str]:
|
|
return {"message": "Hello, World!"}
|
|
|
|
|
|
@app.get("/api/data")
|
|
@rate_limit(100, 60) # 100 requests per minute
|
|
async def get_data(_: Request) -> dict[str, str]:
|
|
return {"data": "Some important data"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
import uvicorn
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Quickstart 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)
|