File size: 1,310 Bytes
b4f2238 | 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 | """Module for request context."""
from __future__ import annotations
import contextvars
import uuid
request_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("request_id", default=None)
trace_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("trace_id", default=None)
span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("span_id", default=None)
def generate_request_id() -> str:
return str(uuid.uuid4())[:8]
def set_request_context(
*,
request_id: str | None = None,
trace_id: str | None = None,
span_id: str | None = None,
) -> None:
if request_id is not None:
request_id_var.set(request_id)
if trace_id is not None:
trace_id_var.set(trace_id)
if span_id is not None:
span_id_var.set(span_id)
def clear_request_context() -> None:
request_id_var.set(None)
trace_id_var.set(None)
span_id_var.set(None)
def get_request_id(*, fallback_generate: bool = True) -> str | None:
value = request_id_var.get()
if value:
return value
if not fallback_generate:
return None
value = generate_request_id()
request_id_var.set(value)
return value
def get_trace_ids() -> tuple[str | None, str | None]:
return trace_id_var.get(), span_id_var.get()
|