pull sfn trace id generation out into a helper

Author: avedmala

datadog_lambda/tracing.py

@@ -380,8 +380,8 @@ def _parse_high_64_bits(trace_tags: str) -> str:
     return ""
 
 
-def _sfn_context_to_parent_id(context: dict) -> int:
-    # todo: testme
+# todo: testme
+def _generate_sfn_parent_id(context: dict) -> int:
     execution_id = context.get("Execution").get("Id")
     state_name = context.get("State").get("Name")
     state_entered_time = context.get("State").get("EnteredTime")
@@ -391,6 +391,18 @@ def _sfn_context_to_parent_id(context: dict) -> int:
     )
 
 
+# take the higher 64 bits as _dd.p.tid tag and use hex to encode
+# [2:] to remove '0x' in the hex str
+# todo testme
+# returning 128 bits since 128bit traceId will be break up into
+# traditional traceId and _dd.p.tid tag
+# https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
+def _generate_sfn_trace_id(pre_hash: str, part: str):
+    if part == HIGHER_64_BITS:
+        return hex(_deterministic_sha256_hash(pre_hash, part))[2:]
+    return _deterministic_sha256_hash(pre_hash, part)
+
+
 def extract_context_from_step_functions(event, lambda_context):
     """
     Only extract datadog trace context when Step Functions Context Object is injected
@@ -410,29 +422,18 @@ def extract_context_from_step_functions(event, lambda_context):
                 if high_64_bit_trace_id:
                     meta["_dd.p.tid"] = high_64_bit_trace_id
             else:  # sfn root
-                trace_id = _deterministic_sha256_hash(
-                    dd_data.get("RootExecutionId"), LOWER_64_BITS
+                root_execution_id = dd_data.get("RootExecutionId")
+                trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS)
+                meta["_dd.p.tid"] = _generate_sfn_trace_id(
+                    root_execution_id, HIGHER_64_BITS
                 )
-                meta["_dd.p.tid"] = hex(
-                    _deterministic_sha256_hash(
-                        dd_data.get("RootExecutionId"), HIGHER_64_BITS
-                    )
-                )[2:]
 
-            parent_id = _sfn_context_to_parent_id(dd_data)
+            parent_id = _generate_sfn_parent_id(dd_data)
         else:
             execution_id = event.get("Execution").get("Id")
-            # returning 128 bits since 128bit traceId will be break up into
-            # traditional traceId and _dd.p.tid tag
-            # https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
-            trace_id = _deterministic_sha256_hash(execution_id, LOWER_64_BITS)
-            # take the higher 64 bits as _dd.p.tid tag and use hex to encode
-            # [2:] to remove '0x' in the hex str
-            meta["_dd.p.tid"] = hex(
-                _deterministic_sha256_hash(execution_id, HIGHER_64_BITS)
-            )[2:]
-
-            parent_id = _sfn_context_to_parent_id(event)
+            trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
+            meta["_dd.p.tid"] = _generate_sfn_trace_id(execution_id, HIGHER_64_BITS)
+            parent_id = _generate_sfn_parent_id(event)
 
         sampling_priority = SamplingPriority.AUTO_KEEP
         return Context(
@@ -454,8 +455,9 @@ def is_legacy_lambda_step_function(event):
         return False
 
     event = event.get("Payload")
-    return "_datadog" in event or (
-        "Execution" in event and "StateMachine" in event and "State" in event
+    return isinstance(event, dict) and (
+        "_datadog" in event
+        or ("Execution" in event and "StateMachine" in event and "State" in event)
     )
 
 

tests/test_tracing.py

@@ -665,10 +665,7 @@ def test_step_function_trace_data_lambda_root(self):
                     "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
                 },
                 "x-datadog-trace-id": "5821803790426892636",
-                "x-datadog-sampling-priority": "1",
                 "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
-                "traceparent": "00-672a7cb10000000050cb33b3c06ae95c-5fda9d8d1d1373f9-01",
-                "tracestate": "dd=p:5fda9d8d1d1373f9;s:1;t.dm:-0;t.tid:672a7cb100000000",
                 "serverless-version": "v2",
             }
         }
@@ -737,47 +734,6 @@ def test_step_function_trace_data_sfn_root(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 = {
-            "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",
-            },
-            "StateMachine": {},
-            "State": {
-                "Name": "my-awesome-state",
-                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
-            },
-        }
-        self.assertFalse(is_legacy_lambda_step_function(sf_event))
-
-        other_event = ["foo", "bar"]
-        self.assertFalse(is_legacy_lambda_step_function(other_event))
-
 
 class TestXRayContextConversion(unittest.TestCase):
     def test_convert_xray_trace_id(self):