feat(trace): support external trace id propagation (#22623)
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>pull/21682/merge
qfl
7 months ago
committed by
GitHub
parent
c987001a19
commit
841e53dbbe
@ -0,0 +1,42 @@
|
|||||||
|
import re
|
||||||
|
from collections.abc import Mapping
|
||||||
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
|
||||||
|
def is_valid_trace_id(trace_id: str) -> bool:
|
||||||
|
"""
|
||||||
|
Check if the trace_id is valid.
|
||||||
|
|
||||||
|
Requirements: 1-128 characters, only letters, numbers, '-', and '_'.
|
||||||
|
"""
|
||||||
|
return bool(re.match(r"^[a-zA-Z0-9\-_]{1,128}$", trace_id))
|
||||||
|
|
||||||
|
|
||||||
|
def get_external_trace_id(request: Any) -> Optional[str]:
|
||||||
|
"""
|
||||||
|
Retrieve the trace_id from the request.
|
||||||
|
|
||||||
|
Priority: header ('X-Trace-Id'), then parameters, then JSON body. Returns None if not provided or invalid.
|
||||||
|
"""
|
||||||
|
trace_id = request.headers.get("X-Trace-Id")
|
||||||
|
if not trace_id:
|
||||||
|
trace_id = request.args.get("trace_id")
|
||||||
|
if not trace_id and getattr(request, "is_json", False):
|
||||||
|
json_data = getattr(request, "json", None)
|
||||||
|
if json_data:
|
||||||
|
trace_id = json_data.get("trace_id")
|
||||||
|
if isinstance(trace_id, str) and is_valid_trace_id(trace_id):
|
||||||
|
return trace_id
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict:
|
||||||
|
"""
|
||||||
|
Extract 'external_trace_id' from args.
|
||||||
|
|
||||||
|
Returns a dict suitable for use in extras. Returns an empty dict if not found.
|
||||||
|
"""
|
||||||
|
trace_id = args.get("external_trace_id")
|
||||||
|
if trace_id:
|
||||||
|
return {"external_trace_id": trace_id}
|
||||||
|
return {}
|
||||||
@ -0,0 +1,86 @@
|
|||||||
|
import pytest
|
||||||
|
|
||||||
|
from core.helper.trace_id_helper import extract_external_trace_id_from_args, get_external_trace_id, is_valid_trace_id
|
||||||
|
|
||||||
|
|
||||||
|
class DummyRequest:
|
||||||
|
def __init__(self, headers=None, args=None, json=None, is_json=False):
|
||||||
|
self.headers = headers or {}
|
||||||
|
self.args = args or {}
|
||||||
|
self.json = json
|
||||||
|
self.is_json = is_json
|
||||||
|
|
||||||
|
|
||||||
|
class TestTraceIdHelper:
|
||||||
|
"""Test cases for trace_id_helper.py"""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("trace_id", "expected"),
|
||||||
|
[
|
||||||
|
("abc123", True),
|
||||||
|
("A-B_C-123", True),
|
||||||
|
("a" * 128, True),
|
||||||
|
("", False),
|
||||||
|
("a" * 129, False),
|
||||||
|
("abc!@#", False),
|
||||||
|
("空格", False),
|
||||||
|
("with space", False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_is_valid_trace_id(self, trace_id, expected):
|
||||||
|
"""Test trace_id validation for various cases"""
|
||||||
|
assert is_valid_trace_id(trace_id) is expected
|
||||||
|
|
||||||
|
def test_get_external_trace_id_from_header(self):
|
||||||
|
"""Should extract valid trace_id from header"""
|
||||||
|
req = DummyRequest(headers={"X-Trace-Id": "abc123"})
|
||||||
|
assert get_external_trace_id(req) == "abc123"
|
||||||
|
|
||||||
|
def test_get_external_trace_id_from_args(self):
|
||||||
|
"""Should extract valid trace_id from args if header missing"""
|
||||||
|
req = DummyRequest(args={"trace_id": "abc123"})
|
||||||
|
assert get_external_trace_id(req) == "abc123"
|
||||||
|
|
||||||
|
def test_get_external_trace_id_from_json(self):
|
||||||
|
"""Should extract valid trace_id from JSON body if header and args missing"""
|
||||||
|
req = DummyRequest(is_json=True, json={"trace_id": "abc123"})
|
||||||
|
assert get_external_trace_id(req) == "abc123"
|
||||||
|
|
||||||
|
def test_get_external_trace_id_priority(self):
|
||||||
|
"""Header > args > json priority"""
|
||||||
|
req = DummyRequest(
|
||||||
|
headers={"X-Trace-Id": "header_id"},
|
||||||
|
args={"trace_id": "args_id"},
|
||||||
|
is_json=True,
|
||||||
|
json={"trace_id": "json_id"},
|
||||||
|
)
|
||||||
|
assert get_external_trace_id(req) == "header_id"
|
||||||
|
req2 = DummyRequest(args={"trace_id": "args_id"}, is_json=True, json={"trace_id": "json_id"})
|
||||||
|
assert get_external_trace_id(req2) == "args_id"
|
||||||
|
req3 = DummyRequest(is_json=True, json={"trace_id": "json_id"})
|
||||||
|
assert get_external_trace_id(req3) == "json_id"
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"req",
|
||||||
|
[
|
||||||
|
DummyRequest(headers={"X-Trace-Id": "!!!"}),
|
||||||
|
DummyRequest(args={"trace_id": "!!!"}),
|
||||||
|
DummyRequest(is_json=True, json={"trace_id": "!!!"}),
|
||||||
|
DummyRequest(),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_get_external_trace_id_invalid(self, req):
|
||||||
|
"""Should return None for invalid or missing trace_id"""
|
||||||
|
assert get_external_trace_id(req) is None
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("args", "expected"),
|
||||||
|
[
|
||||||
|
({"external_trace_id": "abc123"}, {"external_trace_id": "abc123"}),
|
||||||
|
({"other": "value"}, {}),
|
||||||
|
({}, {}),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_extract_external_trace_id_from_args(self, args, expected):
|
||||||
|
"""Test extraction of external_trace_id from args mapping"""
|
||||||
|
assert extract_external_trace_id_from_args(args) == expected
|
||||||
Loading…
Reference in New Issue