feat: improved logging with correlation ID

Author: JanssenBrm

app/config/logger.py

@@ -3,6 +3,7 @@
 import sys
 
 
+from app.middleware.correlation_id import correlation_id_ctx
 from loguru import logger
 
 
@@ -14,21 +15,34 @@ class InterceptHandler(logging.Handler):
     def emit(self, record):
         try:
             level = logger.level(record.levelname).name
+            corr_id = correlation_id_ctx.get()
         except ValueError:
             level = record.levelno
+        except LookupError:
+            corr_id = None
+
         frame, depth = logging.currentframe(), 2
         while frame and frame.f_code.co_filename == logging.__file__:
             frame = frame.f_back
             depth += 1
         logger.opt(depth=depth, exception=record.exc_info).log(
             level, record.getMessage()
         )
+        if corr_id:
+            logger.bind(correlation_id=corr_id)
+
+
+def correlation_id_filter(record):
+    # Always inject the current correlation ID into the log record
+    record["extra"]["correlation_id"] = correlation_id_ctx.get()
+    return True
 
 
 def setup_logging():
     logger.remove()  # remove default handler
     env = os.getenv("APP_ENV", "development")
 
+    logger.configure(extra={"correlation_id": None})
     if env == "production":
         # JSON logs for ELK
         logger.add(
@@ -37,6 +51,7 @@ def setup_logging():
             backtrace=False,
             diagnose=False,
             level="INFO",
+            filter=correlation_id_filter,
         )
     else:
         # Pretty logs for dev
@@ -45,13 +60,15 @@ def setup_logging():
             colorize=True,
             format=(
                 "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
+                "<i>{extra[correlation_id]}</i> | "
                 "<level>{level: <8}</level> | "
-                "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
+                "<cyan>{name}</cyan>:<cyan>{function}</cyan> :<cyan>{line}</cyan> - "
                 "{message}"
             ),
             backtrace=True,
             diagnose=True,
             level="DEBUG",
+            filter=correlation_id_filter,
         )
 
     for name in (

app/main.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI
 
+from app.middleware.correlation_id import add_correlation_id
 from app.platforms.dispatcher import load_processing_platforms
 from .config.logger import setup_logging
 from .config.settings import settings
@@ -15,6 +16,8 @@
     version="1.0.0",
 )
 
+app.middleware("http")(add_correlation_id)
+
 # Register Keycloak - must be done after FastAPI app creation
 # keycloak.register(app, prefix="/auth")  # mounts OIDC endpoints for login if needed
 

app/middleware/correlation_id.py

@@ -0,0 +1,33 @@
+import uuid
+import contextvars
+from fastapi import Request
+from loguru import logger
+
+# Context var to store correlation ID per request
+correlation_id_ctx = contextvars.ContextVar("correlation_id", default=None)
+
+
+async def add_correlation_id(request: Request, call_next):
+    # Try to reuse correlation ID from client headers, else generate one
+    corr_id = request.headers.get("X-Correlation-ID", str(uuid.uuid4()))
+    token = correlation_id_ctx.set(corr_id)
+
+    # Log at start
+    logger.bind(correlation_id=corr_id).info(
+        f"Incoming request: {request.method} - {request.url.path}",
+        method=request.method,
+        path=request.url.path,
+    )
+
+    response = await call_next(request)
+
+    # Add correlation ID to response headers
+    response.headers["X-Correlation-ID"] = corr_id
+
+    # Log at end
+    logger.bind(correlation_id=corr_id).info(
+        f"Request completed with status code {response.status_code}",
+        status_code=response.status_code,
+    )
+    correlation_id_ctx.reset(token)  # Reset after request
+    return response