zanewalker
da496746bb
- Core rate limiting with multiple algorithms (sliding window, token bucket, etc.) - SQLite and memory backends - Decorator and dependency injection patterns - Middleware support - Example usage files
90 lines
1.9 KiB
Python
90 lines
1.9 KiB
Python
"""Data models for rate limiting."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
|
|
class KeyType(str, Enum):
|
|
"""Type of key extraction for rate limiting."""
|
|
|
|
IP = "ip"
|
|
USER = "user"
|
|
API_KEY = "api_key"
|
|
ENDPOINT = "endpoint"
|
|
CUSTOM = "custom"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RateLimitInfo:
|
|
"""Information about the current rate limit state."""
|
|
|
|
limit: int
|
|
remaining: int
|
|
reset_at: float
|
|
retry_after: float | None = None
|
|
window_size: float = 60.0
|
|
|
|
def to_headers(self) -> dict[str, str]:
|
|
"""Convert rate limit info to HTTP headers."""
|
|
headers: dict[str, str] = {
|
|
"X-RateLimit-Limit": str(self.limit),
|
|
"X-RateLimit-Remaining": str(max(0, self.remaining)),
|
|
"X-RateLimit-Reset": str(int(self.reset_at)),
|
|
}
|
|
if self.retry_after is not None:
|
|
headers["Retry-After"] = str(int(self.retry_after))
|
|
return headers
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RateLimitResult:
|
|
"""Result of a rate limit check."""
|
|
|
|
allowed: bool
|
|
info: RateLimitInfo
|
|
key: str
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TokenBucketState:
|
|
"""State for token bucket algorithm."""
|
|
|
|
tokens: float
|
|
last_update: float
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class SlidingWindowState:
|
|
"""State for sliding window algorithm."""
|
|
|
|
timestamps: list[float] = field(default_factory=list)
|
|
count: int = 0
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class FixedWindowState:
|
|
"""State for fixed window algorithm."""
|
|
|
|
count: int
|
|
window_start: float
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class LeakyBucketState:
|
|
"""State for leaky bucket algorithm."""
|
|
|
|
water_level: float
|
|
last_update: float
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class BackendRecord:
|
|
"""Generic record stored in backends."""
|
|
|
|
key: str
|
|
data: dict[str, Any]
|
|
expires_at: float
|