158 lines
5.3 KiB
Python
158 lines
5.3 KiB
Python
"""Correlation id + observability middleware.
|
|
|
|
Every inbound request gets a ``correlation_id`` which is:
|
|
1. Injected as a structlog context variable for every log line.
|
|
2. Returned to the caller via the ``X-Correlation-Id`` response header.
|
|
3. Propagated to downstream calls (events, workers) that happen within
|
|
the request lifecycle.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass
|
|
from typing import Awaitable, Callable
|
|
|
|
import structlog
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse, Response
|
|
|
|
CORRELATION_HEADER = "X-Correlation-Id"
|
|
TENANT_HEADER = "X-ASF-Tenant-Id"
|
|
USER_HEADER = "X-ASF-User-Id"
|
|
|
|
# Deterministic UUID namespace for dev-trust string IDs.
|
|
# Any non-UUID string supplied via X-ASF headers is mapped to a stable UUIDv5
|
|
# so that the domain model (which uses UUID IDs internally) stays consistent
|
|
# while still allowing human-friendly strings like "t-demo-001" during local dev.
|
|
_ASP_NAMESPACE = uuid.UUID("9a8f2d58-14b4-40e1-93d3-1e4c2f0e6a22")
|
|
|
|
|
|
def _coerce_uuid(value: str) -> uuid.UUID:
|
|
"""Accept raw UUID hex strings or arbitrary opaque identifiers (UUID5 fallback)."""
|
|
try:
|
|
return uuid.UUID(value)
|
|
except ValueError:
|
|
return uuid.uuid5(_ASP_NAMESPACE, value)
|
|
|
|
correlation_id_var: ContextVar[str | None] = ContextVar(
|
|
"correlation_id", default=None
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class RequestContext:
|
|
"""Tenant + identity context extracted from the authenticated call.
|
|
|
|
For the bootstrap phase we accept ``X-ASF-Tenant-Id`` / ``X-ASF-User-Id``
|
|
headers directly when ``dev_trust_tenant_header`` is true in settings.
|
|
|
|
In production these fields MUST come from a validated JWT / session —
|
|
NEVER from raw client headers.
|
|
"""
|
|
|
|
tenant_id: uuid.UUID | None
|
|
user_id: uuid.UUID | None
|
|
correlation_id: str
|
|
roles: list[str]
|
|
|
|
def require_tenant(self) -> uuid.UUID:
|
|
if self.tenant_id is None:
|
|
from fastapi import HTTPException
|
|
raise HTTPException(status_code=401, detail="Missing authenticated tenant context")
|
|
return self.tenant_id
|
|
|
|
def require_user(self) -> uuid.UUID:
|
|
if self.user_id is None:
|
|
from fastapi import HTTPException
|
|
raise HTTPException(status_code=401, detail="Missing authenticated user")
|
|
return self.user_id
|
|
|
|
|
|
request_context_var: ContextVar[RequestContext | None] = ContextVar(
|
|
"request_context", default=None
|
|
)
|
|
|
|
|
|
def get_current_context() -> RequestContext:
|
|
ctx = request_context_var.get()
|
|
if ctx is None:
|
|
from fastapi import HTTPException
|
|
raise HTTPException(status_code=401, detail="No request context available")
|
|
return ctx
|
|
|
|
|
|
class CorrelationMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(
|
|
self,
|
|
request: Request,
|
|
call_next: Callable[[Request], Awaitable[Response]],
|
|
) -> Response:
|
|
cid = request.headers.get(CORRELATION_HEADER) or uuid.uuid4().hex
|
|
correlation_id_var.set(cid)
|
|
|
|
with structlog.contextvars.bound_contextvars(correlation_id=cid):
|
|
try:
|
|
ctx = _extract_context(request, cid)
|
|
except _ContextError as exc:
|
|
# BaseHTTPMiddleware runs before FastAPI exception handlers.
|
|
# Build the error response directly instead of raising HTTPException.
|
|
resp = JSONResponse(
|
|
status_code=exc.status_code,
|
|
content={"error": exc.kind, "message": exc.message},
|
|
)
|
|
resp.headers[CORRELATION_HEADER] = cid
|
|
return resp
|
|
|
|
request_context_var.set(ctx)
|
|
structlog.contextvars.bind_contextvars(
|
|
tenant_id=str(ctx.tenant_id) if ctx.tenant_id else None,
|
|
user_id=str(ctx.user_id) if ctx.user_id else None,
|
|
)
|
|
|
|
response = await call_next(request)
|
|
response.headers[CORRELATION_HEADER] = cid
|
|
return response
|
|
|
|
|
|
class _ContextError(Exception):
|
|
"""Raised internally by _extract_context to short-circuit auth failures."""
|
|
|
|
def __init__(self, status_code: int, kind: str, message: str) -> None:
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.kind = kind
|
|
self.message = message
|
|
|
|
|
|
def _extract_context(request: Request, cid: str) -> RequestContext:
|
|
from ..config import get_settings
|
|
|
|
s = get_settings()
|
|
tenant_id: uuid.UUID | None = None
|
|
user_id: uuid.UUID | None = None
|
|
roles: list[str] = []
|
|
|
|
# ---- Bootstrap-trust mode for local dev ----
|
|
if s.dev_trust_tenant_header and TENANT_HEADER in request.headers:
|
|
raw = request.headers[TENANT_HEADER].strip()
|
|
if not raw:
|
|
raise _ContextError(400, "invalid_header", f"Empty {TENANT_HEADER}")
|
|
tenant_id = _coerce_uuid(raw)
|
|
|
|
if s.dev_trust_user_header and USER_HEADER in request.headers:
|
|
raw = request.headers[USER_HEADER].strip()
|
|
if not raw:
|
|
raise _ContextError(400, "invalid_header", f"Empty {USER_HEADER}")
|
|
user_id = _coerce_uuid(raw)
|
|
roles = ["admin"]
|
|
|
|
return RequestContext(
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
correlation_id=cid,
|
|
roles=roles,
|
|
)
|