feat: implement correlation ID tracking and RFC 7807 error handling

Author: yash-blip-ys

docker-compose.yaml

@@ -14,6 +14,7 @@ services:
       - ADMIN_ROLE_NAME=$ADMIN_ROLE_NAME
       - ES_USER=$ES_USER
       - ES_PASSWORD=$ES_PASSWORD
+      - PYTHONUNBUFFERED=1
     ports:
       - ${AIOD_REST_PORT}:8000
     command: python main.py

src/error_handling/error_handling.py

@@ -1,14 +1,11 @@
 import json
 import logging
 import traceback
-import uuid
 from http import HTTPStatus
 
 from fastapi import HTTPException, status
-from pydantic import BaseModel
 from starlette.responses import JSONResponse
 
-
 def as_http_exception(exception: Exception) -> HTTPException:
     if isinstance(exception, HTTPException):
         return exception
@@ -21,33 +18,46 @@ def as_http_exception(exception: Exception) -> HTTPException:
         ),
     )
 
-
-class ErrorSchema(BaseModel):
-    detail: str
-    reference: str
-
-
 async def http_exception_handler(request, exc):
-    reference = uuid.uuid4().hex
-    error = ErrorSchema(detail=exc.detail, reference=reference)
-    content = error.dict()
-
+    # 1.access the 'Single Source of Truth' id from your middleware this is the 'correlation_id' we set in the CorrelationIdMiddleware
+    correlation_id = getattr(request.state, "correlation_id", "unknown")
+
+    # 2.hey team built the rfc7807 'Problem Details' structure this should help
+    
+    content = {
+        "type": "about:blank",
+        "title": HTTPStatus(exc.status_code).phrase,
+        "status": exc.status_code,
+        "detail": exc.detail,
+        "instance": request.url.path,
+        "correlation_id": correlation_id 
+    }
+
+    # 3.enhanced logging logic log the error with all relevant details, including the correlation id, request method, path, and exception details. This will help in debugging and tracing issues effectively.
     body_content = "<Data Stream with unknown content>"
     if not request._stream_consumed:
-        body = await request.body()
-        body_content = json.dumps(json.loads(body)) if body else ""
-
-    log_message = str(
-        dict(
-            reference=reference,
-            exception=f"{str(exc)!r}",
-            method=request.scope["method"],
-            path=request.scope["path"],
-            body=body_content,
-        )
-    )
-    log_level = logging.DEBUG
-    if exc.status_code == HTTPStatus.INTERNAL_SERVER_ERROR:
-        log_level = logging.WARNING
-    logging.log(log_level, log_message)
-    return JSONResponse(content, status_code=exc.status_code)
+        try:
+            body = await request.body()
+            body_content = json.loads(body) if body else ""
+        except Exception:
+            body_content = "<Unparseable Body>"
+
+    log_data = {
+        "correlation_id": correlation_id,
+        "status_code": exc.status_code,
+        "method": request.method,
+        "path": request.url.path,
+        "exception": str(exc.detail),
+        "body": body_content,
+    }
+
+    #log as WARNING for server errors (500+) and info for (400+) just to avoid log noise, but still capture all details for debugging when needed.
+    log_level = logging.ERROR if exc.status_code >= 500 else logging.INFO
+    logging.log(log_level, f"API Error: {json.dumps(log_data)}")
+
+    #4.return with the specific 'problem+json' media type as per rfc7807
+    return JSONResponse(
+        content=content,
+        status_code=exc.status_code,
+        media_type="application/problem+json"
+    )

src/main.py

@@ -4,6 +4,19 @@
 Note: order matters for overloaded paths
 (https://fastapi.tiangolo.com/tutorial/path-params/#order-matters).
 """
+import uuid
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request
+class CorrelationIdMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        #1.it  generate the unique id i  tried to use only single id for the request and response cycle, so that we can easily trace the logs and the response for a particular request.
+        correlation_id = str(uuid.uuid4())
+        request.state.correlation_id = correlation_id
+        #2.pass the request to the next person in line
+        response = await call_next(request)
+        #3.inject the correlation id into response headers for client side tracking (Egress) stamped the id on the way out so the user sees it
+        response.headers["X-Correlation-ID"] = getattr(request.state, "correlation_id", "not-started")
+        return response
 
 import argparse
 import logging
@@ -12,6 +25,9 @@
 from importlib.metadata import version as pkg_version, PackageNotFoundError
 import uvicorn
 from fastapi import Depends, FastAPI, HTTPException
+from starlette.exceptions import HTTPException as StarletteHTTPException
+from fastapi.exceptions import HTTPException as FastAPIHTTPException
+from error_handling.error_handling import http_exception_handler
 from fastapi.responses import HTMLResponse
 from sqlmodel import select, SQLModel
 from starlette.requests import Request
@@ -33,7 +49,7 @@
 from setup_logger import setup_logger
 from taxonomies.synchronize_taxonomy import synchronize_taxonomy_from_file
 from triggers import disable_review_process, enable_review_process
-from error_handling import http_exception_handler
+from error_handling.error_handling import http_exception_handler
 from routers import (
     resource_routers,
     parent_routers,
@@ -53,6 +69,16 @@
     add_deprecation_and_sunset_middleware,
     Version,
 )
+import logging
+import sys
+
+#just to be sure we can look for particular correlation ids in the logs without having to parse the entire log line
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(levelname)s: %(message)s",
+    stream=sys.stdout, 
+    force=True  #this overrides any existing hidden configs
+)
 
 
 def add_routes(app: FastAPI, version: Version, url_prefix=""):
@@ -133,6 +159,8 @@ def create_app() -> FastAPI:
     except PackageNotFoundError:
         dist_version = "dev"
     app = build_app(url_prefix=DEV_CONFIG.get("url_prefix", ""), version=dist_version)
+    
+    
     return app
 
 
@@ -159,6 +187,7 @@ def build_app(*, url_prefix: str = "", version: str = "dev"):
         version="latest",
         **kwargs,
     )
+    main_app.add_middleware(CorrelationIdMiddleware)
     versioned_apps = [
         (
             FastAPI(
@@ -173,15 +202,20 @@ def build_app(*, url_prefix: str = "", version: str = "dev"):
     ]
     for app, version in [(main_app, Version.LATEST)] + versioned_apps:
         add_routes(app, version=version)
-        app.add_exception_handler(HTTPException, http_exception_handler)
+        
+        app.exception_handlers[FastAPIHTTPException] = http_exception_handler
+        #this is needed to catch exceptions raised by Starlette, such as 404s for non existent endpoints which are not caught by FastAPI HTTPException handler
+        app.exception_handlers[StarletteHTTPException] = http_exception_handler
+        app.add_exception_handler(404, http_exception_handler)
+        
         add_deprecation_and_sunset_middleware(app)
         add_version_to_openapi(app)
 
     Instrumentator().instrument(main_app).expose(
         main_app, endpoint="/metrics", include_in_schema=False
     )
-    # Since all traffic goes through the main app, this middleware only
-    # needs to be registered with the main app and not the mounted apps.
+    #Since all traffic goes through the main app this middleware only
+    #needs to be registered with the main app and not the mounted apps
     main_app.add_middleware(AccessLogMiddleware)
 
     for app, _ in versioned_apps:
@@ -231,4 +265,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()