chore: update processing of eventbridge DSM

Author: jeastham1993

datadog_lambda/tracing.py

@@ -99,15 +99,16 @@ def _dsm_set_eventbridge_checkpoint(context_json, detail_type):
         return
 
     try:
-        from ddtrace.internal.datastreams import data_streams_processor
-        from ddtrace.internal.datastreams.processor import PROPAGATION_KEY_BASE_64
+        from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
+        from ddtrace.data_streams import ddtrace as ddtrace_data_streams
 
-        processor = data_streams_processor()
+        processor = getattr(ddtrace_data_streams.tracer, "data_streams_processor", None)
         if not processor:
             return
 
-        carrier_get = lambda k: context_json and context_json.get(k)  # noqa: E731
-        processor.decode_pathway_b64(carrier_get(PROPAGATION_KEY_BASE_64))
+        processor.decode_pathway_b64(
+            context_json.get(PROPAGATION_KEY_BASE_64) if context_json else None
+        )
 
         tags = ["direction:in", "topic:" + detail_type, "type:eventbridge"]
         if config.dsm_exchange_name:
@@ -289,9 +290,13 @@ def extract_context_from_sqs_or_sns_event_or_context(
 
     # EventBridge => SQS
     try:
-        context = _extract_context_from_eventbridge_sqs_event(event)
-        if _is_context_complete(context):
-            return context
+        context, is_eventbridge_sqs = _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.")
 
@@ -390,24 +395,35 @@ 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.
     """
-    records = event.get("Records")
+    records = event.get("Records") or []
+    if not records:
+        return None, False
+
     first_record = records[0]
     body_str = first_record.get("body")
     body = json.loads(body_str)
     detail = body.get("detail")
-    # If `detail` is missing this is not an EventBridge -> SQS event; raising
-    # here lets the caller fall back to the regular SQS extraction path before
-    # any DSM checkpoint is set, avoiding double counting for plain SQS events.
+    if not isinstance(detail, dict):
+        return None, False
+
     dd_context = detail.get("_datadog")
 
     # The event has been confirmed as EventBridge -> SQS. Set a consume
     # checkpoint for every record in the batch. The message is consumed from
     # the SQS queue, so it follows SQS conventions (type:sqs, topic:queue ARN).
     if config.data_streams_enabled:
+        _dsm_set_checkpoint(dd_context, "sqs", first_record.get("eventSourceARN", ""))
         for record in records:
+            if record is first_record:
+                continue
             try:
                 record_body = json.loads(record.get("body"))
-                record_context = (record_body.get("detail") or {}).get("_datadog")
+                record_detail = record_body.get("detail")
+                record_context = (
+                    record_detail.get("_datadog")
+                    if isinstance(record_detail, dict)
+                    else None
+                )
                 _dsm_set_checkpoint(
                     record_context, "sqs", record.get("eventSourceARN", "")
                 )
@@ -418,13 +434,13 @@ def _extract_context_from_eventbridge_sqs_event(event):
 
     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
         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
 
 
 def extract_context_from_eventbridge_event(event, lambda_context):

tests/test_tracing.py

@@ -3651,19 +3651,23 @@ def test_kinesis_data_streams_disabled(self):
     # EVENTBRIDGE -> SQS TESTS
 
     @staticmethod
-    def _eventbridge_sqs_record(queue_arn, pathway_ctx):
-        body = {
-            "detail-type": "MyDetailType",
-            "source": "my.event.source",
-            "detail": {
-                "_datadog": {
+    def _eventbridge_sqs_record(queue_arn, pathway_ctx, include_trace_headers=True):
+        dd_context = {"dd-pathway-ctx-base64": pathway_ctx}
+        if include_trace_headers:
+            dd_context.update(
+                {
                     # Complete trace context so the extractor returns early and
                     # does not fall through to the regular SQS path.
                     "x-datadog-trace-id": "12345",
                     "x-datadog-parent-id": "67890",
                     "x-datadog-sampling-priority": "1",
-                    "dd-pathway-ctx-base64": pathway_ctx,
                 }
+            )
+        body = {
+            "detail-type": "MyDetailType",
+            "source": "my.event.source",
+            "detail": {
+                "_datadog": dd_context
             },
         }
         return {
@@ -3710,6 +3714,35 @@ def test_eventbridge_sqs_checkpoints_all_records(self):
         self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2))
         self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "ctx-2")
 
+    @patch(
+        "datadog_lambda.tracing.extract_context_from_lambda_context",
+        return_value=Context(trace_id=111, span_id=222, sampling_priority=1),
+    )
+    def test_eventbridge_sqs_incomplete_context_uses_single_checkpoint(
+        self, mock_extract_context
+    ):
+        queue_arn = "arn:aws:sqs:us-east-1:123456789012:eb-queue"
+        event = {
+            "Records": [
+                self._eventbridge_sqs_record(
+                    queue_arn, "ctx-only", include_trace_headers=False
+                )
+            ]
+        }
+
+        context = extract_context_from_sqs_or_sns_event_or_context(
+            event, self.lambda_context, parse_event_source(event)
+        )
+
+        self.assertEqual(context.trace_id, 111)
+        self.assertEqual(context.span_id, 222)
+        mock_extract_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(self.mock_checkpoint.call_count, 1)
+        args, _ = self.mock_checkpoint.call_args
+        self.assertEqual(args[0], "sqs")
+        self.assertEqual(args[1], queue_arn)
+        self.assertEqual(args[2]("dd-pathway-ctx-base64"), "ctx-only")
+
     @patch("datadog_lambda.config.Config.data_streams_enabled", False)
     def test_eventbridge_sqs_data_streams_disabled(self):
         queue_arn = "arn:aws:sqs:us-east-1:123456789012:eb-queue"
@@ -3726,12 +3759,12 @@ class TestEventBridgeDSMLogic(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
         self.mock_processor = Mock()
-        processor_patcher = patch(
-            "ddtrace.internal.datastreams.data_streams_processor",
-            return_value=self.mock_processor,
+        tracer_patcher = patch(
+            "ddtrace.data_streams.ddtrace.tracer",
+            new=SimpleNamespace(data_streams_processor=self.mock_processor),
         )
-        processor_patcher.start()
-        self.addCleanup(processor_patcher.stop)
+        tracer_patcher.start()
+        self.addCleanup(tracer_patcher.stop)
         config_patcher = patch(
             "datadog_lambda.config.Config.data_streams_enabled", True
         )