Include unique ID in 500 responses and link to Datadog trace

[amacfie]

Logging tools like Datadog or Sentry can record unhandled exceptions that lead to HTTP 500 responses, but what if a user receives a 500 response and you want to look up the log entry corresponding to that specific request? The following code adds a x-error-id header to 500 responses for unhandled exceptions. The value is a unique ID which is also saved as a tag associated with the trace in Datadog.

import sys
import traceback

from fastapi import FastAPI
from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint,
)
from starlette.requests import Request


app = FastAPI()

def add_unhandled_exception_handler(app: FastAPI):
    async def gotta_catch_em_all(
        request: Request, call_next: RequestResponseEndpoint
    ):
        try:
            return await call_next(request)
        except Exception as ex:
            # Datadog's instrumentation already attaches info to the trace about
            # unhandled exceptions but we're bypassing that so we do it manually
            # here. see https://github.com/DataDog/dd-trace-py/blob/b50bc0f812d90d2d8bc1f822e977cc86831d7769/ddtrace/contrib/asgi/middleware.py#L200
            (exc_type, exc_val, exc_tb) = sys.exc_info()
            span = tracer.current_span()
            assert span is not None
            span.set_exc_info(exc_type, exc_val, exc_tb)

            # since we're not re-raising the exception, also print error
            traceback.print_exception(ex, file=sys.stderr)

            # we're using the span ID as a unique ID but you could use any
            # random string
            span.set_tag("error_id", span.span_id)
            return Response(
                status_code=500, headers={"x-error-id": str(span.span_id)}
            )

    # this is what the `app.middleware` function does but we add to end of list
    # instead of front. if we added to the front, we'd be outside of the Datadog
    # middleware and there'd be no current span.
    app.user_middleware.append(
        Middleware(BaseHTTPMiddleware, dispatch=gotta_catch_em_all)
    )
    app.middleware_stack = app.build_middleware_stack()

add_unhandled_exception_handler(app)

[amacfie]

With current Starlette (v0.26), remove app.middleware_stack = app.build_middleware_stack()