Cover string/b64 msg attribute trace extraction

datadog_lambda/tracing.py

@@ -223,7 +223,25 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
             first_record.get("Sns", {}).get("MessageAttributes", {}),
         )
         dd_payload = msg_attributes.get("_datadog", {})
-        dd_json_data = dd_payload.get("stringValue", dd_payload.get("Value", r"{}"))
+        # SQS uses dataType and binaryValue/stringValue
+        # SNS uses Type and Value
+        dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", ""))
+        if dd_json_data_type == "Binary":
+            dd_json_data = dd_payload.get(
+                "binaryValue",
+                dd_payload.get("Value", r"{}"),
+            )
+            dd_json_data = base64.b64decode(dd_json_data)
+        elif dd_json_data_type == "String":
+            dd_json_data = dd_payload.get(
+                "stringValue",
+                dd_payload.get("Value", r"{}"),
+            )
+        else:
+            logger.debug(
+                "Datadog Lambda Python only supports extracting trace"
+                "context from String or Binary SQS/SNS message attributes"
+            )
         dd_data = json.loads(dd_json_data)
         trace_id = dd_data.get(TraceHeader.TRACE_ID)
         parent_id = dd_data.get(TraceHeader.PARENT_ID)

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 python = ">=3.6.0,<4"
 datadog = "^0.41.0"
 wrapt = "^1.11.2"
-ddtrace = "^0.58.1"
+ddtrace = "^0.59.2"
 importlib_metadata = {version = "^1.0", python = "<3.8"}
 boto3 = { version = "^1.10.33", optional = true }
 typing_extensions = {version = "^4.0", python = "<3.8"}

tests/event_samples/sns-b64-msg-attribute.json

@@ -0,0 +1,27 @@
+{
+    "Records": [
+        {
+            "EventSource": "aws:sns",
+            "EventVersion": "1.0",
+            "EventSubscriptionArn": "arn:aws:sns:sa-east-1:601427279990:serverlessTracingTopicPy:224b60ba-befc-4830-ad96-f1f0ac94eb04",
+            "Sns": {
+                "Type": "Notification",
+                "MessageId": "87056a47-f506-5d77-908b-303605d3b197",
+                "TopicArn": "arn:aws:sns:sa-east-1:601427279990:serverlessTracingTopicPy",
+                "Subject": null,
+                "Message": "Asynchronously invoking a Lambda function with SNS.",
+                "Timestamp": "2022-01-31T14:13:41.637Z",
+                "SignatureVersion": "1",
+                "Signature": "BmwnJb0Ku2KgQef9QOgaSSTwLyUsbkRq90lzD5Vn4mAcRUOq2ForfMOYbxMB6idljWIWy9t/jK4AIMxPGk/eOGiRcENx3BvAcGcoDayBRFY13+xUGaPn5Lfoht/ZJ7/hmCgFWKRa8ooATZL+AwGAw6Id8qzf0R3M3k2asy5Vxa4ODKiFW9OzWY/zFgsYJhddR3JrQl9YOMRyIobNNHT96o1TwjGsSUTEemrxA6jQtb3QbardEKO+2SuataLEZki7gE2D2sA300WqZecumI339q7la+OIj6VDGDwFoppE2sh8hzJYXAH7oo11giwltE0V3/eLFCVhsE8Y1KD/yDPPsA==",
+                "SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-7ff5318490ec183fbaddaa2a969abfda.pem",
+                "UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:601427279990:serverlessTracingTopicPy:224b60ba-befc-4830-ad96-f1f0ac94eb04",
+                "MessageAttributes": {
+                    "_datadog": {
+                        "Type": "Binary",
+                        "Value": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI0OTQ4Mzc3MzE2MzU3MjkxNDIxIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjY3NDY5OTgwMTUwMzc0Mjk1MTIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0="
+                    }
+                }
+            }
+        }
+    ]
+}

tests/integration/snapshots/logs/async-metrics_python36.log

@@ -1114,6 +1114,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/async-metrics_python37.log

@@ -1114,6 +1114,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/async-metrics_python38.log

@@ -1114,6 +1114,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/async-metrics_python39.log

@@ -1114,6 +1114,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/sync-metrics_python36.log

@@ -1185,6 +1185,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/sync-metrics_python37.log

@@ -1185,6 +1185,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/sync-metrics_python38.log

@@ -1185,6 +1185,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/integration/snapshots/logs/sync-metrics_python39.log

@@ -1185,6 +1185,7 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "
           "_inferred_span.tag_source": "self"
         },
         "metrics": {
+          "_dd.agent_psr": 1,
           "system.pid": "XXXX",
           "_sampling_priority_v1": 1,
           "_dd.top_level": 1

tests/test_tracing.py

@@ -253,7 +253,8 @@ def test_with_sqs_distributed_datadog_trace_data(self):
                                     TraceHeader.PARENT_ID: "321",
                                     TraceHeader.SAMPLING_PRIORITY: "1",
                                 }
-                            )
+                            ),
+                            "dataType": "String",
                         }
                     },
                     "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
@@ -796,9 +797,9 @@ def test_create_inferred_span_from_api_gateway_websocket_disconnect_event(self):
         self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self")
         self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync")
 
