diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 81118848..15dec7f1 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -281,7 +281,12 @@ def _before(self, event, context): # Create a Datadog X-Ray subsegment with the trace context if dd_context and trace_context_source == TraceContextSource.EVENT: create_dd_dummy_metadata_subsegment( - dd_context, XraySubsegment.TRACE_KEY + { + "trace-id": str(dd_context.trace_id), + "parent-id": str(dd_context.span_id), + "sampling-priority": str(dd_context.sampling_priority), + }, + XraySubsegment.TRACE_KEY, ) if dd_tracing_enabled: diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 57c5de45..97547a26 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -9,7 +9,7 @@ import datadog_lambda.wrapper as wrapper from datadog_lambda.metric import lambda_metric from datadog_lambda.thread_stats_writer import ThreadStatsWriter -from ddtrace import Span +from ddtrace import Span, tracer def get_mock_context( @@ -580,3 +580,70 @@ def test_some_envs_should_depend_on_dd_tracing_enabled(self): self.assertFalse(decorator.make_inferred_span) self.assertFalse(decorator.encode_authorizer_context) self.assertFalse(decorator.decode_authorizer_context) + + +class TestLambdaWrapperWithTraceContext(unittest.TestCase): + xray_root = "1-5e272390-8c398be037738dc042009320" + xray_parent = "94ae789b969f1cc5" + xray_daemon_envvar = "localhost:1234" + xray_trace_envvar = ( + f"Root={xray_root};Parent={xray_parent};Sampled=1;Lineage=c6c5b1b9:0" + ) + + @patch( + "os.environ", + { + "AWS_XRAY_DAEMON_ADDRESS": xray_daemon_envvar, + "_X_AMZN_TRACE_ID": xray_trace_envvar, + }, + ) + def test_event_bridge_sqs_payload(self): + patcher = patch("datadog_lambda.xray.send") + mock_send = patcher.start() + self.addCleanup(patcher.stop) + + def handler(event, context): + return tracer.current_trace_context() + + wrapper.dd_tracing_enabled = True + wrapped_handler = wrapper.datadog_lambda_wrapper(handler) + + event_trace_id = 3047453991382739997 + event_parent_id = 3047453991382739997 + event = { + "headers": { + "traceparent": ( + f"00-0000000000000000{hex(event_trace_id)[2:]}-{hex(event_parent_id)[2:]}-01" + ), + "tracestate": "dd=s:1;t.dm:-1", + "x-datadog-trace-id": str(event_trace_id), + "x-datadog-parent-id": str(event_parent_id), + "x-datadog-sampling-priority": "1", + }, + } + context = get_mock_context() + + result = wrapped_handler(event, context) + aws_lambda_span = wrapped_handler.span + + self.assertIsNotNone(result) + self.assertEqual(result.trace_id, event_trace_id) + self.assertEqual(result.span_id, aws_lambda_span.span_id) + self.assertEqual(result.sampling_priority, 1) + mock_send.assert_called_once() + (_, raw_payload), _ = mock_send.call_args + payload = json.loads(raw_payload[33:]) # strip formatting prefix + self.assertEqual(self.xray_root, payload["trace_id"]) + self.assertEqual(self.xray_parent, payload["parent_id"]) + self.assertDictEqual( + { + "datadog": { + "trace": { + "trace-id": str(event_trace_id), + "parent-id": str(event_parent_id), + "sampling-priority": "1", + }, + }, + }, + payload["metadata"], + )