fix(datastreams): support datastream propagation with raw message delivery enabled [backport 1.17] (#6799)

Backport b2bd70f from #6735 to 1.17.

When a subscription is set up between SNS and a consumer (such as SQS),
there is an option to enable "Raw Message Delivery" which changes the
format of delivered data, primarily the MessageAttributes and
MessageBody
sections.

Prior to this change, datastream context was dropped/ignored on the
receive because the code assumed the format of delivered messages when
Raw Delivery was disabled.

After this change, the patch handles both raw message delivery and
non-raw delivery.


## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Co-authored-by: Teague Bick <[email protected]>

Author: github-actions[bot]

ddtrace/contrib/botocore/patch.py

@@ -93,6 +93,14 @@ class TraceInjectionDecodingError(Exception):
     pass
 
 
+def _encode_data(trace_data):
+    """
+    This method exists solely to enable us to patch the value in tests, since
+    moto doesn't support auto-encoded SNS -> SQS as binary with RawDelivery enabled
+    """
+    return json.dumps(trace_data)
+
+
 def inject_trace_data_to_message_attributes(trace_data, entry, endpoint_service=None):
     # type: (Dict[str, str], Dict[str, Any], Optional[str]) -> None
     """
@@ -109,12 +117,12 @@ def inject_trace_data_to_message_attributes(trace_data, entry, endpoint_service=
         if endpoint_service == "sqs":
             # Use String since changing this to Binary would be a breaking
             # change as other tracers expect this to be a String.
-            entry["MessageAttributes"]["_datadog"] = {"DataType": "String", "StringValue": json.dumps(trace_data)}
+            entry["MessageAttributes"]["_datadog"] = {"DataType": "String", "StringValue": _encode_data(trace_data)}
         elif endpoint_service == "sns":
             # Use Binary since SNS subscription filter policies fail silently
             # with JSON strings https://github.com/DataDog/datadog-lambda-js/pull/269
             # AWS will encode our value if it sees "Binary"
-            entry["MessageAttributes"]["_datadog"] = {"DataType": "Binary", "BinaryValue": json.dumps(trace_data)}
+            entry["MessageAttributes"]["_datadog"] = {"DataType": "Binary", "BinaryValue": _encode_data(trace_data)}
         else:
             log.warning("skipping trace injection, endpoint is not SNS or SQS")
     else:
@@ -589,6 +597,18 @@ def patched_api_call(original_func, instance, args, kwargs):
                                 # The message originated from SQS
                                 message_body = message
                                 context_json = json.loads(message_body["MessageAttributes"]["_datadog"]["StringValue"])
+                            elif (
+                                "MessageAttributes" in message
+                                and "_datadog" in message["MessageAttributes"]
+                                and "BinaryValue" in message["MessageAttributes"]["_datadog"]
+                            ):
+                                # Raw message delivery
+                                message_body = message
+                                context_json = json.loads(
+                                    message_body["MessageAttributes"]["_datadog"]["BinaryValue"].decode()
+                                )
+                            else:
+                                log.debug("DataStreams did not handle message: %r", message)
 
                             pathway = context_json.get(PROPAGATION_KEY_BASE_64, None) if context_json else None
                             ctx = pin.tracer.data_streams_processor.decode_pathway_b64(pathway)

releasenotes/notes/datastreams-fix-sns-raw-delivery-044e514a84e132be.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    data_streams: This fix resolves an issue where data stream context propagation would not propagate via SNS if raw
+    message delivery was enabled.

tests/contrib/botocore/test.py

@@ -2,6 +2,7 @@
 import datetime
 import io
 import json
+import sys
 import unittest
 import zipfile
 
@@ -1087,10 +1088,29 @@ def test_double_patch(self):
         assert spans
         assert len(spans) == 1
 
-    @mock_sns
-    @mock_sqs
     @TracerTestCase.run_in_subprocess(env_overrides=dict(DD_DATA_STREAMS_ENABLED="True"))
     def test_data_streams_sns_to_sqs(self):
+        self._test_data_streams_sns_to_sqs(False)
+
+    @mock.patch.object(sys.modules["ddtrace.contrib.botocore.patch"], "_encode_data")
+    @TracerTestCase.run_in_subprocess(env_overrides=dict(DD_DATA_STREAMS_ENABLED="True"))
+    def test_data_streams_sns_to_sqs_raw_delivery(self, mock_encode):
+        """
+        Moto doesn't currently handle raw delivery message handling quite correctly.
+        In the real world, AWS will encode data for us. Moto does not.
+
+        So, we patch our code here to encode the data
+        """
+
+        def _moto_compatible_encode(trace_data):
+            return base64.b64encode(json.dumps(trace_data).encode("utf-8"))
+
+        mock_encode.side_effect = _moto_compatible_encode
+        self._test_data_streams_sns_to_sqs(True)
+
+    @mock_sns
+    @mock_sqs
+    def _test_data_streams_sns_to_sqs(self, use_raw_delivery):
         # DEV: We want to mock time to ensure we only create a single bucket
         with mock.patch("time.time") as mt:
             mt.return_value = 1642544540
@@ -1104,7 +1124,14 @@ def test_data_streams_sns_to_sqs(self):
                 sqs_url = self.sqs_test_queue["QueueUrl"]
                 url_parts = sqs_url.split("/")
                 sqs_arn = "arn:aws:sqs:{}:{}:{}".format("us-east-1", url_parts[-2], url_parts[-1])
-                sns.subscribe(TopicArn=topic_arn, Protocol="sqs", Endpoint=sqs_arn)
+                subscription = sns.subscribe(TopicArn=topic_arn, Protocol="sqs", Endpoint=sqs_arn)
+
+                if use_raw_delivery:
+                    sns.set_subscription_attributes(
+                        SubscriptionArn=subscription["SubscriptionArn"],
+                        AttributeName="RawMessageDelivery",
+                        AttributeValue="true",
+                    )
 
                 Pin.get_from(sns).clone(tracer=self.tracer).onto(sns)
                 Pin.get_from(self.sqs_client).clone(tracer=self.tracer).onto(self.sqs_client)