fix: extract SNS dsm context in mixed eventbridge/SNS sqs batches

When EventBridge and SNS are both routed to the same SQS queue and their
messages are batched into a single event payload, the SNS DSM context was
dropped.

_dsm_set_eventbridge_sqs_batch_checkpoints classifies every record in the
batch and falls back to _extract_sqs_record_message_attribute_context for
non-EventBridge records. That helper only read record.messageAttributes and
never unwrapped an SNS notification carried in the SQS body (an SNS => SQS
subscription without raw message delivery), so the SNS carrier resolved to
None. Because the batch contained an EventBridge delivery, dsm_handled was
True for the whole batch and the SNS-aware fallback path further down in
extract_context_from_sqs_or_sns_event_or_context never ran.

Teach the helper to unwrap the SNS envelope the same way the main SQS/SNS
extraction path already does, so each record in a mixed batch is classified
correctly regardless of whether it is a direct SQS send or an SNS delivery.
The shared attribute decoding and envelope detection are factored out into
_decode_dd_message_attribute and _parse_sns_notification_from_sqs_body.

Author: jeastham1993

datadog_lambda/tracing.py

@@ -526,8 +526,51 @@ def _extract_eventbridge_sqs_record_context(record):
 
 
 def _extract_sqs_record_message_attribute_context(record):
-    msg_attributes = record.get("messageAttributes") or {}
-    dd_payload = msg_attributes.get("_datadog")
+    """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 msg_attributes is None:
+        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
 

tests/test_tracing.py

@@ -4103,6 +4103,57 @@ def test_eventbridge_sqs_direct_first_record_still_checkpoints_eventbridge(self)
         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