File size: 2,586 Bytes
f8b48da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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
# A ContextVar is the right tool: it's safe across threads and async tasks
# and is automatically scoped per asyncio task.
_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",
]
|