Add extractor for Eventbridge context (#202)

* Update expected breaking change date (#114)

* change inferred_span to _inferred_span

* Add extractor for eventbridge trace context

* Add another test to test eventbridge extraction

* Get tracing.Literal for pythons that don't have it already

* Add _datadog to eventbridge extractor

* Update integration tests

* Remove init complete and main start logs

Co-authored-by: Jorie Helwig <[email protected]>
Co-authored-by: Harvinder Ghotra <[email protected]>

Author: 3 people

README.md

@@ -8,7 +8,7 @@
 
 Datadog Lambda Library for Python (3.6, 3.7, 3.8, and 3.9) enables enhanced Lambda metrics, distributed tracing, and custom metric submission from AWS Lambda functions.  
 
-**IMPORTANT NOTE:** AWS Lambda is expected to receive a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **March 31, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to the latest.
+**IMPORTANT NOTE:** AWS Lambda is expected to recieve a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **March 31, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to the latest.
 
 ## Installation
 

datadog_lambda/tracing.py

@@ -7,7 +7,13 @@
 import os
 import json
 from datetime import datetime, timezone
-from typing import Optional, Literal, Dict
+from typing import Optional, Dict
+
+try:
+    from typing import Literal
+except ImportError:
+    # Literal was added to typing in python 3.8
+    from typing_extensions import Literal
 
 from datadog_lambda.constants import (
     SamplingPriority,
@@ -227,6 +233,25 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
+def extract_context_from_eventbridge_event(event, lambda_context):
+    """
+    Extract datadog trace context from an EventBridge message's Details.
+    Details is often a weirdly escaped almost-JSON string. Here we have to correct for that.
+    """
+    try:
+        detail = event["detail"]
+        dd_context = detail.get("_datadog")
+        print(f"dd_context is {dd_context}")
+        if not dd_context:
+            return extract_context_from_lambda_context(lambda_context)
+        trace_id = dd_context.get(TraceHeader.TRACE_ID)
+        parent_id = dd_context.get(TraceHeader.PARENT_ID)
+        sampling_priority = dd_context.get(TraceHeader.SAMPLING_PRIORITY)
+        return trace_id, parent_id, sampling_priority
+    except Exception:
+        return extract_context_from_lambda_context(lambda_context)
+
+
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function
@@ -253,6 +278,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
     """
     global dd_trace_context
     trace_context_source = None
+    event_source = parse_event_source(event)
 
     if extractor is not None:
         (
@@ -272,6 +298,12 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
             parent_id,
             sampling_priority,
         ) = extract_context_from_sqs_or_sns_event_or_context(event, lambda_context)
+    elif event_source.equals(EventTypes.EVENTBRIDGE):
+        (
+            trace_id,
+            parent_id,
+            sampling_priority,
+        ) = extract_context_from_eventbridge_event(event, lambda_context)
     else:
         trace_id, parent_id, sampling_priority = extract_context_from_lambda_context(
             lambda_context
@@ -775,7 +807,7 @@ def create_function_execution_span(
 
 
 class InferredSpanInfo(object):
-    BASE_NAME = "inferred_span"
+    BASE_NAME = "_inferred_span"
     SYNCHRONICITY = f"{BASE_NAME}.synchronicity"
     TAG_SOURCE = f"{BASE_NAME}.tag_source"
 

pyproject.toml

@@ -29,6 +29,7 @@ wrapt = "^1.11.2"
 ddtrace = "^0.50.0"
 importlib_metadata = {version = "^1.0", python = "<3.8"}
 boto3 = { version = "^1.10.33", optional = true }
+typing_extensions = {version = "^4.0", python = "<3.8"}
 requests = { version ="^2.22.0", optional = true }
 nose2 = { version= "^0.9.1", optional = true }
 flake8 = { version = "^3.7.9", optional = true }

scripts/run_integration_tests.sh

@@ -212,7 +212,9 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 sed -E "s/(\"system\.pid\"\: )[0-9\.\-]+/\1\"XXXX\"/g" |
                 sed -E "s/(\"runtime-id\"\: \")[a-z0-9\.\-]+/\1XXXX/g" |
                 sed -E "s/(\"datadog_lambda\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g" |
-                sed -E "s/(\"dd_trace\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g"
+                sed -E "s/(\"dd_trace\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g" |
+                sed -E "/init complete at epoch/d" |
+                sed -E "/main started at epoch/d"
         )
 
         if [ ! -f $function_snapshot_path ]; then

tests/event_samples/eventbridge-custom.json

@@ -8,6 +8,11 @@
   "region": "sa-east-1",
   "resources": [],
   "detail": {
-    "foo": "bar"
+    "foo": "bar",
+    "_datadog": {
+      "x-datadog-trace-id": "12345",
+      "x-datadog-parent-id": "67890",
+      "x-datadog-sampling-priority": "2"
+    }
   }
 }