Step Functions Legacy Lambda Span Linking (#511)

This change will check if events fall into the Legacy Lambda case and parse them accordingly so that we can extract the trace context and infer a span link between the Step Function and the downstream Legacy Lambda

Author: avedmala

Diff for: datadog_lambda/tracing.py

@@ -407,6 +407,14 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
+def is_legacy_lambda_step_function(event):
+    """
+    Check if the event is a step function that called a legacy lambda
+    """
+    event = event.get("Payload", {})
+    return "Execution" in event and "StateMachine" in event and "State" in event
+
+
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function

Diff for: datadog_lambda/wrapper.py

@@ -44,6 +44,7 @@
     is_authorizer_response,
     tracer,
     propagator,
+    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -277,6 +278,8 @@ 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(

Diff for: tests/test_tracing.py

@@ -40,6 +40,7 @@
     service_mapping as global_service_mapping,
     propagator,
     emit_telemetry_on_exception_outside_of_handler,
+    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import EventTypes
 
@@ -647,6 +648,33 @@ def test_step_function_trace_data(self):
             expected_context,
         )
 
+    def test_is_legacy_lambda_step_function(self):
+        sf_event = {
+            "Payload": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                },
+            }
+        }
+        self.assertTrue(is_legacy_lambda_step_function(sf_event))
+
+        sf_event = {
+            "Execution": {
+                "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+            },
+            "StateMachine": {},
+            "State": {
+                "Name": "my-awesome-state",
+                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+            },
+        }
+        self.assertFalse(is_legacy_lambda_step_function(sf_event))
+
 
 class TestXRayContextConversion(unittest.TestCase):
     def test_convert_xray_trace_id(self):