APMSVLS-65 Extract Trace Context For AppSync Events

datadog_lambda/tracing.py

@@ -537,6 +537,25 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
+def extract_context_from_request_header(event, lambda_context):
+    """
+    Attempt to extract Datadog trace context from an event payload where Datadog information
+    is under the event["request"]["header"] section. Intended to replace some of the
+    need for customers to provide custom extractor functions in their code for workflows such as
+    RUM -> AppSync -> Lambda.
+    """
+    try:
+        nested_json = event.get("request").get("headers")
+        context = propagator.extract(nested_json)
+
+        if not _is_context_complete(context):
+            context = extract_context_from_lambda_context(lambda_context)
+    except Exception:
+        context = extract_context_from_lambda_context(lambda_context)
+
+    return context
+
+
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function
@@ -642,7 +661,7 @@ def extract_dd_trace_context(
     elif event_source.equals(EventTypes.STEPFUNCTIONS):
         context = extract_context_from_step_functions(event, lambda_context)
     else:
-        context = extract_context_from_lambda_context(lambda_context)
+        context = extract_context_from_request_header(event, lambda_context)
 
     if _is_context_complete(context):
         logger.debug("Extracted Datadog trace context from event or context")

tests/event_samples/rum-appsync.json

@@ -0,0 +1,53 @@
+{
+    "identity": "None",
+    "info": {
+        "fieldName": "getItems",
+        "parentTypeName": "Query",
+        "selectionSetGraphQL": "{\n  id\n}",
+        "selectionSetList":["id"]
+    }, 
+    "prev": "None",
+    "request": { 
+        "domainName": "None",
+        "headers": {
+            "accept": "*/*",
+            "accept-encoding": "gzip, deflate, br, zstd",
+            "accept-language": "en-US,en;q=0.9",
+            "cloudfront-forwarded-proto": "https",
+            "cloudfront-is-desktop-viewer": "True",
+            "cloudfront-is-mobile-viewer": "False",
+            "cloudfront-is-smarttv-viewer":"False", 
+            "cloudfront-is-tablet-viewer":" False",
+            "cloudfront-viewer-asn": "6461",
+            "cloudfront-viewer-country": "US",
+            "content-length": "47",
+            "content-type": "application/graphql",
+            "host": "4aowrg2uhvbw5mn7osu6searqi.appsync-api.us-east-1.amazonaws.com",
+            "origin": "http://localhost:5173", 
+            "priority": "u=1, i",
+            "referer": "http://localhost:5173/",
+            "sec-ch-ua": "\"Chromium\";v=\"140\", \"Not=A?Brand\";v=\"24\", \"Google Chrome\";v=\"140\"",
+            "sec-ch-ua-mobile": "?0",
+            "sec-ch-ua-platform": "macOS",
+            "sec-fetch-dest": "empty",
+            "sec-fetch-mode": "cors",
+            "sec-fetch-site": "cross-site",
+            "traceparent": "00-0000000000000000d9f454c80b9a529a-73ac6ca3427073a3-01",
+            "tracestate": "dd=s:1;o:rum",
+            "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36","via":"2.0 62f9f3967e93a923f21c8acb20cf10b6.cloudfront.net (CloudFront)",
+            "x-amz-cf-id": "Femhicb_Vbva-J8qWjdI4hKMmqusCeQSp207UGyY3u8VOUrdE8BBvg==",
+            "x-amzn-appsync-is-vpce-request": "False",
+            "x-amzn-remote-ip": "64.124.12.19",
+            "x-amzn-requestid": "1ee1669a-eda3-4d4f-911c-35f74ebef31d",
+            "x-amzn-trace-id": "Root=1-68cdb7e9-438bf88259d7ce3535723bdd",
+            "x-datadog-origin": "rum",
+            "x-datadog-parent-id": "67890",
+            "x-datadog-sampling-priority": "1",
+            "x-datadog-trace-id": "12345",
+            "x-forwarded-for": "64.124.12.19,15.158.225.229",
+            "x-forwarded-port": "443",
+            "x-forwarded-proto": "https"
+        }
+    },
+    "source": "None"
+}

tests/test_tracing.py

@@ -179,6 +179,14 @@ def _wrap(*args, **kwargs):
         ),
     ),
     ("lambda-url", None),
