94 lines
2.4 KiB
Python
94 lines
2.4 KiB
Python
"""
|
|
Request ID middleware for request tracing.
|
|
|
|
Generates or propagates X-Request-ID header for distributed tracing.
|
|
"""
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Optional
|
|
|
|
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.responses import Response
|
|
|
|
from logging_config import request_id_var
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RequestIDMiddleware(BaseHTTPMiddleware):
|
|
"""
|
|
HTTP middleware for request ID generation and propagation.
|
|
|
|
- Extracts X-Request-ID from incoming request headers
|
|
- Generates a new UUID if not present
|
|
- Sets request_id in context var for logging
|
|
- Adds X-Request-ID to response headers
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
app,
|
|
header_name: str = "X-Request-ID",
|
|
generator: Optional[callable] = None,
|
|
):
|
|
"""
|
|
Initialize request ID middleware.
|
|
|
|
Args:
|
|
app: FastAPI application.
|
|
header_name: Header name for request ID.
|
|
generator: Optional custom ID generator function.
|
|
"""
|
|
super().__init__(app)
|
|
self.header_name = header_name
|
|
self.generator = generator or (lambda: str(uuid.uuid4()))
|
|
|
|
async def dispatch(self, request: Request, call_next) -> Response:
|
|
"""
|
|
Process request with request ID.
|
|
|
|
Args:
|
|
request: Incoming HTTP request.
|
|
call_next: Next middleware/handler in chain.
|
|
|
|
Returns:
|
|
HTTP response with X-Request-ID header.
|
|
"""
|
|
# Get or generate request ID
|
|
request_id = request.headers.get(self.header_name)
|
|
if not request_id:
|
|
request_id = self.generator()
|
|
|
|
# Set in request state for access in handlers
|
|
request.state.request_id = request_id
|
|
|
|
# Set in context var for logging
|
|
token = request_id_var.set(request_id)
|
|
|
|
try:
|
|
# Process request
|
|
response = await call_next(request)
|
|
|
|
# Add request ID to response
|
|
response.headers[self.header_name] = request_id
|
|
|
|
return response
|
|
finally:
|
|
# Reset context var
|
|
request_id_var.reset(token)
|
|
|
|
|
|
def get_request_id(request: Request) -> Optional[str]:
|
|
"""
|
|
Get request ID from request state.
|
|
|
|
Args:
|
|
request: FastAPI request object.
|
|
|
|
Returns:
|
|
Request ID string or None.
|
|
"""
|
|
return getattr(request.state, "request_id", None)
|