113 lines
3.6 KiB
Python
113 lines
3.6 KiB
Python
"""写接口审计上下文工具。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextvars import ContextVar, Token
|
|
from dataclasses import dataclass, replace
|
|
from datetime import datetime
|
|
from functools import wraps
|
|
from inspect import iscoroutinefunction
|
|
from typing import Any, Awaitable, Callable, Optional, TypeVar, cast
|
|
|
|
from .timezone import CN_TZ
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
VALID_CALLERS = {"system", "user", "ai"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AuditRequestContext:
|
|
trace_id: Optional[str] = None
|
|
interface_path: Optional[str] = None
|
|
caller: str = "system"
|
|
method: Optional[str] = None
|
|
tenant_id: Optional[int] = None
|
|
user_id: Optional[int] = None
|
|
|
|
|
|
_audit_context: ContextVar[AuditRequestContext] = ContextVar(
|
|
"pmflow_audit_request_context",
|
|
default=AuditRequestContext(),
|
|
)
|
|
|
|
|
|
def normalize_caller(value: Optional[str]) -> str:
|
|
normalized = str(value or "").strip().lower()
|
|
return normalized if normalized in VALID_CALLERS else "system"
|
|
|
|
|
|
def get_audit_context() -> AuditRequestContext:
|
|
return _audit_context.get()
|
|
|
|
|
|
def set_audit_context(context: AuditRequestContext) -> Token[AuditRequestContext]:
|
|
return _audit_context.set(
|
|
AuditRequestContext(
|
|
trace_id=(str(context.trace_id)[:64] if context.trace_id else None),
|
|
interface_path=(str(context.interface_path)[:255] if context.interface_path else None),
|
|
caller=normalize_caller(context.caller),
|
|
method=(str(context.method)[:16] if context.method else None),
|
|
tenant_id=int(context.tenant_id) if context.tenant_id is not None else None,
|
|
user_id=int(context.user_id) if context.user_id is not None else None,
|
|
)
|
|
)
|
|
|
|
|
|
def reset_audit_context(token: Token[AuditRequestContext]) -> None:
|
|
_audit_context.reset(token)
|
|
|
|
|
|
def update_audit_context(
|
|
*,
|
|
trace_id: Optional[str] = None,
|
|
interface_path: Optional[str] = None,
|
|
caller: Optional[str] = None,
|
|
method: Optional[str] = None,
|
|
tenant_id: Optional[int] = None,
|
|
user_id: Optional[int] = None,
|
|
) -> AuditRequestContext:
|
|
current = get_audit_context()
|
|
next_ctx = replace(
|
|
current,
|
|
trace_id=current.trace_id if trace_id is None else trace_id,
|
|
interface_path=current.interface_path if interface_path is None else interface_path,
|
|
caller=current.caller if caller is None else normalize_caller(caller),
|
|
method=current.method if method is None else method,
|
|
tenant_id=current.tenant_id if tenant_id is None else int(tenant_id),
|
|
user_id=current.user_id if user_id is None else int(user_id),
|
|
)
|
|
_audit_context.set(next_ctx)
|
|
return next_ctx
|
|
|
|
|
|
def format_beijing_datetime(dt: Optional[datetime]) -> Optional[str]:
|
|
if dt is None:
|
|
return None
|
|
localized = dt.replace(tzinfo=CN_TZ) if dt.tzinfo is None else dt.astimezone(CN_TZ)
|
|
return localized.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def audited_write(*, caller: Optional[str] = None) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
|
"""写接口装饰器：覆写当前请求 caller 审计上下文。"""
|
|
|
|
normalized = normalize_caller(caller)
|
|
|
|
def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
if not iscoroutinefunction(func):
|
|
return func
|
|
|
|
@wraps(func)
|
|
async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
current = get_audit_context()
|
|
token = set_audit_context(replace(current, caller=normalized))
|
|
try:
|
|
return await cast(Callable[..., Awaitable[Any]], func)(*args, **kwargs)
|
|
finally:
|
|
reset_audit_context(token)
|
|
|
|
return _async_wrapper
|
|
|
|
return _decorator
|