| """Thread-local request context for the current request_id.
|
|
|
| This module stores the current request's identifier in a
|
| :class:`contextvars.ContextVar` so any code running inside the request --
|
| including deeply nested service code, background tasks, and observability
|
| helpers -- can attach the request_id to log records and traces without the
|
| caller having to thread it through every function signature.
|
|
|
| The :class:`RequestContext` context manager sets the value on entry and
|
| restores the previous value on exit, so nested contexts (e.g. a span
|
| triggering an internal sub-request) work as expected.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import contextvars
|
| import uuid
|
|
|
|
|
|
|
| _request_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
| "request_id", default=None
|
| )
|
|
|
|
|
| class RequestContext:
|
| """Context manager that sets the current request_id.
|
|
|
| Example::
|
|
|
| with RequestContext() as request_id:
|
| logger.info("processing") # log record is auto-tagged
|
|
|
| # Or with an explicit id:
|
| with RequestContext("req-abc123"):
|
| ...
|
| """
|
|
|
| def __init__(self, request_id: str | None = None) -> None:
|
| self._explicit = request_id
|
| self._generated: str | None = None
|
| self._token: contextvars.Token[str | None] | None = None
|
|
|
| def __enter__(self) -> str:
|
| rid = self._explicit or str(uuid.uuid4())
|
| if self._explicit is None:
|
| self._generated = rid
|
| self._token = _request_id_var.set(rid)
|
| return rid
|
|
|
| def __exit__(self, *exc_info: object) -> None:
|
| if self._token is not None:
|
| _request_id_var.reset(self._token)
|
| self._token = None
|
|
|
|
|
| def get_request_id() -> str | None:
|
| """Return the current request_id, or None if no context is active."""
|
| return _request_id_var.get()
|
|
|
|
|
| def set_request_id(request_id: str) -> RequestContext:
|
| """Return a context manager that sets the current request_id for its block.
|
|
|
| Convenience wrapper for use with ``with``::
|
|
|
| with set_request_id("req-123"):
|
| do_work()
|
| """
|
| return RequestContext(request_id)
|
|
|
|
|
| def clear_request_id() -> None:
|
| """Reset the current request_id to None.
|
|
|
| Mainly useful in tests.
|
| """
|
| _request_id_var.set(None)
|
|
|
|
|
| __all__ = [
|
| "RequestContext",
|
| "clear_request_id",
|
| "get_request_id",
|
| "set_request_id",
|
| ]
|
|
|
