Check for the new direct invocation trace format and use that (#146)

* Check for the new direct invocation trace format and use that

Author: agocs

datadog_lambda/tracing.py

@@ -129,16 +129,27 @@ def extract_context_from_lambda_context(lambda_context):
     dd_trace libraries inject this trace context on synchronous invocations
     """
     client_context = lambda_context.client_context
-
+    trace_id = None
+    parent_id = None
+    sampling_priority = None
     if client_context and client_context.custom:
-        dd_data = client_context.custom.get("_datadog", {})
-        trace_id = dd_data.get(TraceHeader.TRACE_ID)
-        parent_id = dd_data.get(TraceHeader.PARENT_ID)
-        sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
+        if "_datadog" in client_context.custom:
+            # Legacy trace propagation dict
+            dd_data = client_context.custom.get("_datadog", {})
+            trace_id = dd_data.get(TraceHeader.TRACE_ID)
+            parent_id = dd_data.get(TraceHeader.PARENT_ID)
+            sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
+        elif (
+            TraceHeader.TRACE_ID in client_context.custom
+            and TraceHeader.PARENT_ID in client_context.custom
+            and TraceHeader.SAMPLING_PRIORITY in client_context.custom
+        ):
+            # New trace propagation keys
+            trace_id = client_context.custom.get(TraceHeader.TRACE_ID)
+            parent_id = client_context.custom.get(TraceHeader.PARENT_ID)
+            sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY)
 
-        return trace_id, parent_id, sampling_priority
-
-    return None, None, None
+    return trace_id, parent_id, sampling_priority
 
 
 def extract_context_from_http_event_or_context(event, lambda_context):

tests/test_tracing.py

@@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self):
             XraySubsegment.NAMESPACE,
         )
 
-    def test_with_client_context_datadog_trace_data(self):
+    def test_with_legacy_client_context_datadog_trace_data(self):
         lambda_ctx = get_mock_context(
             custom={
                 "_datadog": {
@@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self):
             XraySubsegment.NAMESPACE,
         )
 
+    def test_with_new_client_context_datadog_trace_data(self):
+        lambda_ctx = get_mock_context(
+            custom={
+                TraceHeader.TRACE_ID: "666",
+                TraceHeader.PARENT_ID: "777",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+            }
+        )
+        ctx, source = extract_dd_trace_context({}, lambda_ctx)
+        self.assertEqual(source, "event")
+        self.assertDictEqual(
+            ctx,
+            {
+                "trace-id": "666",
+                "parent-id": "777",
+                "sampling-priority": "1",
+            },
+        )
+        self.assertDictEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "666",
+                TraceHeader.PARENT_ID: "65535",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+            },
+        )
+        create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
+        self.mock_xray_recorder.begin_subsegment.assert_called()
+        self.mock_xray_recorder.end_subsegment.assert_called()
+        self.mock_current_subsegment.put_metadata.assert_called_with(
+            XraySubsegment.TRACE_KEY,
+            {"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
+            XraySubsegment.NAMESPACE,
+        )
+
     def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
         lambda_ctx = get_mock_context()
         extract_dd_trace_context(