53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, replace
|
|
|
|
from app_runtime.contracts.trace import TraceContextRecord
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ActiveTraceContext:
|
|
record: TraceContextRecord
|
|
step: str = ""
|
|
|
|
|
|
class TraceContextStore:
|
|
def __init__(self) -> None:
|
|
self._current: ContextVar[ActiveTraceContext | None] = ContextVar("trace_current", default=None)
|
|
self._stack: ContextVar[tuple[ActiveTraceContext, ...]] = ContextVar("trace_stack", default=())
|
|
|
|
def current(self) -> ActiveTraceContext | None:
|
|
return self._current.get()
|
|
|
|
def current_trace_id(self) -> str | None:
|
|
active = self.current()
|
|
return active.record.trace_id if active else None
|
|
|
|
def push(self, record: TraceContextRecord) -> ActiveTraceContext:
|
|
active = self.current()
|
|
stack = self._stack.get()
|
|
if active is not None:
|
|
self._stack.set(stack + (active,))
|
|
updated = ActiveTraceContext(record=record)
|
|
self._current.set(updated)
|
|
return updated
|
|
|
|
def pop(self) -> ActiveTraceContext | None:
|
|
stack = self._stack.get()
|
|
if not stack:
|
|
self._current.set(None)
|
|
return None
|
|
previous = stack[-1]
|
|
self._stack.set(stack[:-1])
|
|
self._current.set(previous)
|
|
return previous
|
|
|
|
def set_step(self, step: str) -> ActiveTraceContext | None:
|
|
active = self.current()
|
|
if active is None:
|
|
return None
|
|
updated = replace(active, step=step)
|
|
self._current.set(updated)
|
|
return updated
|