rate limit redis pool and initial setup

Author: igorbenav

README.md

@@ -64,7 +64,7 @@
 - [ ] Docs for other databases (MysQL, SQLite)
 
 #### Features
-- [ ] Add a Rate Limiter decorator
+- [ ] Add a Rate Limiter dependency
 - [ ] Add mongoDB support
 
 #### Tests

src/app/api/dependencies.py

@@ -35,7 +35,7 @@ async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)], db: An
 
     return user
 
-def get_current_superuser(current_user: Annotated[User, Depends(get_current_user)]) -> User:
+async def get_current_superuser(current_user: Annotated[User, Depends(get_current_user)]) -> User:
     if not current_user.is_superuser:
         raise privileges_exception
 

src/app/core/cache.py

@@ -1,8 +1,6 @@
 from typing import Callable, Union, List, Dict, Any
 import functools
 import json
-from uuid import UUID
-from datetime import datetime
 import re
 
 from fastapi import Request, Response

src/app/core/config.py

@@ -84,6 +84,12 @@ class RedisQueueSettings(BaseSettings):
     REDIS_QUEUE_PORT: str = config("REDIS_QUEUE_PORT", default=6379)
 
 
+class RedisRateLimiterSettings(BaseSettings):
+    REDIS_RATE_LIMITER_HOST: str = config("REDIS_RATE_LIMITER_HOST", default="localhost")
+    REDIS_RATE_LIMITER_PORT: int = config("REDIS_RATE_LIMITER_PORT", default=6379)
+    REDIS_RATE_LIMITER_URL: str = f"redis://{REDIS_RATE_LIMITER_HOST}:{REDIS_RATE_LIMITER_PORT}"
+
+
 class EnvironmentOption(Enum):
     LOCAL = "local"
     STAGING = "staging"
@@ -103,6 +109,7 @@ class Settings(
     RedisCacheSettings,
     ClientSideCacheSettings,
     RedisQueueSettings,
+    RedisRateLimiterSettings,
     EnvironmentSettings
 ):
     pass

src/app/core/rate_limit.py

@@ -0,0 +1,9 @@
+import time
+
+from fastapi import HTTPException, Request, Depends, status
+from redis.asyncio import Redis, ConnectionPool
+
+from app.api.dependencies import get_current_user
+
+pool: ConnectionPool | None = None
+client: Redis | None = None

src/app/core/setup.py

@@ -8,7 +8,7 @@
 
 
 from app.api.dependencies import get_current_superuser
-from app.core import cache, queue
+from app.core import cache, queue, rate_limit
 from app.core.config import settings
 from app.core.database import Base
 from app.core.database import async_engine as engine
@@ -18,6 +18,7 @@
     AppSettings, 
     ClientSideCacheSettings, 
     RedisQueueSettings,
+    RedisRateLimiterSettings,
     EnvironmentOption,
     EnvironmentSettings
 )
@@ -49,6 +50,16 @@ async def close_redis_queue_pool():
     await queue.pool.aclose()
 
 
+# -------------- rate limit --------------
+async def create_redis_rate_limit_pool():
+    rate_limit.pool = redis.ConnectionPool.from_url(settings.REDIS_RATE_LIMIT_URL)
+    rate_limit.client = redis.Redis.from_pool(rate_limit.pool)
+
+
+async def close_redis_rate_limit_pool():
+    await rate_limit.client.aclose()
+
+
 # -------------- application --------------
 async def set_threadpool_tokens(number_of_tokens=100):
     limiter = anyio.to_thread.current_default_thread_limiter()
@@ -133,6 +144,10 @@ def create_application(router: APIRouter, settings, **kwargs) -> FastAPI:
         application.add_event_handler("startup", create_redis_queue_pool)
         application.add_event_handler("shutdown", close_redis_queue_pool)
 
+    if isinstance(settings, RedisRateLimiterSettings):
+        application.add_event_handler("startup", create_redis_rate_limit_pool)
+        application.add_event_handler("shutdown", close_redis_rate_limit_pool)
+
     if isinstance(settings, EnvironmentSettings):
         if settings.ENVIRONMENT != EnvironmentOption.PRODUCTION:
             docs_router = APIRouter()