fix legacy lambda parsing with new header

Author: avedmala

datadog_lambda/tracing.py

@@ -432,8 +432,13 @@ 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
+    if not isinstance(event, dict) or "Payload" not in event:
+        return False
+
+    event = event.get("Payload")
+    return "_datadog" in event or (
+        "Execution" in event and "StateMachine" in event and "State" in event
+    )
 
 
 def extract_context_custom_extractor(extractor, event, lambda_context):
@@ -688,6 +693,7 @@ def create_inferred_span(
     event_source: _EventSource = None,
     decode_authorizer_context: bool = True,
 ):
+    logger.debug("abhinav event %s", event)
     if event_source is None:
         event_source = parse_event_source(event)
     try:

tests/test_tracing.py

@@ -738,6 +738,17 @@ def test_is_legacy_lambda_step_function(self):
         }
         self.assertTrue(is_legacy_lambda_step_function(sf_event))
 
+        sf_event = {
+            "Payload": {
+                "_datadog": {
+                    "x-datadog-trace-id-hash": "fed93f8c162880cb9aa90fcd1f8395383835841d5470d30215f3dd52906ebc58",
+                    "x-datadog-parent-id-hash": "c5eb94cc9220ab5783e1db53debd54b8c93f6f2a3eae1c680d7b849f2d34e551",
+                    "serverless-version": "v2",
+                }
+            }
+        }
+        self.assertTrue(is_legacy_lambda_step_function(sf_event))
+
         sf_event = {
             "Execution": {
                 "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",