72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
"""
|
|
Correlation ID Middleware for FastAPI
|
|
Adds unique request ID for distributed tracing
|
|
"""
|
|
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.responses import Response
|
|
import uuid
|
|
import time
|
|
from typing import Callable
|
|
|
|
class CorrelationIdMiddleware(BaseHTTPMiddleware):
|
|
"""
|
|
Middleware that adds correlation ID to all requests for distributed tracing
|
|
"""
|
|
|
|
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
|
# Check for existing correlation ID in headers
|
|
existing_id = (
|
|
request.headers.get('x-correlation-id') or
|
|
request.headers.get('x-request-id') or
|
|
request.headers.get('traceparent')
|
|
)
|
|
|
|
# Generate new ID if not provided
|
|
correlation_id = existing_id or self._generate_correlation_id()
|
|
|
|
# Attach to request state
|
|
request.state.correlation_id = correlation_id
|
|
request.state.request_start_time = time.time()
|
|
|
|
# Process request
|
|
response = await call_next(request)
|
|
|
|
# Add correlation ID to response headers
|
|
response.headers['X-Correlation-ID'] = correlation_id
|
|
response.headers['X-Request-ID'] = correlation_id
|
|
|
|
return response
|
|
|
|
@staticmethod
|
|
def _generate_correlation_id() -> str:
|
|
"""
|
|
Generate a unique correlation ID
|
|
Format: timestamp-random (e.g., 1703423456789-a1b2c3d4)
|
|
"""
|
|
timestamp = int(time.time() * 1000)
|
|
random_part = uuid.uuid4().hex[:8]
|
|
return f"{timestamp}-{random_part}"
|
|
|
|
def get_correlation_id(request: Request) -> str:
|
|
"""
|
|
Get correlation ID from request (for use in services)
|
|
|
|
Args:
|
|
request: FastAPI request object
|
|
|
|
Returns:
|
|
Correlation ID string
|
|
"""
|
|
return getattr(request.state, 'correlation_id', 'unknown')
|
|
|
|
def generate_correlation_id() -> str:
|
|
"""
|
|
Generate a new correlation ID
|
|
|
|
Returns:
|
|
Correlation ID string
|
|
"""
|
|
return CorrelationIdMiddleware._generate_correlation_id()
|
|
|