feat: add support for event bridge DSM context extraction

README.md

@@ -30,6 +30,7 @@ Besides the environment variables supported by dd-trace-py, the datadog-lambda-p
 | DD_CAPTURE_LAMBDA_PAYLOAD | [Captures incoming and outgoing AWS Lambda payloads][1] in the Datadog APM spans for Lambda invocations. | `false` |
 | DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH | Determines the level of detail captured from AWS Lambda payloads, which are then assigned as tags for the `aws.lambda` span. It specifies the nesting depth of the JSON payload structure to process. Once the specified maximum depth is reached, the tag's value is set to the stringified value of any nested elements beyond this level.  <br> For example, given the input payload: <pre>{<br>  "lv1" : {<br>    "lv2": {<br>      "lv3": "val"<br>    }<br>  }<br>}</pre> If the depth is set to `2`, the resulting tag's key is set to `function.request.lv1.lv2` and the value is `{\"lv3\": \"val\"}`. <br> If the depth is set to `0`, the resulting tag's key is set to `function.request` and value is `{\"lv1\":{\"lv2\":{\"lv3\": \"val\"}}}` | `10` |
 | DD_EXCEPTION_REPLAY_ENABLED | When set to `true`, the Lambda will run with Error Tracking Exception Replay enabled, capturing local variables. | `false` |
+| DD_DSM_EXCHANGE_NAME | If you are using Datadog Data Streams Monitoring and propagating context through Amazon Event Bridge, this variable specifies the name of the EventBridge EventBus. The name of the bus is *not* passed through in the event payload and therefore cannot be automatically inferred at consume time. | `false` |
 
 
 ## Opening Issues

datadog_lambda/config.py

@@ -96,6 +96,10 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     data_streams_enabled = _get_env(
         "DD_DATA_STREAMS_ENABLED", "false", as_bool, depends_on_tracing=True
     )
+    # EventBridge bus name used as the DSM `exchange` tag. The bus name is not
+    # present in the inbound event, so it must be provided explicitly to allow
+    # the consume checkpoint to pair with the EventBridge produce checkpoint.
+    dsm_exchange_name = _get_env("DD_DSM_EXCHANGE_NAME")
     appsec_enabled = _get_env("DD_APPSEC_ENABLED", "false", as_bool)
     sca_enabled = _get_env("DD_APPSEC_SCA_ENABLED", "false", as_bool)
 

datadog_lambda/tracing.py

@@ -83,6 +83,47 @@ def _dsm_set_checkpoint(context_json, event_type, arn):
         )
 
 
