support step function context in eventbridge extraction

Author: avedmala

datadog_lambda/tracing.py

@@ -320,12 +320,17 @@ def extract_context_from_eventbridge_event(event, lambda_context):
     """
     Extract datadog trace context from an EventBridge message's Details.
     This is only possible if Details is a JSON string.
+
+    If we find a Step Function context, try to extract the trace context from
+    that header.
     """
     try:
         detail = event.get("detail")
         dd_context = detail.get("_datadog")
         if not dd_context:
             return extract_context_from_lambda_context(lambda_context)
+        if is_step_function_event(dd_context):
+            return extract_context_from_step_functions(detail, lambda_context)
         return propagator.extract(dd_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
@@ -424,7 +429,7 @@ def _generate_sfn_trace_id(execution_id: str, part: str):
 def extract_context_from_step_functions(event, lambda_context):
     """
     Only extract datadog trace context when Step Functions Context Object is injected
-    into lambda's event dict.
+    into lambda's event dict. Unwrap "Payload" if it exists to handle Legacy Lambda cases.
 
     If '_datadog' header is present, we have two cases:
       1. Root is a Lambda and we use its traceID
@@ -435,6 +440,8 @@ def extract_context_from_step_functions(event, lambda_context):
     object.
     """
     try:
+        event = event.get("Payload", event)
+
         meta = {}
         dd_data = event.get("_datadog")
 
@@ -472,18 +479,30 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
-def is_legacy_lambda_step_function(event):
+def is_step_function_event(event):
     """
-    Check if the event is a step function that called a legacy lambda
+    Check if the event is a step function that invoked the current lambda.
+
+    The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a
+    "_datadog" for JSONata style context propagation.
+
+    The actual event must contain "Execution", "StateMachine", and "State" fields.
     """
-    if not isinstance(event, dict) or "Payload" not in event:
-        return False
+    event = event.get("Payload", event)
+
+    # JSONPath style
+    if all(field in event for field in ("Execution", "StateMachine", "State")):
+        return True
+
+    # JSONata style
+    if "_datadog" in event:
+        event = event["_datadog"]
+        return all(
+            field in event
+            for field in ("Execution", "StateMachine", "State", "serverless-version")
+        )
 
-    event = event.get("Payload")
-    return isinstance(event, dict) and (
-        "_datadog" in event
-        or ("Execution" in event and "StateMachine" in event and "State" in event)
-    )
+    return False
 
 
 def extract_context_custom_extractor(extractor, event, lambda_context):
@@ -1320,6 +1339,10 @@ def create_inferred_span_from_eventbridge_event(event, context):
     if span:
         span.set_tags(tags)
     span.start = dt.replace(tzinfo=timezone.utc).timestamp()
+
+    # Since inferred span will later parent Lambda, preserve Lambda's current parent
+    span.parent_id = dd_trace_context.span_id
+
     return span
 
 

datadog_lambda/trigger.py

@@ -10,6 +10,8 @@
 from enum import Enum
 from typing import Any
 
+from datadog_lambda.tracing import is_step_function_event
+
 
 class _stringTypedEnum(Enum):
     """
@@ -146,9 +148,7 @@ def parse_event_source(event: dict) -> _EventSource:
     if event.get("source") == "aws.events" or has_event_categories:
         event_source = _EventSource(EventTypes.CLOUDWATCH_EVENTS)
 
-    if (
-        "_datadog" in event and event.get("_datadog").get("serverless-version") == "v1"
-    ) or ("Execution" in event and "StateMachine" in event and "State" in event):
+    if is_step_function_event(event):
         event_source = _EventSource(EventTypes.STEPFUNCTIONS)
 
     event_record = get_first_record(event)

datadog_lambda/wrapper.py

@@ -45,7 +45,6 @@
     is_authorizer_response,
     tracer,
     propagator,
-    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -279,8 +278,6 @@ def _before(self, event, context):
             self.response = None
             set_cold_start(init_timestamp_ns)
             submit_invocations_metric(context)
-            if is_legacy_lambda_step_function(event):
-                event = event["Payload"]
             self.trigger_tags = extract_trigger_tags(event, context)
             # Extract Datadog trace context and source from incoming requests
             dd_context, trace_context_source, event_source = extract_dd_trace_context(

tests/test_tracing.py

@@ -42,7 +42,7 @@
     service_mapping as global_service_mapping,
     propagator,
     emit_telemetry_on_exception_outside_of_handler,
-    is_legacy_lambda_step_function,
+    is_step_function_event,
 )
 from datadog_lambda.trigger import EventTypes
 
@@ -836,6 +836,55 @@ def test_step_function_trace_data_sfn_root(self):
             expected_context,
         )
 
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_event_bridge(self):
+        lambda_ctx = get_mock_context()
+        sfn_event = {
+            "_datadog": {
+                "Execution": {
+                    "StartTime": "2025-03-11T01:16:31.408Z",
+                    "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+                    "RedriveCount": 0,
+                    "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
+                    "Name": "eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+                },
+                "StateMachine": {
+                    "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine",
+                    "Name": "abhinav-inner-state-machine",
+                },
+                "State": {
+                    "EnteredTime": "2025-03-11T01:16:31.448Z",
+                    "RetryCount": 0,
+                    "Name": "EventBridge PutEvents",
+                },
+                "serverless-version": "v1",
+                "RootExecutionId": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+            }
+        }
+        ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
+        self.assertEqual(source, "event")
+        expected_context = Context(
+            trace_id=4521899030418994483,
+            span_id=6880978411788117524,
+            sampling_priority=1,
+            meta={"_dd.p.tid": "12d1270d99cc5e03"},
+        )
+        self.assertEqual(ctx, expected_context)
+        self.assertEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "4521899030418994483",
+                TraceHeader.PARENT_ID: "2685222157636933868",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+                TraceHeader.TAGS: "_dd.p.tid=12d1270d99cc5e03",
+            },
+        )
+        create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
+        self.mock_send_segment.assert_called_with(
+            XraySubsegment.TRACE_KEY,
+            expected_context,
+        )
+
 
 class TestXRayContextConversion(unittest.TestCase):
     def test_convert_xray_trace_id(self):
@@ -2282,6 +2331,69 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
             if len(result_in_binary) == 66:  # "0b" + 64 bits.
                 self.assertTrue(result_in_binary.startswith("0b0"))
 
+    def test_is_step_function_event_jsonata(self):
+        event = {
+            "_datadog": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                    "RedriveCount": 0,
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                    "RetryCount": 0,
+                },
+                "x-datadog-trace-id": "5821803790426892636",
+                "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
+                "serverless-version": "v1",
+            }
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_jsonpath(self):
+        event = {
+            "Execution": {
+                "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                "RedriveCount": 0,
+            },
+            "StateMachine": {},
+            "State": {
+                "Name": "my-awesome-state",
+                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                "RetryCount": 0,
+            },
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_legacy_lambda(self):
+        event = {
+            "Payload": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                    "RedriveCount": 0,
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                    "RetryCount": 0,
+                },
+            }
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_dd_header(self):
+        event = {
+            "_datadog": {
+                "x-datadog-trace-id": "5821803790426892636",
+                "x-datadog-parent-id": "5821803790426892636",
+                "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
+                "x-datadog-sampling-priority": "1",
+            }
+        }
+        self.assertFalse(is_step_function_event(event))
+
 
 class TestExceptionOutsideHandler(unittest.TestCase):
     @patch("datadog_lambda.tracing.dd_tracing_enabled", True)