34 lines
835 B
Python
34 lines
835 B
Python
import time
|
|
import redis.asyncio as redis
|
|
from ..core import get_settings
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
class RateLimitService:
|
|
def __init__(self):
|
|
self.redis: redis.Redis | None = None
|
|
|
|
async def connect(self):
|
|
self.redis = redis.from_url(settings.redis_url)
|
|
|
|
async def disconnect(self):
|
|
if self.redis:
|
|
await self.redis.close()
|
|
|
|
async def is_allowed(self, key: str, limit: int | None = None) -> bool:
|
|
if not self.redis:
|
|
return True
|
|
limit = limit or settings.rate_limit_per_minute
|
|
now = int(time.time())
|
|
window_key = f"rl:{key}:{now // 60}"
|
|
|
|
count = await self.redis.incr(window_key)
|
|
if count == 1:
|
|
await self.redis.expire(window_key, 60)
|
|
|
|
return count <= limit
|
|
|
|
|
|
rate_limit_service = RateLimitService()
|