+def _dsm_set_eventbridge_checkpoint(context_json, detail_type):
+    """Set a DSM consume checkpoint for an EventBridge event.
+
+    When ddtrace >= 4.14 is present the public ``tags`` parameter is used to
+    attach the ``exchange:`` edge tag (sourced from ``DD_DSM_EXCHANGE_NAME``).
+    On older installs the checkpoint is still emitted, just without the
+    exchange tag.
+    """
+    if not config.data_streams_enabled or not detail_type:
+        return
+
+    try:
+        from ddtrace.data_streams import set_consume_checkpoint
+
+        carrier_get = lambda k: context_json and context_json.get(k)  # noqa: E731
+        try:
+            tags = (
+                ["exchange:" + config.dsm_exchange_name]
+                if config.dsm_exchange_name
+                else None
+            )
+            set_consume_checkpoint(
+                "eventbridge",
+                detail_type,
+                carrier_get,
+                manual_checkpoint=False,
+                tags=tags,
+            )
+        except TypeError:
+            # ddtrace < 4.14 has no `tags` parameter. Retry without it, but
+            # keep `manual_checkpoint=False` so the checkpoint stays
+            # consistent with every other consume checkpoint.
+            set_consume_checkpoint(
+                "eventbridge", detail_type, carrier_get, manual_checkpoint=False
+            )
+    except Exception as e:
+        logger.debug(
+            f"DSM:Failed to set consume checkpoint for eventbridge {detail_type}: {e}"
+        )
+
+
 def _convert_xray_trace_id(xray_trace_id):
     """
     Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
@@ -251,11 +292,21 @@ def extract_context_from_sqs_or_sns_event_or_context(
     source_arn = ""
     event_type = "sqs" if event_source.equals(EventTypes.SQS) else "sns"
 
-    # EventBridge => SQS
+    # EventBridge => SQS. `dsm_handled` is True when the batch contained at
+    # least one EventBridge delivery and DSM checkpoints were already set
+    # per-record; in that case the regular SQS path below must not set another
+    # checkpoint for the first record or it would be double counted.
+    dsm_handled = False
     try:
-        context = _extract_context_from_eventbridge_sqs_event(event)
-        if _is_context_complete(context):
-            return context
+        (
+            context,
+            is_eventbridge_sqs,
+            dsm_handled,
+        ) = _extract_context_from_eventbridge_sqs_event(event)
+        if is_eventbridge_sqs:
+            if _is_context_complete(context):
+                return context
+            return extract_context_from_lambda_context(lambda_context)
     except Exception:
         logger.debug("Failed extracting context as EventBridge to SQS.")
 
@@ -312,7 +363,8 @@ def extract_context_from_sqs_or_sns_event_or_context(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
                 context = propagator.extract(dd_data)
-                _dsm_set_checkpoint(dd_data, event_type, source_arn)
+                if not dsm_handled:
+                    _dsm_set_checkpoint(dd_data, event_type, source_arn)
                 return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
@@ -337,12 +389,14 @@ def extract_context_from_sqs_or_sns_event_or_context(
                             sampling_priority=float(x_ray_context["sampled"]),
                         )
         # Still want to set a DSM checkpoint even if DSM context not propagated
-        _dsm_set_checkpoint(None, event_type, source_arn)
+        if not dsm_handled:
+            _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
         # Still want to set a DSM checkpoint even if DSM context not propagated
-        _dsm_set_checkpoint(None, event_type, source_arn)
+        if not dsm_handled:
+            _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
 
 
@@ -353,22 +407,190 @@ def _extract_context_from_eventbridge_sqs_event(event):
 
     This is only possible if first record in `Records` contains a
     `body` field which contains the EventBridge `detail` as a JSON string.
+
+    Returns a tuple ``(context, is_eventbridge_sqs, dsm_handled)``:
+
+    * ``context`` / ``is_eventbridge_sqs`` describe the trace context and are
+      derived only from the first record, since that is the record whose trace
+      context becomes the Lambda's parent.
+    * ``dsm_handled`` reports whether this function already set DSM checkpoints
+      for the batch. It is ``True`` whenever *any* record in the batch is an
+      EventBridge delivery, so the caller must not set its own SQS checkpoint
+      (which would double count the first record).
     """
-    first_record = event.get("Records")[0]
-    body_str = first_record.get("body")
-    body = json.loads(body_str)
-    detail = body.get("detail")
-    dd_context = detail.get("_datadog")
+    records = event.get("Records")
+    if not records:
+        return None, False, False
+
+    first_record = records[0]
+    dd_context, is_eventbridge_sqs = _extract_eventbridge_sqs_record_context(
+        first_record
+    )
+
+    # Set a consume checkpoint for every record in the batch whenever the batch
+    # contains at least one EventBridge delivery. Each record is classified
+    # independently so a mixed batch (EventBridge deliveries alongside direct
+    # SQS sends, in either order) uses the correct carrier per record. The
+    # message is consumed from the SQS queue, so it follows SQS conventions
+    # (type:sqs, topic:queue ARN).
+    dsm_handled = _dsm_set_eventbridge_sqs_batch_checkpoints(records)
+
+    if not is_eventbridge_sqs:
+        return None, False, dsm_handled
 
     if is_step_function_event(dd_context):
         try:
-            return extract_context_from_step_functions(dd_context, None)
+            return (
+                extract_context_from_step_functions(dd_context, None),
+                True,
+                dsm_handled,
+            )
         except Exception:
             logger.debug(
                 "Failed to extract Step Functions context from EventBridge to SQS event."
             )
 
