Initial commit: 000-platform project skeleton

sdk/trace.py

@@ -0,0 +1,79 @@
+"""链路追踪模块"""
+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