added root arn case

Author: avedmala

datadog_lambda/tracing.py

@@ -384,13 +384,20 @@ def extract_context_from_step_functions(event, lambda_context):
         meta = {}
 
         if "_datadog" in event:
-            # use the trace ID from the top-most parent when it exists
             trace_header = event.get("_datadog")
-            trace_id = int(trace_header.get(TraceHeader.TRACE_ID))
-            tags = trace_header.get(TraceHeader.TAGS)
-            for tag in tags.split(","):
-                tag_key, tag_val = tag.split("=")
-                meta[tag_key] = tag_val
+            if TraceHeader.TRACE_ID in trace_header:
+                # use the trace ID from the top-most parent when it exists
+                trace_id = int(trace_header.get(TraceHeader.TRACE_ID))
+                tags = trace_header.get(TraceHeader.TAGS, "")
+                for tag in tags.split(","):
+                    tag_key, tag_val = tag.split("=")
+                    meta[tag_key] = tag_val
+            elif "x-datadog-execution-arn" in trace_header:
+                root_execution_id = trace_header.get("x-datadog-execution-arn")
+                trace_id = _deterministic_sha256_hash(root_execution_id, LOWER_64_BITS)
+                meta["_dd.p.tid"] = hex(
+                    _deterministic_sha256_hash(root_execution_id, HIGHER_64_BITS)
+                )[2:]
         else:
             # returning 128 bits since 128bit traceId will be break up into
             # traditional traceId and _dd.p.tid tag

tests/test_tracing.py

@@ -652,7 +652,7 @@ def test_step_function_trace_data(self):
         )
 
     @with_trace_propagation_style("datadog")
-    def test_step_function_trace_data_with_header(self):
+    def test_step_function_trace_data_with_trace_header(self):
         lambda_ctx = get_mock_context()
         sf_event = {
             "Execution": {
@@ -692,6 +692,46 @@ def test_step_function_trace_data_with_header(self):
             expected_context,
         )
 
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_with_arn_header(self):
+        lambda_ctx = get_mock_context()
+        sf_event = {
+            "Execution": {
+                "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+            },
+            "StateMachine": {},
+            "State": {
+                "Name": "my-awesome-state",
+                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+            },
+            "_datadog": {
+                "x-datadog-execution-arn": "ca7383bc-e370-4a85-a266-a4686bd7d00f"
+            },
+        }
+        ctx, source, event_source = extract_dd_trace_context(sf_event, lambda_ctx)
+        self.assertEqual(source, "event")
+        expected_context = Context(
+            trace_id=6970872619724504833,
+            span_id=6880978411788117524,
+            sampling_priority=1,
+            meta={"_dd.p.tid": "71dab8f4d4629263"},
+        )
+        self.assertEqual(ctx, expected_context)
+        self.assertEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "6970872619724504833",
+                TraceHeader.PARENT_ID: "10713633173203262661",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+                TraceHeader.TAGS: "_dd.p.tid=71dab8f4d4629263",
+            },
+        )
+        create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
+        self.mock_send_segment.assert_called_with(
+            XraySubsegment.TRACE_KEY,
+            expected_context,
+        )
+
     def test_is_legacy_lambda_step_function(self):
         sf_event = {
             "Payload": {