80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
"""链路追踪模块"""
|
|
import uuid
|
|
from contextvars import ContextVar
|
|
from typing import Optional
|
|
|
|
# 上下文变量存储trace_id
|
|
_trace_id_var: ContextVar[Optional[str]] = ContextVar("trace_id", default=None)
|
|
_tenant_id_var: ContextVar[Optional[int]] = ContextVar("tenant_id", default=None)
|
|
_user_id_var: ContextVar[Optional[int]] = ContextVar("user_id", default=None)
|
|
|
|
|
|
def generate_trace_id() -> str:
|
|
"""生成新的trace_id"""
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
def get_trace_id() -> Optional[str]:
|
|
"""获取当前trace_id"""
|
|
return _trace_id_var.get()
|
|
|
|
|
|
def set_trace_id(trace_id: str) -> None:
|
|
"""设置当前trace_id"""
|
|
_trace_id_var.set(trace_id)
|
|
|
|
|
|
def get_tenant_id() -> Optional[int]:
|
|
"""获取当前租户ID"""
|
|
return _tenant_id_var.get()
|
|
|
|
|
|
def set_tenant_id(tenant_id: int) -> None:
|
|
"""设置当前租户ID"""
|
|
_tenant_id_var.set(tenant_id)
|
|
|
|
|
|
def get_user_id() -> Optional[int]:
|
|
"""获取当前用户ID"""
|
|
return _user_id_var.get()
|
|
|
|
|
|
def set_user_id(user_id: int) -> None:
|
|
"""设置当前用户ID"""
|
|
_user_id_var.set(user_id)
|
|
|
|
|
|
class TraceContext:
|
|
"""链路追踪上下文管理器
|
|
|
|
使用示例:
|
|
with TraceContext(tenant_id=1, user_id=100) as ctx:
|
|
print(ctx.trace_id)
|
|
# 在此上下文中的所有操作都会使用相同的trace_id
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
trace_id: Optional[str] = None,
|
|
tenant_id: Optional[int] = None,
|
|
user_id: Optional[int] = None
|
|
):
|
|
self.trace_id = trace_id or generate_trace_id()
|
|
self.tenant_id = tenant_id
|
|
self.user_id = user_id
|
|
self._tokens = []
|
|
|
|
def __enter__(self):
|
|
self._tokens.append(_trace_id_var.set(self.trace_id))
|
|
if self.tenant_id is not None:
|
|
self._tokens.append(_tenant_id_var.set(self.tenant_id))
|
|
if self.user_id is not None:
|
|
self._tokens.append(_user_id_var.set(self.user_id))
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
for token in self._tokens:
|
|
# 重置为之前的值
|
|
pass
|
|
return False
|