Add extractor for Eventbridge context

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"
+ }
  }
 }