44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
from __future__ import annotations
|
||
|
||
from contextvars import ContextVar, Token
|
||
|
||
|
||
_job_id_var: ContextVar[str | None] = ContextVar("connecthub_job_id", default=None)
|
||
_job_log_id_var: ContextVar[int | None] = ContextVar("connecthub_job_log_id", default=None)
|
||
|
||
|
||
JobContextTokens = tuple[Token[str | None], Token[int | None]]
|
||
|
||
|
||
def set_job_context(*, job_id: str | None, job_log_id: int | None) -> JobContextTokens:
|
||
"""
|
||
设置当前执行上下文（用于日志隔离）。
|
||
返回 tokens，便于在 finally 中 reset 回原值。
|
||
"""
|
||
t1 = _job_id_var.set(job_id)
|
||
t2 = _job_log_id_var.set(job_log_id)
|
||
return (t1, t2)
|
||
|
||
|
||
def clear_job_context(tokens: JobContextTokens | None = None) -> None:
|
||
"""
|
||
清理当前执行上下文。
|
||
- 若提供 tokens：reset 回进入上下文前的值（推荐）
|
||
- 否则：直接置空
|
||
"""
|
||
if tokens is not None:
|
||
_job_id_var.reset(tokens[0])
|
||
_job_log_id_var.reset(tokens[1])
|
||
return
|
||
_job_id_var.set(None)
|
||
_job_log_id_var.set(None)
|
||
|
||
|
||
def get_job_id() -> str | None:
|
||
return _job_id_var.get()
|
||
|
||
|
||
def get_job_log_id() -> int | None:
|
||
return _job_log_id_var.get()
|
||
|