Skip to content

Commit 25780a8

Browse files
avedmalaavedmala
committed
use default propagator.extract
1 parent 4e0afdb commit 25780a8

File tree

3 files changed

+30
-23
lines changed

3 files changed

+30
-23
lines changed

‎datadog_lambda/tracing.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -383,22 +383,27 @@ def extract_context_from_step_functions(event, lambda_context):
383383
state_entered_time = event.get("State").get("EnteredTime")
384384
meta = {}
385385

386+
trace_id = None
386387
if "_datadog" in event:
387388
trace_header = event.get("_datadog")
388-
if TraceHeader.TRACE_ID in trace_header:
389-
# use the trace ID from the top-most parent when it exists
390-
trace_id = int(trace_header.get(TraceHeader.TRACE_ID))
389+
explicit_context = propagator.extract(trace_header)
390+
if _is_context_complete(explicit_context):
391391
tags = trace_header.get(TraceHeader.TAGS, "")
392392
for tag in tags.split(","):
393393
tag_key, tag_val = tag.split("=")
394394
meta[tag_key] = tag_val
395-
elif "x-datadog-execution-arn" in trace_header:
396-
root_execution_id = trace_header.get("x-datadog-execution-arn")
395+
explicit_context.span_id = _deterministic_sha256_hash(
396+
f"{execution_id}#{state_name}#{state_entered_time}", HIGHER_64_BITS
397+
)
398+
return explicit_context
399+
if "x-datadog-execution-arn" in trace_header:
400+
root_execution_id = trace_header.get("x-datadog-root-execution-arn")
397401
trace_id = _deterministic_sha256_hash(root_execution_id, LOWER_64_BITS)
398402
meta["_dd.p.tid"] = hex(
399403
_deterministic_sha256_hash(root_execution_id, HIGHER_64_BITS)
400404
)[2:]
401-
else:
405+
406+
if not trace_id:
402407
# returning 128 bits since 128bit traceId will be break up into
403408
# traditional traceId and _dd.p.tid tag
404409
# https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240

‎tests/test_tracing.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -664,26 +664,28 @@ def test_step_function_trace_data_with_trace_header(self):
664664
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
665665
},
666666
"_datadog": {
667-
"x-datadog-trace-id": "4061173386180447114",
668-
"x-datadog-tags": "_dd.p.tid=aac3639aa724716",
667+
"x-datadog-trace-id": "11742251842529032210",
668+
"x-datadog-parent-id": "13977111940858727778",
669+
"x-datadog-sampling-priority": "1",
670+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6728f8ec00000000"
669671
},
670672
}
671673
ctx, source, event_source = extract_dd_trace_context(sf_event, lambda_ctx)
672674
self.assertEqual(source, "event")
673675
expected_context = Context(
674-
trace_id=4061173386180447114,
676+
trace_id=137123224175615787006624409899264538642,
675677
span_id=6880978411788117524,
676678
sampling_priority=1,
677-
meta={"_dd.p.tid": "aac3639aa724716"},
679+
meta={'_dd.p.dm': '-0', "_dd.p.tid": "6728f8ec00000000"},
678680
)
679681
self.assertEqual(ctx, expected_context)
680682
self.assertEqual(
681683
get_dd_trace_context(),
682684
{
683-
TraceHeader.TRACE_ID: "4061173386180447114",
685+
TraceHeader.TRACE_ID: "11742251842529032210",
684686
TraceHeader.PARENT_ID: "10713633173203262661",
685687
TraceHeader.SAMPLING_PRIORITY: "1",
686-
TraceHeader.TAGS: "_dd.p.tid=aac3639aa724716",
688+
TraceHeader.TAGS: "_dd.p.dm=-0,_dd.p.tid=6728f8ec00000000",
687689
},
688690
)
689691
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
@@ -705,25 +707,25 @@ def test_step_function_trace_data_with_arn_header(self):
705707
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
706708
},
707709
"_datadog": {
708-
"x-datadog-execution-arn": "ca7383bc-e370-4a85-a266-a4686bd7d00f"
710+
"x-datadog-root-execution-arn": "ca7383bc-e370-4a85-a266-a4686bd7d00f"
709711
},
710712
}
711713
ctx, source, event_source = extract_dd_trace_context(sf_event, lambda_ctx)
712714
self.assertEqual(source, "event")
713715
expected_context = Context(
714-
trace_id=6970872619724504833,
716+
trace_id=3675572987363469717,
715717
span_id=6880978411788117524,
716718
sampling_priority=1,
717-
meta={"_dd.p.tid": "71dab8f4d4629263"},
719+
meta={"_dd.p.tid": "e987c84b36b11ab"},
718720
)
719721
self.assertEqual(ctx, expected_context)
720722
self.assertEqual(
721723
get_dd_trace_context(),
722724
{
723-
TraceHeader.TRACE_ID: "6970872619724504833",
725+
TraceHeader.TRACE_ID: "3675572987363469717",
724726
TraceHeader.PARENT_ID: "10713633173203262661",
725727
TraceHeader.SAMPLING_PRIORITY: "1",
726-
TraceHeader.TAGS: "_dd.p.tid=71dab8f4d4629263",
728+
TraceHeader.TAGS: "_dd.p.tid=e987c84b36b11ab",
727729
},
728730
)
729731
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)

‎tests/test_xray.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@ def test_get_xray_host_port_success(self):
3434

3535
def test_send_segment_sampled_out(self):
3636
os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
37-
os.environ[
38-
"_X_AMZN_TRACE_ID"
39-
] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0;Lineage=c6c5b1b9:0"
37+
os.environ["_X_AMZN_TRACE_ID"] = (
38+
"Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0;Lineage=c6c5b1b9:0"
39+
)
4040

4141
with patch(
4242
"datadog_lambda.xray.sock.send", MagicMock(return_value=None)
@@ -47,9 +47,9 @@ def test_send_segment_sampled_out(self):
4747

4848
def test_send_segment_sampled(self):
4949
os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
50-
os.environ[
51-
"_X_AMZN_TRACE_ID"
52-
] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
50+
os.environ["_X_AMZN_TRACE_ID"] = (
51+
"Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
52+
)
5353
with patch(
5454
"datadog_lambda.xray.sock.send", MagicMock(return_value=None)
5555
) as mock_send:

0 commit comments

Comments
 (0)