Fastapi - Ratelimiting
Fastapi & Rate Limiting
Let’s test a simple Rate Limiting Function found on the Net …
main.py
Main App with Rate Limiting Function
# main.py
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from Token_bucket import TokenBucket
app = FastAPI()
class RateLimiterMiddleware(BaseHTTPMiddleware):
def __init__(self, app, bucket: TokenBucket):
super().__init__(app)
self.bucket = bucket
async def dispatch(self, request: Request, call_next):
if self.bucket.take_token():
return await call_next(request)
# Return a JSON response for rate limit exceeded
return JSONResponse(status_code=429, content={"detail": "Rate limit exceeded"})
# Token bucket with a capacity of 3 and refill rate of 1 token per second
bucket = TokenBucket(capacity=3, refill_rate=3)
# Apply the rate limiter middleware
app.add_middleware(RateLimiterMiddleware, bucket=bucket)
@app.get("/")
async def read_root():
return {"message": "Hello World"}
Token_bucket.py
# Token_bucket.py
import time
class TokenBucket:
def __init__(self, capacity, refill_rate):
self.capacity = capacity
self.refill_rate = refill_rate
self.tokens = capacity
self.last_refill = time.time()
def add_tokens(self):
now = time.time()
if self.tokens < self.capacity:
tokens_to_add = (now - self.last_refill) * self.refill_rate
self.tokens = min (self.capacity, self.tokens + tokens_to_add)
self.last_refill=now
def take_token(self):
self.add_tokens()
if self.tokens >= 1:
self.tokens -=1
return True
return False
Test Script
Let’s produce some requests …
