Skip to content

Commit 5c01c2e

Browse files
authored
Check for the new direct invocation trace format and use that (#146)
* Check for the new direct invocation trace format and use that
1 parent 7fcf3fa commit 5c01c2e

File tree

2 files changed

+55
-9
lines changed

2 files changed

+55
-9
lines changed

datadog_lambda/tracing.py

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -129,16 +129,27 @@ def extract_context_from_lambda_context(lambda_context):
129129
dd_trace libraries inject this trace context on synchronous invocations
130130
"""
131131
client_context = lambda_context.client_context
132-
132+
trace_id = None
133+
parent_id = None
134+
sampling_priority = None
133135
if client_context and client_context.custom:
134-
dd_data = client_context.custom.get("_datadog", {})
135-
trace_id = dd_data.get(TraceHeader.TRACE_ID)
136-
parent_id = dd_data.get(TraceHeader.PARENT_ID)
137-
sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
136+
if "_datadog" in client_context.custom:
137+
# Legacy trace propagation dict
138+
dd_data = client_context.custom.get("_datadog", {})
139+
trace_id = dd_data.get(TraceHeader.TRACE_ID)
140+
parent_id = dd_data.get(TraceHeader.PARENT_ID)
141+
sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
142+
elif (
143+
TraceHeader.TRACE_ID in client_context.custom
144+
and TraceHeader.PARENT_ID in client_context.custom
145+
and TraceHeader.SAMPLING_PRIORITY in client_context.custom
146+
):
147+
# New trace propagation keys
148+
trace_id = client_context.custom.get(TraceHeader.TRACE_ID)
149+
parent_id = client_context.custom.get(TraceHeader.PARENT_ID)
150+
sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY)
138151

139-
return trace_id, parent_id, sampling_priority
140-
141-
return None, None, None
152+
return trace_id, parent_id, sampling_priority
142153

143154

144155
def extract_context_from_http_event_or_context(event, lambda_context):

tests/test_tracing.py

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self):
272272
XraySubsegment.NAMESPACE,
273273
)
274274

275-
def test_with_client_context_datadog_trace_data(self):
275+
def test_with_legacy_client_context_datadog_trace_data(self):
276276
lambda_ctx = get_mock_context(
277277
custom={
278278
"_datadog": {
@@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self):
309309
XraySubsegment.NAMESPACE,
310310
)
311311

312+
def test_with_new_client_context_datadog_trace_data(self):
313+
lambda_ctx = get_mock_context(
314+
custom={
315+
TraceHeader.TRACE_ID: "666",
316+
TraceHeader.PARENT_ID: "777",
317+
TraceHeader.SAMPLING_PRIORITY: "1",
318+
}
319+
)
320+
ctx, source = extract_dd_trace_context({}, lambda_ctx)
321+
self.assertEqual(source, "event")
322+
self.assertDictEqual(
323+
ctx,
324+
{
325+
"trace-id": "666",
326+
"parent-id": "777",
327+
"sampling-priority": "1",
328+
},
329+
)
330+
self.assertDictEqual(
331+
get_dd_trace_context(),
332+
{
333+
TraceHeader.TRACE_ID: "666",
334+
TraceHeader.PARENT_ID: "65535",
335+
TraceHeader.SAMPLING_PRIORITY: "1",
336+
},
337+
)
338+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
339+
self.mock_xray_recorder.begin_subsegment.assert_called()
340+
self.mock_xray_recorder.end_subsegment.assert_called()
341+
self.mock_current_subsegment.put_metadata.assert_called_with(
342+
XraySubsegment.TRACE_KEY,
343+
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
344+
XraySubsegment.NAMESPACE,
345+
)
346+
312347
def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
313348
lambda_ctx = get_mock_context()
314349
extract_dd_trace_context(

0 commit comments

Comments
 (0)