Refactoring change to use extract_context_from_http_event_or_context

Author: rithikanarayan

datadog_lambda/tracing.py

@@ -537,25 +537,6 @@ 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
@@ -646,6 +627,14 @@ def extract_dd_trace_context(
 
     if extractor is not None:
         context = extract_context_custom_extractor(extractor, event, lambda_context)
+    elif isinstance(event, (set, dict)) and "request" in event:
+        if "headers" in event.get("request"):
+            context = extract_context_from_http_event_or_context(
+                event.get("request"),
+                lambda_context,
+                event_source,
+                decode_authorizer_context=False,
+            )
     elif isinstance(event, (set, dict)) and "headers" in event:
         context = extract_context_from_http_event_or_context(
             event, lambda_context, event_source, decode_authorizer_context
@@ -661,7 +650,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_request_header(event, lambda_context)
+        context = extract_context_from_lambda_context(lambda_context)
 
     if _is_context_complete(context):
         logger.debug("Extracted Datadog trace context from event or context")

publish_test_layer.sh

@@ -0,0 +1,21 @@
+REGION=us-east-1 PYTHON_VERSION=$PYTHON_VERSION ARCH=$ARCH ./scripts/build_layers.sh
+
+LAYER_NAME="Python${PYTHON_VERSION/3./3}-RITHIKA"
+LAYER_ZIP=".layers/datadog_lambda_py-${ARCH}-${PYTHON_VERSION}.zip"
+REGION="us-east-1"
+
+echo "zipping package files"
+PYTHON_VERSION=$PYTHON_VERSION ARCH=$ARCH ./scripts/build_layers.sh
+
+echo "publishing layer"
+
+layer_arn=$(
+    aws-vault exec sso-serverless-sandbox-account-admin -- \
+        aws lambda publish-layer-version --layer-name $LAYER_NAME \
+            --description "Datadog Tracer Lambda Layer for Python" \
+            --zip-file "fileb://$LAYER_ZIP" \
+            --region "$REGION" \
+            --output json \
+                | jq -r '.LayerVersionArn')
+
+echo "new python layer published $layer_arn"

tests/integration/input_events/appsync.json

@@ -0,0 +1,49 @@
+{
+    "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-forwarded-for": "64.124.12.19,15.158.225.229",
+            "x-forwarded-port": "443",
+            "x-forwarded-proto": "https"
+        }
+    },
+    "source": "None"
+}