fix: detect sqs carriers per record in eventbridge batches

Author: jeastham1993

datadog_lambda/tracing.py

@@ -290,9 +290,7 @@ def extract_context_from_sqs_or_sns_event_or_context(
 
     # EventBridge => SQS
     try:
-        context, is_eventbridge_sqs = _extract_context_from_eventbridge_sqs_event(
-            event
-        )
+        context, is_eventbridge_sqs = _extract_context_from_eventbridge_sqs_event(event)
         if is_eventbridge_sqs:
             if _is_context_complete(context):
                 return context
@@ -400,37 +398,25 @@ def _extract_context_from_eventbridge_sqs_event(event):
         return None, False
 
     first_record = records[0]
-    body_str = first_record.get("body")
-    body = json.loads(body_str)
-    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")
-    ):
+    dd_context, is_eventbridge_sqs = _extract_eventbridge_sqs_record_context(
+        first_record
+    )
+    if not is_eventbridge_sqs:
         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_detail = record_body.get("detail")
-                record_context = (
-                    record_detail.get("_datadog")
-                    if isinstance(record_detail, dict)
-                    else None
+                record_context, is_eventbridge_record = (
+                    _extract_eventbridge_sqs_record_context(record)
                 )
+                if not is_eventbridge_record:
+                    record_context = _extract_sqs_record_message_attribute_context(
+                        record
+                    )
                 _dsm_set_checkpoint(
                     record_context, "sqs", record.get("eventSourceARN", "")
                 )
@@ -450,6 +436,46 @@ def _extract_context_from_eventbridge_sqs_event(event):
     return propagator.extract(dd_context), True
 
 
+def _extract_eventbridge_sqs_record_context(record):
+    body_str = record.get("body")
+    body = json.loads(body_str)
+    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):
     """
     Extract datadog trace context from an EventBridge message's Details.

tests/test_tracing.py

@@ -3703,9 +3703,7 @@ def _eventbridge_sqs_record(queue_arn, pathway_ctx, include_trace_headers=True):
         body = {
             "detail-type": "MyDetailType",
             "source": "my.event.source",
-            "detail": {
-                "_datadog": dd_context
-            },
+            "detail": {"_datadog": dd_context},
         }
         return {
             "eventSourceARN": queue_arn,
@@ -3751,6 +3749,44 @@ 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")
 
+    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")
+
     @patch(
         "datadog_lambda.tracing.extract_context_from_lambda_context",
         return_value=Context(trace_id=111, span_id=222, sampling_priority=1),