-    return propagator.extract(dd_context)
+    return propagator.extract(dd_context), True, dsm_handled
+
+
+def _dsm_set_eventbridge_sqs_batch_checkpoints(records):
+    """Set a per-record SQS DSM consume checkpoint for an EventBridge -> SQS
+    batch.
+
+    Returns ``True`` when the batch contains at least one EventBridge delivery
+    (and checkpoints were therefore this function's responsibility), otherwise
+    ``False`` so the caller can fall back to its regular SQS checkpoint path.
+    Each record is classified independently: EventBridge records use the
+    carrier embedded in ``body.detail._datadog`` while other records fall back
+    to the SQS ``messageAttributes._datadog`` carrier.
+    """
+    if not config.data_streams_enabled:
+        return False
+
+    record_carriers = []
+    batch_has_eventbridge = False
+    for record in records:
+        record_context, is_eventbridge_record = _extract_eventbridge_sqs_record_context(
+            record
+        )
+        if is_eventbridge_record:
+            batch_has_eventbridge = True
+        else:
+            try:
+                record_context = _extract_sqs_record_message_attribute_context(record)
+            except Exception:
+                record_context = None
+        record_carriers.append(record_context)
+
+    if not batch_has_eventbridge:
+        return False
+
+    for record, record_context in zip(records, record_carriers):
+        try:
+            _dsm_set_checkpoint(record_context, "sqs", record.get("eventSourceARN", ""))
+        except Exception:
+            logger.debug(
+                "Failed to set DSM checkpoint for an EventBridge to SQS record."
+            )
+
+    return True
+
+
+def _extract_eventbridge_sqs_record_context(record):
+    """Classify a single SQS record as an EventBridge delivery and return its
+    DSM carrier.
+
+    Returns a tuple ``(dd_context, is_eventbridge)``. A record is only treated
+    as EventBridge when its ``body`` is a JSON object carrying the EventBridge
+    envelope fields (``detail`` object plus ``detail-type`` and ``source``). A
+    non-JSON or non-envelope body is not an error here: it simply means the
+    record is a regular SQS message, so the caller can fall back to the SQS
+    message attribute carrier instead.
+    """
+    body_str = record.get("body")
+    try:
+        body = json.loads(body_str)
+    except (ValueError, TypeError):
+        return None, False
+
+    if not isinstance(body, dict):
+        return None, False
+
+    detail = body.get("detail")
+    if not (
+        isinstance(detail, dict) and body.get("detail-type") and body.get("source")
+    ):
+        return None, False
+
+    return detail.get("_datadog"), True
+
+
+def _extract_sqs_record_message_attribute_context(record):
+    """Return the ``_datadog`` carrier for a non-EventBridge SQS record.
+
+    A record in an SQS batch may itself be an SNS notification (an SNS => SQS
+    subscription without raw message delivery), in which case the attributes
+    live in the SNS envelope inside ``body`` rather than in the record's own
+    ``messageAttributes``. Both shapes are handled here so that a mixed batch
+    (EventBridge deliveries alongside SNS deliveries on the same queue) does
+    not lose the SNS DSM context.
+    """
+    msg_attributes = record.get("messageAttributes")
+
+    if not msg_attributes:
+        sns_record = _parse_sns_notification_from_sqs_body(record) or {}
+        msg_attributes = sns_record.get("MessageAttributes") or {}
+
+    return _decode_dd_message_attribute(msg_attributes.get("_datadog"))
+
+
+def _parse_sns_notification_from_sqs_body(record):
+    """Return the SNS envelope carried in an SQS record's ``body``, if any.
+
+    Returns ``None`` when the body is not an SNS notification, which simply
+    means the record is a direct SQS send.
+    """
+    try:
+        body = json.loads(record.get("body"))
+    except (ValueError, TypeError):
+        return None
+
+    if (
+        isinstance(body, dict)
+        and body.get("Type", "") == "Notification"
+        and "TopicArn" in body
+    ):
+        return body
+
+    return None
+
+
+def _decode_dd_message_attribute(dd_payload):
+    """Decode a ``_datadog`` SQS/SNS message attribute into a carrier dict.
+
+    SQS uses ``dataType`` with ``stringValue``/``binaryValue`` while SNS uses
+    ``Type`` with ``Value``; both are supported.
+    """
+    if not dd_payload:
+        return None
+
+    dd_json_data = None
+    dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType")
+    if dd_json_data_type == "Binary":
+        import base64
+
+        dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value")
+        if dd_json_data:
+            dd_json_data = base64.b64decode(dd_json_data)
+    elif dd_json_data_type == "String":
+        dd_json_data = dd_payload.get("stringValue") or dd_payload.get("Value")
+    else:
+        logger.debug(
+            "Datadog Lambda Python only supports extracting trace"
+            "context from String or Binary SQS/SNS message attributes"
+        )
+
+    return json.loads(dd_json_data) if dd_json_data else None
 
 
 def extract_context_from_eventbridge_event(event, lambda_context):
@@ -380,8 +602,11 @@ def extract_context_from_eventbridge_event(event, lambda_context):
     that header.
     """
     try:
-        detail = event.get("detail")
+        detail = event.get("detail") or {}
         dd_context = detail.get("_datadog")
+
+        _dsm_set_eventbridge_checkpoint(dd_context, event.get("detail-type"))
+
         if not dd_context:
             return extract_context_from_lambda_context(lambda_context)