From 1d9d09f25f0504b137f31d373bcb48ef6c4c532e Mon Sep 17 00:00:00 2001 From: Carl Dunham Date: Wed, 3 Jan 2024 14:01:52 -0800 Subject: [PATCH] allow for empty X-Parent (#820) The `X-Parent` header could be empty, and will be for the root span. This PR allows for that. --- baseplate/__init__.py | 7 ++----- baseplate/frameworks/pyramid/__init__.py | 2 +- baseplate/frameworks/thrift/__init__.py | 2 +- tests/unit/core_tests.py | 12 ++++++++++++ 4 files changed, 16 insertions(+), 7 deletions(-) diff --git a/baseplate/__init__.py b/baseplate/__init__.py index f9275600c..b12977ae2 100644 --- a/baseplate/__init__.py +++ b/baseplate/__init__.py @@ -130,9 +130,9 @@ def new(cls) -> "TraceInfo": @classmethod def from_upstream( cls, - trace_id: Optional[str], + trace_id: str, parent_id: Optional[str], - span_id: Optional[str], + span_id: str, sampled: Optional[bool], flags: Optional[int], ) -> "TraceInfo": @@ -153,9 +153,6 @@ def from_upstream( if span_id is None: raise ValueError("invalid span_id") - if parent_id is None: - raise ValueError("invalid parent_id") - if sampled is not None and not isinstance(sampled, bool): raise ValueError("invalid sampled value") diff --git a/baseplate/frameworks/pyramid/__init__.py b/baseplate/frameworks/pyramid/__init__.py index 842931d4f..20070d0be 100644 --- a/baseplate/frameworks/pyramid/__init__.py +++ b/baseplate/frameworks/pyramid/__init__.py @@ -402,7 +402,7 @@ def _get_trace_info(self, headers: Mapping[str, str]) -> TraceInfo: flags = headers.get("X-Flags", None) return TraceInfo.from_upstream( headers["X-Trace"], - headers["X-Parent"], + headers.get("X-Parent", None), headers["X-Span"], sampled, int(flags) if flags is not None else None, diff --git a/baseplate/frameworks/thrift/__init__.py b/baseplate/frameworks/thrift/__init__.py index 6e878926d..bd67a416c 100644 --- a/baseplate/frameworks/thrift/__init__.py +++ b/baseplate/frameworks/thrift/__init__.py @@ -205,7 +205,7 @@ def call_processor_with_span_context( flags = headers.get(b"Flags", None) trace_info = TraceInfo.from_upstream( headers[b"Trace"].decode(), - headers[b"Parent"].decode(), + headers.get(b"Parent", b"").decode(), headers[b"Span"].decode(), sampled, int(flags) if flags is not None else None, diff --git a/tests/unit/core_tests.py b/tests/unit/core_tests.py index 57e3e56a4..6ee2ade1c 100644 --- a/tests/unit/core_tests.py +++ b/tests/unit/core_tests.py @@ -319,6 +319,11 @@ def test_new_does_not_have_parent_id(self): new_trace_info = TraceInfo.new() self.assertIsNone(new_trace_info.parent_id) + def test_new_sets_trace_and_span_ids(self): + new_trace_info = TraceInfo.new() + self.assertIsNotNone(new_trace_info.trace_id) + self.assertIsNotNone(new_trace_info.span_id) + def test_new_does_not_set_flags(self): new_trace_info = TraceInfo.new() self.assertIsNone(new_trace_info.flags) @@ -341,3 +346,10 @@ def test_from_upstream_handles_no_sampled_or_flags(self): span = TraceInfo.from_upstream(1, 2, 3, None, None) self.assertIsNone(span.sampled) self.assertIsNone(span.flags) + + def test_from_upstream_handles_no_parent(self): + """Verify that `parent_id` can be left out without raising an exception.""" + span = TraceInfo.from_upstream(1, None, 3, None, None) + self.assertIsNone(span.parent_id) + self.assertIsNotNone(span.trace_id) + self.assertIsNotNone(span.span_id)