How can I implement a Correlation ID middleware?

[mths0x5f]

It's common practice in my company to trace log messages related to a unique request with a common UUID between messages. How could I implement these?

I have had a look in Starlette middleware docs but I have no idea how I can add an extra field for logging that has same value for log lines in the same request without some kind of global object.

Has someone done something along these lines in Python?

[dmontagu]

Looks like this might be related?

encode/starlette#420

[mths0x5f]

This was indeed related and I just implemented the way was expecting to! Thank you, @dmontagu

middlewares.py

from contextvars import ContextVar
from uuid import uuid4

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request

CORRELATION_ID_CTX_KEY = 'correlation_id'
REQUEST_ID_CTX_KEY = 'request_id'

_correlation_id_ctx_var: ContextVar[str] = ContextVar(CORRELATION_ID_CTX_KEY, default=None)
_request_id_ctx_var: ContextVar[str] = ContextVar(REQUEST_ID_CTX_KEY, default=None)


def get_correlation_id() -> str:
    return _correlation_id_ctx_var.get()


def get_request_id() -> str:
    return _request_id_ctx_var.get()


class RequestContextLogMiddleware(BaseHTTPMiddleware):

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        correlation_id = _correlation_id_ctx_var.set(request.headers.get('X-Correlation-ID', str(uuid4())))
        request_id = _request_id_ctx_var.set(str(uuid4()))

        response = await call_next(request)
        response.headers['X-Correlation-ID'] = get_correlation_id()
        response.headers['X-Request-ID'] = get_request_id()

        _correlation_id_ctx_var.reset(correlation_id)
        _request_id_ctx_var.reset(request_id)

        return response

logging.py

import logging

from app.middlewares import get_request_id, get_correlation_id
from app.settings import DEBUG


class AppFilter(logging.Filter):
    def filter(self, record):
        record.correlation_id = get_correlation_id()
        record.request_id = get_request_id()
        return True


def setup_logging():
    logger = logging.getLogger()
    syslog = logging.StreamHandler()
    syslog.addFilter(AppFilter())

    formatter = logging.Formatter('%(asctime)s %(process)s %(request_id)s %(correlation_id)s '
                                  '%(levelname)s %(name)s %(message)s')

    syslog.setFormatter(formatter)
    logger.setLevel(logging.DEBUG if DEBUG else logging.WARN)
    logger.addHandler(syslog)

[cetanu]

Sorry for reviving an old issue, but it seems like this middleware does not run when the application encounters an exception, is that expected? If so, what's the canonical way to ensure the middleware runs even before/after an exception handler?

Minimum code to demonstrate:

import uvicorn
from fastapi import FastAPI
from starlette.responses import JSONResponse
from contextvars import ContextVar
from uuid import uuid4

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request

CORRELATION_ID_CTX_KEY = 'correlation_id'
REQUEST_ID_CTX_KEY = 'request_id'

_correlation_id_ctx_var: ContextVar[str] = ContextVar(CORRELATION_ID_CTX_KEY, default=None)
_request_id_ctx_var: ContextVar[str] = ContextVar(REQUEST_ID_CTX_KEY, default=None)


def get_correlation_id() -> str:
    return _correlation_id_ctx_var.get()


def get_request_id() -> str:
    return _request_id_ctx_var.get()


class RequestContextLogMiddleware(BaseHTTPMiddleware):

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        correlation_id = _correlation_id_ctx_var.set(request.headers.get('X-Correlation-ID', str(uuid4())))
        request_id = _request_id_ctx_var.set(str(uuid4()))

        response = await call_next(request)
        response.headers['X-Correlation-ID'] = get_correlation_id()
        response.headers['X-Request-ID'] = get_request_id()

        _correlation_id_ctx_var.reset(correlation_id)
        _request_id_ctx_var.reset(request_id)

        return response


app = FastAPI(
    title='HelloWorld',
    version='0.1.0'
)
app.add_middleware(RequestContextLogMiddleware)

@app.get('/')
def index():
    return 'Hello!'

@app.get('/raise')
def raise_exception():
    raise ValueError('Something random')

@app.exception_handler(500)
async def exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        content={'error': exc.__class__.__name__},
        status_code=500
    )

uvicorn.run(app)

[dmontagu]

@cetanu

I'm not sure I fully understand your intention, but I think it might solve things if you put this line in a try: except: or try: finally: block:

response = await call_next(request)

so it is like:

try:
    response = await call_next(request)
except Exception:
    # do whatever you want to do *only* when there is an exception
    response = Response(...)

or similar.

If you are trying to add headers in the case of an exception, you'll basically need to write an exception handler.

[tomwojcik]

@mths0x5f @cetanu I needed something like this so I made a package.

https://github.com/tomwojcik/starlette-context

I basically packaged what you wrote in here + tested it myself in my small project. Seems to be working. All feedback welcome.

[tiangolo]

I guess this section in the docs is probably useful to adding that custom logic that reads the request data: https://fastapi.tiangolo.com/advanced/custom-request-and-route/#accessing-the-request-body-in-an-exception-handler

[github-actions]

Assuming the original issue was solved, it will be automatically closed now. But feel free to add more comments or create new issues.