upload

backend/app/utils/log_context.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import Any
+
+
+_trace_id_ctx: ContextVar[str | None] = ContextVar("trace_id", default=None)
+_path_ctx: ContextVar[str | None] = ContextVar("path", default=None)
+_method_ctx: ContextVar[str | None] = ContextVar("method", default=None)
+_user_id_ctx: ContextVar[str | None] = ContextVar("user_id", default=None)
+
+
+def set_request_context(*, trace_id: str | None, path: str | None, method: str | None) -> None:
+    _trace_id_ctx.set(trace_id)
+    _path_ctx.set(path)
+    _method_ctx.set(method)
+
+
+def set_user_context(*, user_id: str | None) -> None:
+    _user_id_ctx.set(user_id)
+
+
+def clear_log_context() -> None:
+    _trace_id_ctx.set(None)
+    _path_ctx.set(None)
+    _method_ctx.set(None)
+    _user_id_ctx.set(None)
+
+
+def get_log_context() -> dict[str, Any]:
+    payload: dict[str, Any] = {}
+
+    trace_id = _trace_id_ctx.get()
+    if trace_id:
+        payload["trace_id"] = trace_id
+
+    path = _path_ctx.get()
+    if path:
+        payload["path"] = path
+
+    method = _method_ctx.get()
+    if method:
+        payload["method"] = method
+
+    user_id = _user_id_ctx.get()
+    if user_id:
+        payload["user_id"] = user_id
+
+    return payload