fix: use public dsm checkpoint for eventbridge

Author: jeastham1993

datadog_lambda/tracing.py

@@ -85,39 +85,12 @@ 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.
 
-    Unlike the SQS/SNS/Kinesis helper, the EventBridge edge tags include an
-    `exchange` tag (the bus name) to mirror the produce-side tags so the
-    consume node pairs with the produce node. The bus name is not present in
-    the inbound event, so it is sourced from `DD_DSM_EXCHANGE_NAME` when set.
-    The public `set_consume_checkpoint` helper cannot emit an `exchange` tag,
-    so the lower-level processor API is used directly.
+    Temporary note: richer EventBridge consume checkpoint tagging is being
+    upstreamed into dd-trace-py. Until that lands, keep this on the same
+    public consume checkpoint API used by SQS/SNS/Kinesis and omit the
+    EventBridge exchange tag for now.
     """
-    if not config.data_streams_enabled:
-        return
-
-    if not detail_type:
-        return
-
-    try:
-        from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
-        from ddtrace.data_streams import ddtrace as ddtrace_data_streams
-
-        processor = getattr(ddtrace_data_streams.tracer, "data_streams_processor", None)
-        if not processor:
-            return
-
-        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:
-            tags.append("exchange:" + config.dsm_exchange_name)
-        processor.set_checkpoint(tags)
-    except Exception as e:
-        logger.debug(
-            f"DSM:Failed to set consume checkpoint for eventbridge {detail_type}: {e}"
-        )
+    _dsm_set_checkpoint(context_json, "eventbridge", detail_type)
 
 
 def _convert_xray_trace_id(xray_trace_id):

tests/test_tracing.py

@@ -3831,13 +3831,9 @@ def test_eventbridge_sqs_data_streams_disabled(self):
 class TestEventBridgeDSMLogic(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
-        self.mock_processor = Mock()
-        tracer_patcher = patch(
-            "ddtrace.data_streams.ddtrace.tracer",
-            new=SimpleNamespace(data_streams_processor=self.mock_processor),
-        )
-        tracer_patcher.start()
-        self.addCleanup(tracer_patcher.stop)
+        checkpoint_patcher = patch("ddtrace.data_streams.set_consume_checkpoint")
+        self.mock_checkpoint = checkpoint_patcher.start()
+        self.addCleanup(checkpoint_patcher.stop)
         config_patcher = patch(
             "datadog_lambda.config.Config.data_streams_enabled", True
         )
@@ -3857,51 +3853,61 @@ def test_eventbridge_context_propagated(self):
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        self.mock_processor.decode_pathway_b64.assert_called_once_with("12345")
-        self.mock_processor.set_checkpoint.assert_called_once()
-        (tags,), _ = self.mock_processor.set_checkpoint.call_args
-        self.assertIn("direction:in", tags)
-        self.assertIn("type:eventbridge", tags)
-        self.assertIn("topic:MyDetailType", tags)
-        self.assertFalse(any(t.startswith("exchange:") for t in tags))
+        self.mock_checkpoint.assert_called_once()
+        args, kwargs = self.mock_checkpoint.call_args
+        self.assertEqual(args[0], "eventbridge")
+        self.assertEqual(args[1], "MyDetailType")
+        carrier_get = args[2]
+        self.assertEqual(carrier_get("dd-pathway-ctx-base64"), "12345")
+        self.assertEqual(kwargs, {"manual_checkpoint": False})
 
     @patch("datadog_lambda.config.Config.dsm_exchange_name", "my-event-bus")
-    def test_eventbridge_exchange_tag_from_env(self):
+    def test_eventbridge_exchange_name_ignored_until_upstream_support_lands(self):
         event = self._eventbridge_event()
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        (tags,), _ = self.mock_processor.set_checkpoint.call_args
-        self.assertIn("exchange:my-event-bus", tags)
-        self.assertIn("topic:MyDetailType", tags)
-        self.assertIn("type:eventbridge", tags)
+        self.mock_checkpoint.assert_called_once()
+        args, kwargs = self.mock_checkpoint.call_args
+        self.assertEqual(args[0], "eventbridge")
+        self.assertEqual(args[1], "MyDetailType")
+        carrier_get = args[2]
+        self.assertEqual(carrier_get("dd-pathway-ctx-base64"), "12345")
+        self.assertEqual(kwargs, {"manual_checkpoint": False})
 
     def test_eventbridge_no_detail_type_skips_checkpoint(self):
         event = self._eventbridge_event(detail_type=None)
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        self.mock_processor.set_checkpoint.assert_not_called()
+        self.mock_checkpoint.assert_not_called()
 
     def test_eventbridge_no_dd_context_still_checkpoints(self):
         event = {"detail-type": "MyDetailType", "detail": {}}
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        self.mock_processor.decode_pathway_b64.assert_called_once_with(None)
-        self.mock_processor.set_checkpoint.assert_called_once()
+        self.mock_checkpoint.assert_called_once()
+        args, kwargs = self.mock_checkpoint.call_args
+        carrier_get = args[2]
+        self.assertIsNone(carrier_get("dd-pathway-ctx-base64"))
+        self.assertEqual(kwargs, {"manual_checkpoint": False})
 
     def test_eventbridge_missing_detail_still_checkpoints(self):
         event = {"detail-type": "MyDetailType"}
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        self.mock_processor.set_checkpoint.assert_called_once()
+        self.mock_checkpoint.assert_called_once()
+        args, kwargs = self.mock_checkpoint.call_args
+        carrier_get = args[2]
+        self.assertIsNone(carrier_get("dd-pathway-ctx-base64"))
+        self.assertEqual(kwargs, {"manual_checkpoint": False})
 
     @patch("datadog_lambda.config.Config.data_streams_enabled", False)
     def test_eventbridge_data_streams_disabled(self):
         event = self._eventbridge_event()
 
         extract_context_from_eventbridge_event(event, self.lambda_context)
 
-        self.mock_processor.set_checkpoint.assert_not_called()
+        self.mock_checkpoint.assert_not_called()