diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py index b92c37ee..6487c9e7 100644 --- a/datadog_lambda/config.py +++ b/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) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 5c094027..d6b6d707 100644 --- a/datadog_lambda/tracing.py +++ b/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) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index a74e0139..17c924a7 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1,6 +1,8 @@ import base64 +import contextlib import copy import functools +import inspect import json import traceback import pytest @@ -44,14 +46,15 @@ propagator, emit_telemetry_on_exception_outside_of_handler, _dsm_set_checkpoint, + _dsm_set_eventbridge_checkpoint, extract_context_from_kinesis_event, extract_context_from_sqs_or_sns_event_or_context, + extract_context_from_eventbridge_event, ) from datadog_lambda.trigger import parse_event_source from tests.utils import get_mock_context, ClientContext - function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test" fake_xray_header_value = ( @@ -3087,6 +3090,43 @@ def test_sqs_no_datadog_message_attribute(self): # None indicates no DSM context propagation self.assertEqual(carrier_get("dd-pathway-ctx-base64"), None) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + def test_sqs_detail_body_still_uses_message_attributes(self, mock_extract_context): + dd_data = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": "attr-ctx", + } + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "String", "stringValue": dd_json_data} + }, + "eventSource": "aws:sqs", + "body": json.dumps({"detail": {"application": "payload"}}), + } + ] + } + + context = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + mock_extract_context.assert_not_called() + self.assertEqual(context.trace_id, 12345) + self.assertEqual(context.span_id, 67890) + self.assertEqual(context.sampling_priority, 1) + self.assertEqual(self.mock_checkpoint.call_count, 1) + args, _ = self.mock_checkpoint.call_args + self.assertEqual(args[0], "sqs") + self.assertEqual(args[1], "arn:aws:sqs:us-east-1:123456789012:test-queue") + self.assertEqual(args[2]("dd-pathway-ctx-base64"), "attr-ctx") + def test_sqs_empty_datadog_message_attribute(self): event = { "Records": [ @@ -3919,3 +3959,446 @@ def test_kinesis_data_streams_disabled(self): arn = "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream" _dsm_set_checkpoint(context_json, event_type, arn) + + # EVENTBRIDGE -> SQS TESTS + + @staticmethod + 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", + } + ) + body = { + "detail-type": "MyDetailType", + "source": "my.event.source", + "detail": {"_datadog": dd_context}, + } + return { + "eventSourceARN": queue_arn, + "eventSource": "aws:sqs", + "body": json.dumps(body), + } + + def test_eventbridge_sqs_context_propagated(self): + queue_arn = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + event = {"Records": [self._eventbridge_sqs_record(queue_arn, "12345")]} + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + # EventBridge -> SQS is consumed from the queue, so it uses SQS tags. + 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) + carrier_get = args[2] + self.assertEqual(carrier_get("dd-pathway-ctx-base64"), "12345") + + def test_eventbridge_sqs_checkpoints_all_records(self): + arn1 = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + arn2 = "arn:aws:sqs:us-east-1:123456789012:eb-queue-2" + event = { + "Records": [ + self._eventbridge_sqs_record(arn1, "ctx-1"), + self._eventbridge_sqs_record(arn2, "ctx-2"), + ] + } + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + self.assertEqual(self.mock_checkpoint.call_count, 2) + first_args, _ = self.mock_checkpoint.call_args_list[0] + second_args, _ = self.mock_checkpoint.call_args_list[1] + self.assertEqual((first_args[0], first_args[1]), ("sqs", arn1)) + self.assertEqual(first_args[2]("dd-pathway-ctx-base64"), "ctx-1") + self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2)) + self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "ctx-2") + + def test_eventbridge_sqs_mixed_batch_uses_per_record_carriers(self): + arn1 = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + arn2 = "arn:aws:sqs:us-east-1:123456789012:direct-queue" + second_dd_data = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": "sqs-ctx", + } + event = { + "Records": [ + self._eventbridge_sqs_record(arn1, "eb-ctx"), + { + "eventSourceARN": arn2, + "eventSource": "aws:sqs", + "body": json.dumps({"message": "direct sqs payload"}), + "messageAttributes": { + "_datadog": { + "dataType": "String", + "stringValue": json.dumps(second_dd_data), + } + }, + }, + ] + } + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + self.assertEqual(self.mock_checkpoint.call_count, 2) + first_args, _ = self.mock_checkpoint.call_args_list[0] + second_args, _ = self.mock_checkpoint.call_args_list[1] + self.assertEqual((first_args[0], first_args[1]), ("sqs", arn1)) + self.assertEqual(first_args[2]("dd-pathway-ctx-base64"), "eb-ctx") + self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2)) + self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "sqs-ctx") + + def test_eventbridge_sqs_direct_first_record_still_checkpoints_eventbridge(self): + # First record is a direct SQS send, a later record is an EventBridge + # delivery. Every record must still be classified and checkpointed + # independently, and the first record must not be double counted. + arn1 = "arn:aws:sqs:us-east-1:123456789012:direct-queue" + arn2 = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + first_dd_data = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": "sqs-ctx", + } + event = { + "Records": [ + { + "eventSourceARN": arn1, + "eventSource": "aws:sqs", + "body": json.dumps({"message": "direct sqs payload"}), + "messageAttributes": { + "_datadog": { + "dataType": "String", + "stringValue": json.dumps(first_dd_data), + } + }, + }, + self._eventbridge_sqs_record(arn2, "eb-ctx"), + ] + } + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + self.assertEqual(self.mock_checkpoint.call_count, 2) + first_args, _ = self.mock_checkpoint.call_args_list[0] + second_args, _ = self.mock_checkpoint.call_args_list[1] + self.assertEqual((first_args[0], first_args[1]), ("sqs", arn1)) + self.assertEqual(first_args[2]("dd-pathway-ctx-base64"), "sqs-ctx") + self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2)) + self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "eb-ctx") + + def test_eventbridge_sqs_non_json_body_falls_back_to_sqs_attributes(self): + # A later direct SQS record can legitimately carry a non-JSON body + # while its DSM carrier lives in messageAttributes._datadog. The + # unparseable body must not prevent that record's checkpoint. + arn1 = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + arn2 = "arn:aws:sqs:us-east-1:123456789012:direct-queue" + second_dd_data = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": "sqs-ctx", + } + event = { + "Records": [ + self._eventbridge_sqs_record(arn1, "eb-ctx"), + { + "eventSourceARN": arn2, + "eventSource": "aws:sqs", + "body": "plain text, not json", + "messageAttributes": { + "_datadog": { + "dataType": "String", + "stringValue": json.dumps(second_dd_data), + } + }, + }, + ] + } + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + self.assertEqual(self.mock_checkpoint.call_count, 2) + first_args, _ = self.mock_checkpoint.call_args_list[0] + second_args, _ = self.mock_checkpoint.call_args_list[1] + self.assertEqual((first_args[0], first_args[1]), ("sqs", arn1)) + self.assertEqual(first_args[2]("dd-pathway-ctx-base64"), "eb-ctx") + self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2)) + self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "sqs-ctx") + + @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" + event = {"Records": [self._eventbridge_sqs_record(queue_arn, "12345")]} + + extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + self.mock_checkpoint.assert_not_called() + + +class TestEventBridgeDSMLogic(unittest.TestCase): + def setUp(self): + self.lambda_context = get_mock_context() + 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 + ) + config_patcher.start() + self.addCleanup(config_patcher.stop) + + @staticmethod + def _eventbridge_event(detail_type="MyDetailType", pathway_ctx="12345"): + return { + "detail-type": detail_type, + "source": "my.event.source", + "detail": {"_datadog": {"dd-pathway-ctx-base64": pathway_ctx}}, + } + + def test_eventbridge_context_propagated(self): + event = self._eventbridge_event() + + extract_context_from_eventbridge_event(event, self.lambda_context) + + 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, "tags": None}) + + @patch("datadog_lambda.config.Config.dsm_exchange_name", "my-event-bus") + def test_eventbridge_exchange_name_used_when_upstream_support_present(self): + event = self._eventbridge_event() + + extract_context_from_eventbridge_event(event, self.lambda_context) + + 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, "tags": ["exchange:my-event-bus"]} + ) + + 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_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_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, "tags": None}) + + def test_eventbridge_missing_detail_still_checkpoints(self): + event = {"detail-type": "MyDetailType"} + + extract_context_from_eventbridge_event(event, self.lambda_context) + + 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, "tags": None}) + + @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_checkpoint.assert_not_called() + + +def _real_set_consume_checkpoint_supports_tags(): + """True when the installed ddtrace `set_consume_checkpoint` exposes the + `tags` parameter (ddtrace >= 4.14).""" + from ddtrace.data_streams import set_consume_checkpoint + + return "tags" in inspect.signature(set_consume_checkpoint).parameters + + +@contextlib.contextmanager +def _stub_data_streams_processor(mock_processor): + """Drive the *real* `ddtrace.data_streams.set_consume_checkpoint` while + stubbing only the DSM processor via a public (non-`.internal`) seam. + + The public accessor differs by ddtrace version: + + * ddtrace >= 4.x exposes the module-level `ddtrace.data_streams. + data_streams_processor()` factory, which `set_consume_checkpoint` + calls directly. + * ddtrace 3.19.x reads `ddtrace.tracer.data_streams_processor` instead + (populated at startup when DSM is enabled). + + Patching whichever seam the running version uses keeps this test honest: + the real `set_consume_checkpoint` executes, so any upstream change to its + signature or behaviour surfaces here instead of being masked by a mock. + """ + import ddtrace + from ddtrace import data_streams + + ds_enabled = ddtrace.config._data_streams_enabled + ddtrace.config._data_streams_enabled = True + try: + if hasattr(data_streams, "data_streams_processor"): + with patch( + "ddtrace.data_streams.data_streams_processor", + return_value=mock_processor, + ): + yield + else: + with patch( + "ddtrace.tracer.data_streams_processor", + mock_processor, + create=True, + ): + yield + finally: + ddtrace.config._data_streams_enabled = ds_enabled + + +class TestDSMRealDdtraceApi(unittest.TestCase): + """Guard tests that exercise the real ddtrace `set_consume_checkpoint` + public API without mocking it, so a signature/behaviour change upstream + fails CI instead of passing silently against a permissive mock. + """ + + def setUp(self): + config_patcher = patch( + "datadog_lambda.config.Config.data_streams_enabled", True + ) + config_patcher.start() + self.addCleanup(config_patcher.stop) + + def _edge_tags(self, mock_processor): + self.assertTrue( + mock_processor.set_checkpoint.called, + "expected real set_consume_checkpoint to reach processor.set_checkpoint", + ) + return mock_processor.set_checkpoint.call_args.args[0] + + def test_real_api_sqs_checkpoint(self): + mock_processor = Mock() + context_json = {"dd-pathway-ctx-base64": "sqs-ctx"} + arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" + + with _stub_data_streams_processor(mock_processor): + _dsm_set_checkpoint(context_json, "sqs", arn) + + edge_tags = self._edge_tags(mock_processor) + self.assertIn("type:sqs", edge_tags) + self.assertIn("topic:" + arn, edge_tags) + self.assertIn("direction:in", edge_tags) + # manual_checkpoint=False must not emit the manual_checkpoint tag. + self.assertNotIn("manual_checkpoint:true", edge_tags) + + def test_real_api_eventbridge_checkpoint(self): + mock_processor = Mock() + context_json = {"dd-pathway-ctx-base64": "eb-ctx"} + + with _stub_data_streams_processor(mock_processor): + _dsm_set_eventbridge_checkpoint(context_json, "MyDetailType") + + edge_tags = self._edge_tags(mock_processor) + self.assertIn("type:eventbridge", edge_tags) + self.assertIn("topic:MyDetailType", edge_tags) + self.assertIn("direction:in", edge_tags) + self.assertNotIn("manual_checkpoint:true", edge_tags) + + @unittest.skipUnless( + _real_set_consume_checkpoint_supports_tags(), + "installed ddtrace has no `tags` parameter (< 4.14)", + ) + @patch("datadog_lambda.config.Config.dsm_exchange_name", "my-event-bus") + def test_real_api_eventbridge_exchange_tag(self): + # Exercises the ddtrace >= 4.14 `tags` functionality end-to-end: the + # exchange edge tag must actually reach the processor. + mock_processor = Mock() + context_json = {"dd-pathway-ctx-base64": "eb-ctx"} + + with _stub_data_streams_processor(mock_processor): + _dsm_set_eventbridge_checkpoint(context_json, "MyDetailType") + + edge_tags = self._edge_tags(mock_processor) + self.assertIn("type:eventbridge", edge_tags) + self.assertIn("topic:MyDetailType", edge_tags) + self.assertIn("exchange:my-event-bus", edge_tags) + + @unittest.skipUnless( + _real_set_consume_checkpoint_supports_tags(), + "installed ddtrace has no `tags` parameter (< 4.14)", + ) + def test_real_api_eventbridge_no_exchange_tag_when_unset(self): + # With no DD_DSM_EXCHANGE_NAME configured, no exchange tag is emitted + # even on ddtrace versions that support `tags`. + mock_processor = Mock() + context_json = {"dd-pathway-ctx-base64": "eb-ctx"} + + with _stub_data_streams_processor(mock_processor): + _dsm_set_eventbridge_checkpoint(context_json, "MyDetailType") + + edge_tags = self._edge_tags(mock_processor) + self.assertFalse(any(t.startswith("exchange:") for t in edge_tags))