feat: add redis for api caching

Author: Pranavseelam06

app/crud.py

@@ -1,6 +1,23 @@
-from sqlalchemy.orm import Session
 from sqlalchemy import select
+from sqlalchemy.orm import Session
+
 from app.models import APIKey
+from app.redis_client import redis_client, get_cached_api_key, cache_api_key
+
+
+def get_api_key(db: Session, key_value: str):
+    cached_key = get_cached_api_key(redis_client, key_value)
+
+    if cached_key is not None:
+        return cached_key
+
+    api_key = db.execute(
+        select(APIKey).where(APIKey.key == key_value)
+    ).scalar_one_or_none()
+
+    if api_key is None:
+        return None
+
+    cache_api_key(redis_client, key_value, api_key)
 
-def get_api_key(db: Session, key_value: str) -> APIKey | None:
-    return db.execute(select(APIKey).where(APIKey.key == key_value)).scalar_one_or_none()
+    return api_key

app/main.py

@@ -7,6 +7,7 @@
 from app.database import SessionLocal
 from app.models import RequestLog
 from contextlib import asynccontextmanager
+from app.redis_client import redis_client
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):

app/rate_limit_dependency.py

@@ -3,20 +3,23 @@
 from app.models import APIKey
 from app.rate_limiter import TokenBucket
 
-
-
 def enforce_rate_limit(
     request: Request,
     api_key: APIKey = Depends(require_api_key),
 ) -> APIKey:
     buckets = request.app.state.buckets
-    if api_key.key not in buckets:
-        buckets[api_key.key] = TokenBucket(
+
+    bucket_key = f"api_key:{api_key.id}"
+
+    if bucket_key not in buckets:
+        buckets[bucket_key] = TokenBucket(
             capacity=api_key.rate_limit_capacity,
             refill_rate=api_key.rate_limit_refill_rate,
         )
-    bucket = buckets[api_key.key]
+
+    bucket = buckets[bucket_key]
+
     if not bucket.consume():
         raise HTTPException(status_code=429, detail="Rate limit exceeded")
-    return api_key
-        
+
+    return api_key

app/redis_client.py

@@ -0,0 +1,50 @@
+import hashlib
+import json
+import os
+from types import SimpleNamespace
+
+import redis
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+redis_client = redis.Redis.from_url(
+    os.getenv("REDIS_URL", "redis://localhost:6379/0"),
+    decode_responses=True,
+)
+
+
+def api_key_cache_key(raw_api_key: str) -> str:
+    hashed = hashlib.sha256(raw_api_key.encode()).hexdigest()
+    return f"api_key:{hashed}"
+
+
+def cache_api_key(redis_client, key_value: str, api_key) -> None:
+    cache_key = api_key_cache_key(key_value)
+
+    data = {
+        "id": api_key.id,
+        "name": api_key.name,
+        "rate_limit_capacity": api_key.rate_limit_capacity,
+        "rate_limit_refill_rate": api_key.rate_limit_refill_rate,
+    }
+
+    redis_client.setex(
+        cache_key,
+        300,
+        json.dumps(data),
+    )
+
+
+def get_cached_api_key(redis_client, key_value: str):
+    cache_key = api_key_cache_key(key_value)
+
+    cached = redis_client.get(cache_key)
+
+    if cached is None:
+        return None
+
+    data = json.loads(cached)
+
+    return SimpleNamespace(**data)

requirements.txt

@@ -22,6 +22,7 @@ Pygments==2.20.0
 pytest==9.0.3
 python-dotenv==1.2.2
 PyYAML==6.0.3
+redis==8.0.0
 respx==0.23.1
 SQLAlchemy==2.0.49
 starlette==1.0.0