AutoScraperX/core/utils/trace_utils.py

import time
import uuid



# core/utils/trace_utils.py
import uuid
import contextvars

# 定义全局上下文变量，用于存储当前 trace_id
trace_id = contextvars.ContextVar("trace_id", default=None)

def generate_trace_id() -> str:
    """生成唯一 trace_id"""
    return f"{uuid.uuid4().hex[:8]}-{int(time.time() * 1000)}"

def get_current_trace_id() -> str:
    """获取当前上下文的 trace_id"""
    return trace_id.get()

def set_current_trace_id(trace_id: str) -> None:
    """设置当前上下文的 trace_id"""
    trace_id.set(trace_id)

class TraceContext:
    """上下文管理器：自动创建和清理 trace_id"""
    def __init__(self, trace_id: str = None):
        self.trace_id = trace_id or generate_trace_id()
        self.token = None  # 用于退出时恢复上下文

    def __enter__(self) -> str:
        # 保存当前上下文状态（token），并设置新的 trace_id
        self.token = trace_id.set(self.trace_id)
        return self.trace_id

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 退出时恢复之前的上下文（避免影响其他协程/线程）
        if self.token:
            trace_id.reset(self.token)