diff --git a/README.md b/README.md index a1199e8e..fc6c7f2c 100644 --- a/README.md +++ b/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.
For example, given the input payload: 
{
  "lv1" : {
    "lv2": {
      "lv3": "val"
    }
  }
}
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\"}`.
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 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..74168649 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,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) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index a74e0139..74253916 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,497 @@ 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_mixed_batch_with_sns_record_keeps_sns_context(self): + # EventBridge and SNS can both be routed to the same SQS queue and then + # batched into a single event payload. The SNS record's DSM context + # lives in the SNS envelope inside `body` (SNS => SQS subscription + # without raw message delivery), not in `messageAttributes`, and must + # still be extracted even though the batch contains an EventBridge + # delivery. + arn1 = "arn:aws:sqs:us-east-1:123456789012:eb-queue" + arn2 = "arn:aws:sqs:us-east-1:123456789012:sns-queue" + sns_dd_data = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": "sns-ctx", + } + sns_body = { + "Type": "Notification", + "TopicArn": "arn:aws:sns:us-east-1:123456789012:my-topic", + "Message": "hello from sns", + "MessageAttributes": { + "_datadog": { + "Type": "String", + "Value": json.dumps(sns_dd_data), + } + }, + } + event = { + "Records": [ + self._eventbridge_sqs_record(arn1, "eb-ctx"), + { + "eventSourceARN": arn2, + "eventSource": "aws:sqs", + "body": json.dumps(sns_body), + }, + ] + } + + 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") + # The SNS record is consumed from the SQS queue, so it keeps SQS + # conventions, but its carrier must come from the SNS envelope. + self.assertEqual((second_args[0], second_args[1]), ("sqs", arn2)) + self.assertEqual(second_args[2]("dd-pathway-ctx-base64"), "sns-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))