Merge pull request #14 from wednesday-solutions/feat/request-id-injec…

…tion

Feat/request id injection

app.py

@@ -10,10 +10,12 @@
 from routes.users import user
 from fastapi_pagination import add_pagination
 from middlewares.rate_limiter_middleware import RateLimitMiddleware
+from middlewares.request_id_injection import RequestIdInjection
 from pybreaker import CircuitBreakerError
 from dependencies import circuit_breaker
 from utils.slack_notification_utils import send_slack_message
 import traceback
+from middlewares.request_id_injection import request_id_contextvar
 
 # Initializing the swagger docs
 app = FastAPI(
@@ -35,13 +37,14 @@
     allow_headers=["*"],
 )
 app.add_middleware(RateLimitMiddleware)
+app.add_middleware(RequestIdInjection)
 app.include_router(user, prefix="/user")
 
 
 # Default API route
 @app.get("/")
 async def read_main():
-    1/0
+    print('Request ID:', request_id_contextvar.get())
     return {"response": "service up and running..!"}
 
 
@@ -88,8 +91,12 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 async def http_exception_handler(request: Request, exc: Exception):
     error_message = f'Error: {str(exc)}'
     # Include the traceback in the response for debugging purposes
-    traceback_str = traceback.format_exc()
-    send_slack_message({ "text": f'```{traceback_str}```', "request_url": str(request.url), "request_method": str(request.method)})
+    traceback_str = traceback.format_exc(chain=False)
+    send_slack_message(
+        { 
+            "text": f'```\nRequestID: {request_id_contextvar.get()}\nRequest URL: {str(request.url)} \nRequest_method: {str(request.method)} \nTraceback: {traceback_str}```'
+        }
+    )
 
     return JSONResponse(
         status_code=500,

middlewares/request_id_injection.py

@@ -0,0 +1,22 @@
+from fastapi import FastAPI
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from fastapi.responses import JSONResponse
+import contextvars
+import uuid
+
+request_id_contextvar = contextvars.ContextVar("request_id", default=None)
+
+class RequestIdInjection(BaseHTTPMiddleware):
+    def dispatch(self, request: Request, call_next):
+        request_id = str(uuid.uuid4())
+        request_id_contextvar.set(request_id)
+        try:
+            return call_next(request)
+
+        except Exception as ex:
+            print(ex)
+            return JSONResponse(content={"success": False}, status_code=500)
+
+        finally:
+            assert request_id_contextvar.get() == request_id

requirements.txt

@@ -115,4 +115,5 @@ flake8
 black
 pre-commit
 pybreaker
-pytest-asyncio
+pytest-asyncio
+contextvars

routes/users.py

@@ -15,6 +15,7 @@
 from utils.user_utils import get_current_user
 from typing import Annotated
 from fastapi.security import HTTPBearer
+from middlewares.request_id_injection import request_id_contextvar
 
 user = APIRouter()
 
@@ -23,12 +24,14 @@
 
 @user.post("/register", tags=["Users"])
 def register(payload: CreateUser, db: Session = Depends(create_local_session)):
+    print('Request ID:', request_id_contextvar.get())
     response = create_user_dao(data=payload, dbSession=db)
     return response
 
 
 @user.post("/signin", tags=["Users"])
 def login(payload: Login, db: Session = Depends(create_local_session)):
+    print('Request ID:', request_id_contextvar.get())
     response = signin(data=payload, dbSession=db)
     return response
 
@@ -40,17 +43,20 @@ async def profile(
     db: Session = Depends(create_local_session),
     redis=Depends(get_redis),
 ):
+    print('Request ID:', request_id_contextvar.get())
     # Here, you can use 'redis' to fetch or store data in Redis cache
     response = await get_user_dao(user_id, dbSession=db, redis=redis)
     return response
 
 
 @user.get("/", tags=["Users"], response_model=Page[UserOutResponse])
 def list_users(db: Session = Depends(create_local_session)):
+    print('Request ID:', request_id_contextvar.get())
     response = list_users_dao(dbSession=db)
     return response
 
 
 @user.get("/{user_id}/secure-route/", tags=["Users"], dependencies=[Depends(get_current_user)])
 def secure_route(token: Annotated[str, Depends(httpBearerScheme)], user_id: int):
+    print('Request ID:', request_id_contextvar.get())
     return {"message": "If you see this, you're authenticated"}