feat: add support for event bridge DSM context extraction

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,147 @@ 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):
+    msg_attributes = record.get("messageAttributes") or {}
+    dd_payload = msg_attributes.get("_datadog")
+    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 +559,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)