How can I implement a Correlation ID middleware?

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?

[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.

[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)