starlette_context_middleware.py

import logging
from typing import Any, Optional, Union
from uuid import uuid4

from starlette.requests import HTTPConnection, Request
from starlette.responses import Response
from starlette.types import Message
from starlette_context import context
from starlette_context.plugins import Plugin

# middleware that injects a X-Request-ID or X-Correlation-ID, depending if it is missing or not

# https://starlette-context.readthedocs.io/en/latest/plugins.html#implementing-your-own
class ScottyXID(Plugin):

    def __init__(self, header_key):
        super().__init__()
        self.key = header_key

    async def process_request(self, request: Union[Request, HTTPConnection]) -> Optional[Any]:
        # access any part of the request
        if corr_id := request.headers.get(self.key):
            return corr_id
        else:
            new_uuid = uuid4().hex
            logging.warning(f"No {self.key}. Generated a new one: {new_uuid}")
            return new_uuid

    async def enrich_response(self, response: Union[Response, Message]) -> None:
        # can access the populated context here.
        response.headers.append(self.key, context.get(self.key))
        # mutate the response in-place, return nothing.