document-management/backend/dms/common/request_context.py

"""DMS请求ID上下文。"""

from __future__ import annotations

from uuid import UUID, uuid4

from flask import g, request

from dms.common.errors import InvalidArgumentError


REQUEST_ID_HEADER = "X-Request-Id"


def initialize_request_id() -> None:
    """接收合法UUID，未提供时生成新的UUID。"""
    incoming = request.headers.get(REQUEST_ID_HEADER)
    if incoming:
        try:
            UUID(incoming)
        except (ValueError, AttributeError) as exc:
            g.dms_request_id = str(uuid4())
            raise InvalidArgumentError("X-Request-Id必须是有效UUID") from exc
        g.dms_request_id = incoming
        return
    g.dms_request_id = str(uuid4())


def get_request_id() -> str:
    request_id = getattr(g, "dms_request_id", None)
    if request_id is None:
        request_id = str(uuid4())
        g.dms_request_id = request_id
    return request_id