From f8d61739c8a5c1398f78037282e93a55408cc211 Mon Sep 17 00:00:00 2001 From: Daniel Ebrahimian Date: Wed, 29 Nov 2023 12:57:27 +1100 Subject: [PATCH 1/2] Extract xray context span id when it exists --- datadog_lambda/tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 252a4717..7973d91c 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -89,7 +89,7 @@ def _convert_xray_sampling(xray_sampled): return SamplingPriority.USER_KEEP if xray_sampled else SamplingPriority.USER_REJECT -def _get_xray_trace_context(): +def _get_xray_trace_context() -> Optional[Context]: if not is_lambda_context(): return None @@ -604,7 +604,7 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces): ) if merge_xray_traces: xray_context = _get_xray_trace_context() - if xray_context.span_id: + if xray_context and xray_context.span_id: context.span_id = xray_context.span_id tracer.context_provider.activate(context) From 354a61bd3ec5572fa52da053f1bca4015ec10ef0 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Wed, 29 Nov 2023 12:32:24 -0800 Subject: [PATCH 2/2] Add test for no Parent in _X_AMZN_TRACE_ID. --- datadog_lambda/tracing.py | 2 +- tests/test_tracing.py | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 7973d91c..6adf5d00 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -89,7 +89,7 @@ def _convert_xray_sampling(xray_sampled): return SamplingPriority.USER_KEEP if xray_sampled else SamplingPriority.USER_REJECT -def _get_xray_trace_context() -> Optional[Context]: +def _get_xray_trace_context(): if not is_lambda_context(): return None diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 745bf5d1..3a28a2a3 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -14,6 +14,7 @@ from datadog_lambda.constants import ( SamplingPriority, TraceHeader, + TraceContextSource, XraySubsegment, ) from datadog_lambda.tracing import ( @@ -854,6 +855,30 @@ def test_mixed_parent_context_when_merging(self): self.mock_activate.assert_called() self.mock_activate.assert_has_calls([call(expected_context)]) + def test_set_dd_trace_py_root_no_span_id(self): + os.environ["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320" + + lambda_ctx = get_mock_context() + ctx, source, event_type = extract_dd_trace_context( + { + "headers": { + TraceHeader.TRACE_ID: "123", + TraceHeader.PARENT_ID: "321", + TraceHeader.SAMPLING_PRIORITY: "1", + } + }, + lambda_ctx, + ) + set_dd_trace_py_root(TraceContextSource.EVENT, True) + + expected_context = Context( + trace_id=123, # Trace Id from incomming context + span_id=321, # Span Id from incoming context + sampling_priority=1, # Sampling priority from incomming context + ) + self.mock_activate.assert_called() + self.mock_activate.assert_has_calls([call(expected_context)]) + class TestAuthorizerInferredSpans(unittest.TestCase): def setUp(self):