diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index ae48a161..49e1d05f 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -129,16 +129,27 @@ def extract_context_from_lambda_context(lambda_context): dd_trace libraries inject this trace context on synchronous invocations """ client_context = lambda_context.client_context - + trace_id = None + parent_id = None + sampling_priority = None if client_context and client_context.custom: - dd_data = client_context.custom.get("_datadog", {}) - trace_id = dd_data.get(TraceHeader.TRACE_ID) - parent_id = dd_data.get(TraceHeader.PARENT_ID) - sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) + if "_datadog" in client_context.custom: + # Legacy trace propagation dict + dd_data = client_context.custom.get("_datadog", {}) + trace_id = dd_data.get(TraceHeader.TRACE_ID) + parent_id = dd_data.get(TraceHeader.PARENT_ID) + sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) + elif ( + TraceHeader.TRACE_ID in client_context.custom + and TraceHeader.PARENT_ID in client_context.custom + and TraceHeader.SAMPLING_PRIORITY in client_context.custom + ): + # New trace propagation keys + trace_id = client_context.custom.get(TraceHeader.TRACE_ID) + parent_id = client_context.custom.get(TraceHeader.PARENT_ID) + sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY) - return trace_id, parent_id, sampling_priority - - return None, None, None + return trace_id, parent_id, sampling_priority def extract_context_from_http_event_or_context(event, lambda_context): diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 431f339f..64b42e18 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self): XraySubsegment.NAMESPACE, ) - def test_with_client_context_datadog_trace_data(self): + def test_with_legacy_client_context_datadog_trace_data(self): lambda_ctx = get_mock_context( custom={ "_datadog": { @@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self): XraySubsegment.NAMESPACE, ) + def test_with_new_client_context_datadog_trace_data(self): + lambda_ctx = get_mock_context( + custom={ + TraceHeader.TRACE_ID: "666", + TraceHeader.PARENT_ID: "777", + TraceHeader.SAMPLING_PRIORITY: "1", + } + ) + ctx, source = extract_dd_trace_context({}, lambda_ctx) + self.assertEqual(source, "event") + self.assertDictEqual( + ctx, + { + "trace-id": "666", + "parent-id": "777", + "sampling-priority": "1", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "666", + TraceHeader.PARENT_ID: "65535", + TraceHeader.SAMPLING_PRIORITY: "1", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_xray_recorder.begin_subsegment.assert_called() + self.mock_xray_recorder.end_subsegment.assert_called() + self.mock_current_subsegment.put_metadata.assert_called_with( + XraySubsegment.TRACE_KEY, + {"trace-id": "666", "parent-id": "777", "sampling-priority": "1"}, + XraySubsegment.NAMESPACE, + ) + def test_with_complete_datadog_trace_headers_with_mixed_casing(self): lambda_ctx = get_mock_context() extract_dd_trace_context(