-    def test_create_inferred_span_from_sqs_event(self):
-        event_sample_source = "sqs"
-        test_file = event_samples + event_sample_source + ".json"
+    def test_create_inferred_span_from_sqs_event_string_msg_attr(self):
+        event_sample_name = "sqs-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)
         ctx = get_mock_context()
@@ -834,9 +835,47 @@ def test_create_inferred_span_from_sqs_event(self):
         self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self")
         self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "async")
 
-    def test_create_inferred_span_from_sns_event(self):
-        event_sample_source = "sns"
-        test_file = event_samples + event_sample_source + ".json"
+    def test_create_inferred_span_from_sns_event_string_msg_attr(self):
+        event_sample_name = "sns-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
+        with open(test_file, "r") as event:
+            event = json.load(event)
+        ctx = get_mock_context()
+        ctx.aws_request_id = "123"
+        span = create_inferred_span(event, ctx)
+        self.assertEqual(span.get_tag("operation_name"), "aws.sns")
+        self.assertEqual(
+            span.service,
+            "sns",
+        )
+        self.assertEqual(
+            span.get_tag("http.url"),
+            None,
+        )
+        self.assertEqual(span.get_tag("endpoint"), None)
+        self.assertEqual(span.get_tag("http.method"), None)
+        self.assertEqual(
+            span.get_tag("resource_names"),
+            "serverlessTracingTopicPy",
+        )
+        self.assertEqual(span.get_tag("topicname"), "serverlessTracingTopicPy")
+        self.assertEqual(
+            span.get_tag("topic_arn"),
+            "arn:aws:sns:sa-east-1:601427279990:serverlessTracingTopicPy",
+        )
+        self.assertEqual(
+            span.get_tag("message_id"), "87056a47-f506-5d77-908b-303605d3b197"
+        )
+        self.assertEqual(span.get_tag("type"), "Notification")
+        self.assertEqual(span.get_tag("subject"), None)
+        self.assertEqual(span.start, 1643638421.637)
+        self.assertEqual(span.span_type, "web")
+        self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self")
+        self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "async")
+
+    def test_create_inferred_span_from_sns_event_b64_msg_attr(self):
+        event_sample_name = "sns-b64-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)
         ctx = get_mock_context()
@@ -1040,8 +1079,8 @@ def test_extract_dd_trace_context_for_eventbridge(self):
         self.assertEqual(context["trace-id"], "12345")
         self.assertEqual(context["parent-id"], "67890")
 
-    def test_extract_context_from_sqs_event(self):
-        event_sample_source = "sqs"
+    def test_extract_context_from_sqs_event_with_string_msg_attr(self):
+        event_sample_source = "sqs-string-msg-attribute"
         test_file = event_samples + event_sample_source + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)
@@ -1062,8 +1101,19 @@ def test_extract_context_from_sqs_batch_event(self):
         self.assertEqual(context["parent-id"], "7431398482019833808")
         self.assertEqual(context["sampling-priority"], "1")
 
-    def test_extract_context_from_sns_event(self):
-        event_sample_source = "sns"
+    def test_extract_context_from_sns_event_with_string_msg_attr(self):
+        event_sample_source = "sns-string-msg-attribute"
+        test_file = event_samples + event_sample_source + ".json"
+        with open(test_file, "r") as event:
+            event = json.load(event)
+        ctx = get_mock_context()
+        context, source = extract_dd_trace_context(event, ctx)
+        self.assertEqual(context["trace-id"], "4948377316357291421")
+        self.assertEqual(context["parent-id"], "6746998015037429512")
+        self.assertEqual(context["sampling-priority"], "1")
+
+    def test_extract_context_from_sns_event_with_b64_msg_attr(self):
+        event_sample_source = "sns-b64-msg-attribute"
         test_file = event_samples + event_sample_source + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)

tests/test_trigger.py

@@ -206,7 +206,8 @@ def test_event_source_s3(self):
 
     def test_event_source_sns(self):
         event_sample_source = "sns"
-        test_file = event_samples + event_sample_source + ".json"
+        event_sample_name = "sns-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)
         ctx = get_mock_context()
@@ -220,7 +221,8 @@ def test_event_source_sns(self):
 
     def test_event_source_sqs(self):
         event_sample_source = "sqs"
-        test_file = event_samples + event_sample_source + ".json"
+        event_sample_name = "sqs-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         with open(test_file, "r") as event:
             event = json.load(event)
         ctx = get_mock_context()
@@ -455,8 +457,8 @@ def test_extract_trigger_tags_s3(self):
         )
 
     def test_extract_trigger_tags_sns(self):
-        event_sample_source = "sns"
-        test_file = event_samples + event_sample_source + ".json"
+        event_sample_name = "sns-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         ctx = get_mock_context()
         with open(test_file, "r") as event:
             event = json.load(event)
@@ -470,8 +472,8 @@ def test_extract_trigger_tags_sns(self):
         )
 
     def test_extract_trigger_tags_sqs(self):
-        event_sample_source = "sqs"
-        test_file = event_samples + event_sample_source + ".json"
+        event_sample_name = "sqs-string-msg-attribute"
+        test_file = event_samples + event_sample_name + ".json"
         ctx = get_mock_context()
         with open(test_file, "r") as event:
             event = json.load(event)