+    (
+        "rum-appsync", 
+        Context(
+            trace_id=12345,
+            span_id=67890,
+            sampling_priority=1,
+        )
+    ),
     ("s3", None),
     (
         "sns-b64-msg-attribute",
@@ -526,6 +534,77 @@ def test_with_sqs_distributed_w3c_trace_data(self):
         }
         self.common_tests_with_trace_context_extraction_injection(headers, sqs_event)
 
+    @with_trace_propagation_style("datadog")
+    def test_with_request_header_datadog_trace_data(self):
+        """Test trace data extraction from event['request']['headers']"""
+        headers = {
+            TraceHeader.TRACE_ID: "12345",
+            TraceHeader.PARENT_ID: "67890",
+            TraceHeader.SAMPLING_PRIORITY: "1",
+        }
+        request_headers = {
+            "request": { 
+                "domainName": "None",
+                "headers": {
+                    "content-length": "47",
+                    "content-type": "application/graphql",
+                    "host": "4aowrg2uhvbw5mn7osu6searqi.appsync-api.us-east-1.amazonaws.com",
+                    "origin": "http://localhost:5173", 
+                    "priority": "u=1, i",
+                    "referer": "http://localhost:5173/",
+                    "traceparent": "00-0000000000000000d9f454c80b9a529a-73ac6ca3427073a3-01",
+                    "tracestate": "dd=s:1;o:rum",
+                    "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36","via":"2.0 62f9f3967e93a923f21c8acb20cf10b6.cloudfront.net (CloudFront)",
+                    "x-amz-cf-id": "Femhicb_Vbva-J8qWjdI4hKMmqusCeQSp207UGyY3u8VOUrdE8BBvg==",
+                    "x-amzn-appsync-is-vpce-request": "False",
+                    "x-amzn-remote-ip": "64.124.12.19",
+                    "x-amzn-requestid": "1ee1669a-eda3-4d4f-911c-35f74ebef31d",
+                    "x-amzn-trace-id": "Root=1-68cdb7e9-438bf88259d7ce3535723bdd",
+                    "x-datadog-parent-id": "67890",
+                    "x-datadog-sampling-priority": "1",
+                    "x-datadog-trace-id": "12345",
+                    "x-forwarded-for": "64.124.12.19,15.158.225.229",
+                    "x-forwarded-port": "443",
+                    "x-forwarded-proto": "https"
+                }
+            },
+        }
+        self.common_tests_with_trace_context_extraction_injection(headers, request_headers)
+
+    @with_trace_propagation_style("tracecontext")
+    def test_with_request_header_w3_trace_data(self):
+        """Test trace data extraction from event['request']['headers']"""
+        headers = {
+            "traceparent": "00-0000000000000000d9f454c80b9a529a-73ac6ca3427073a3-01",
+            "tracestate": "dd=s:1;o:rum",
+        }
+        request_headers = {
+            "request": { 
+                "domainName": "None",
+                "headers": {
+                    "content-length": "47",
+                    "content-type": "application/graphql",
+                    "host": "4aowrg2uhvbw5mn7osu6searqi.appsync-api.us-east-1.amazonaws.com",
+                    "origin": "http://localhost:5173", 
+                    "priority": "u=1, i",
+                    "referer": "http://localhost:5173/",
+                    "traceparent": "00-0000000000000000d9f454c80b9a529a-73ac6ca3427073a3-01",
+                    "tracestate": "dd=s:1;o:rum",
+                    "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36","via":"2.0 62f9f3967e93a923f21c8acb20cf10b6.cloudfront.net (CloudFront)",
+                    "x-amz-cf-id": "Femhicb_Vbva-J8qWjdI4hKMmqusCeQSp207UGyY3u8VOUrdE8BBvg==",
+                    "x-amzn-appsync-is-vpce-request": "False",
+                    "x-amzn-remote-ip": "64.124.12.19",
+                    "x-amzn-requestid": "1ee1669a-eda3-4d4f-911c-35f74ebef31d",
+                    "x-amzn-trace-id": "Root=1-68cdb7e9-438bf88259d7ce3535723bdd",
+                    "x-forwarded-for": "64.124.12.19,15.158.225.229",
+                    "x-forwarded-port": "443",
+                    "x-forwarded-proto": "https"
+                }
+            },
+        }
+        self.common_tests_with_trace_context_extraction_injection(headers, request_headers)
+
+
     @with_trace_propagation_style("datadog")
     def test_with_legacy_client_context_datadog_trace_data(self):
         headers = {