110 lines
3.1 KiB
Python
110 lines
3.1 KiB
Python
"""Examples demonstrating different storage backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from fastapi_traffic import (
|
|
MemoryBackend,
|
|
RateLimiter,
|
|
RateLimitExceeded,
|
|
SQLiteBackend,
|
|
rate_limit,
|
|
)
|
|
from fastapi_traffic.core.limiter import set_limiter
|
|
|
|
|
|
# Choose backend based on environment
|
|
def get_backend():
|
|
"""Select appropriate backend based on environment."""
|
|
backend_type = os.getenv("RATE_LIMIT_BACKEND", "memory")
|
|
|
|
if backend_type == "sqlite":
|
|
# SQLite - Good for single-instance apps, persists across restarts
|
|
return SQLiteBackend("rate_limits.db")
|
|
|
|
elif backend_type == "redis":
|
|
# Redis - Required for distributed/multi-instance deployments
|
|
# Requires: pip install redis
|
|
try:
|
|
import asyncio
|
|
|
|
from fastapi_traffic import RedisBackend
|
|
|
|
async def create_redis():
|
|
return await RedisBackend.from_url(
|
|
os.getenv("REDIS_URL", "redis://localhost:6379/0"),
|
|
key_prefix="myapp_ratelimit",
|
|
)
|
|
|
|
return asyncio.get_event_loop().run_until_complete(create_redis())
|
|
except ImportError:
|
|
print("Redis not installed, falling back to memory backend")
|
|
return MemoryBackend()
|
|
|
|
else:
|
|
# Memory - Fast, but resets on restart, not shared across instances
|
|
return MemoryBackend()
|
|
|
|
|
|
backend = get_backend()
|
|
limiter = RateLimiter(backend)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
await limiter.initialize()
|
|
set_limiter(limiter)
|
|
yield
|
|
await limiter.close()
|
|
|
|
|
|
app = FastAPI(title="Storage Backends Example", lifespan=lifespan)
|
|
|
|
|
|
@app.exception_handler(RateLimitExceeded)
|
|
async def rate_limit_handler(_: Request, exc: RateLimitExceeded) -> JSONResponse:
|
|
return JSONResponse(
|
|
status_code=429,
|
|
content={"error": "rate_limit_exceeded", "retry_after": exc.retry_after},
|
|
)
|
|
|
|
|
|
@app.get("/api/resource")
|
|
@rate_limit(100, 60)
|
|
async def get_resource(_: Request) -> dict[str, str]:
|
|
return {"message": "Resource data", "backend": type(backend).__name__}
|
|
|
|
|
|
@app.get("/backend-info")
|
|
async def backend_info() -> dict[str, Any]:
|
|
"""Get information about the current backend."""
|
|
info = {
|
|
"backend_type": type(backend).__name__,
|
|
"description": "",
|
|
}
|
|
|
|
if isinstance(backend, MemoryBackend):
|
|
info["description"] = "In-memory storage, fast but ephemeral"
|
|
elif isinstance(backend, SQLiteBackend):
|
|
info["description"] = "SQLite storage, persistent, single-instance"
|
|
else:
|
|
info["description"] = "Redis storage, distributed, multi-instance"
|
|
|
|
return info
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
# Run with different backends:
|
|
# RATE_LIMIT_BACKEND=memory python 03_backends.py
|
|
# RATE_LIMIT_BACKEND=sqlite python 03_backends.py
|
|
# RATE_LIMIT_BACKEND=redis REDIS_URL=redis://localhost:6379/0 python 03_backends.py
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|