add sns -> lambda support

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -1,4 +1,5 @@
 import json
+import base64
 
 from ddtrace.internal.logger import get_logger
 from datadog_lambda import logger
@@ -10,6 +11,8 @@
 def set_dsm_context(event, event_source):
     if event_source.equals(EventTypes.SQS):
         _dsm_set_sqs_context(event)
+    elif event_source.equals(EventTypes.SNS):
+        _dsm_set_sns_context(event)
 
 
 def _dsm_set_context_helper(record, service_type, arn, payload_size):
@@ -57,13 +60,37 @@ def _dsm_set_sqs_context(event):
             logger.error(format_err_with_traceback(e))
 
 
+def _dsm_set_sns_context(event):
+    from ddtrace.internal.datastreams.botocore import calculate_sns_payload_size
+
+    def sns_payload_calculator(record, context_json):
+        return calculate_sns_payload_size(record, context_json)
+
+    def sns_arn_extractor(record):
+        sns_data = record.get("Sns")
+        if not sns_data:
+            return ""
+        return sns_data.get("TopicArn", "")
+
+    _dsm_set_context_helper(event, "sns", sns_arn_extractor, sns_payload_calculator)
+
+
 def _get_dsm_context_from_lambda(message):
     """
     Lambda-specific message formats:
         - message.messageAttributes._datadog.stringValue (SQS -> lambda)
+        - message.Sns.MessageAttributes._datadog.Value.decode() (SNS -> lambda)
     """
     context_json = None
-    message_attributes = message.get("messageAttributes")
+    message_body = message
+
+    if "Sns" in message:
+        message_body = message["Sns"]
+
+    message_attributes = message_body.get("MessageAttributes") or message_body.get(
+        "messageAttributes"
+    )
+
     if not message_attributes:
         log.debug("DataStreams skipped lambda message: %r", message)
         return None
@@ -74,7 +101,11 @@ def _get_dsm_context_from_lambda(message):
 
     datadog_attr = message_attributes["_datadog"]
 
-    if "stringValue" in datadog_attr:
+    if message_body.get("Type") == "Notification":
+        # SNS -> lambda notification
+        if datadog_attr.get("Type") == "Binary":
+            context_json = json.loads(base64.b64decode(datadog_attr["Value"]).decode())
+    elif "stringValue" in datadog_attr:
         # SQS -> lambda
         context_json = json.loads(datadog_attr["stringValue"])
     else:

tests/test_dsm.py

@@ -1,10 +1,12 @@
 import unittest
+import base64
 import json
 from unittest.mock import patch, MagicMock
 
 from datadog_lambda.dsm import (
     set_dsm_context,
     _dsm_set_sqs_context,
+    _dsm_set_sns_context,
     _get_dsm_context_from_lambda,
 )
 from datadog_lambda.trigger import EventTypes, _EventSource
@@ -16,6 +18,10 @@ def setUp(self):
         self.mock_dsm_set_sqs_context = patcher.start()
         self.addCleanup(patcher.stop)
 
+        patcher = patch("datadog_lambda.dsm._dsm_set_sns_context")
+        self.mock_dsm_set_sns_context = patcher.start()
+        self.addCleanup(patcher.stop)
+
         patcher = patch("ddtrace.internal.datastreams.data_streams_processor")
         self.mock_data_streams_processor = patcher.start()
         self.addCleanup(patcher.stop)
@@ -32,6 +38,13 @@ def setUp(self):
         self.mock_calculate_sqs_payload_size.return_value = 100
         self.addCleanup(patcher.stop)
 
+        patcher = patch(
+            "ddtrace.internal.datastreams.botocore.calculate_sns_payload_size"
+        )
+        self.mock_calculate_sns_payload_size = patcher.start()
+        self.mock_calculate_sns_payload_size.return_value = 150
+        self.addCleanup(patcher.stop)
+
         patcher = patch("ddtrace.internal.datastreams.processor.DsmPathwayCodec.decode")
         self.mock_dsm_pathway_codec_decode = patcher.start()
         self.addCleanup(patcher.stop)
@@ -116,6 +129,84 @@ def test_sqs_multiple_records_process_each_record(self):
             self.assertIn("type:sqs", tags)
             self.assertEqual(kwargs["payload_size"], 100)
 
+    def test_sns_event_with_no_records_does_nothing(self):
+        """Test that events where Records is None don't trigger DSM processing"""
+        events_with_no_records = [
+            {},
+            {"Records": None},
+            {"someOtherField": "value"},
+        ]
+
+        for event in events_with_no_records:
+            _dsm_set_sns_context(event)
+            self.mock_data_streams_processor.assert_not_called()
+
+    def test_sns_event_triggers_dsm_sns_context(self):
+        """Test that SNS event sources trigger the SNS-specific DSM context function"""
+        sns_event = {
+            "Records": [
+                {
+                    "EventSource": "aws:sns",
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:my-topic",
+                        "Message": "Hello from SNS!",
+                    },
+                }
+            ]
+        }
+
+        event_source = _EventSource(EventTypes.SNS)
+        set_dsm_context(sns_event, event_source)
+
+        self.mock_dsm_set_sns_context.assert_called_once_with(sns_event)
+
+    def test_sns_multiple_records_process_each_record(self):
+        """Test that each record in an SNS event gets processed individually"""
+        multi_record_event = {
+            "Records": [
+                {
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:topic1",
+                        "Message": "Message 1",
+                    }
+                },
+                {
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:topic2",
+                        "Message": "Message 2",
+                    }
+                },
+                {
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:topic3",
+                        "Message": "Message 3",
+                    }
+                },
+            ]
+        }
+
+        mock_context = MagicMock()
+        self.mock_dsm_pathway_codec_decode.return_value = mock_context
+
+        _dsm_set_sns_context(multi_record_event)
+
+        self.assertEqual(mock_context.set_checkpoint.call_count, 3)
+
+        calls = mock_context.set_checkpoint.call_args_list
+        expected_arns = [
+            "arn:aws:sns:us-east-1:123456789012:topic1",
+            "arn:aws:sns:us-east-1:123456789012:topic2",
+            "arn:aws:sns:us-east-1:123456789012:topic3",
+        ]
+
+        for i, call in enumerate(calls):
+            args, kwargs = call
+            tags = args[0]
+            self.assertIn("direction:in", tags)
+            self.assertIn(f"topic:{expected_arns[i]}", tags)
+            self.assertIn("type:sns", tags)
+            self.assertEqual(kwargs["payload_size"], 150)
+
 
 class TestGetDSMContext(unittest.TestCase):
     def test_sqs_to_lambda_string_value_format(self):
@@ -164,6 +255,43 @@ def test_sqs_to_lambda_string_value_format(self):
         assert result["x-datadog-parent-id"] == "321987654"
         assert result["dd-pathway-ctx"] == "test-pathway-ctx"
 
+    def test_sns_to_lambda_format(self):
+        """Test format: message.Sns.MessageAttributes._datadog.Value.decode() (SNS -> lambda)"""
+        trace_context = {
+            "x-datadog-trace-id": "111111111",
+            "x-datadog-parent-id": "222222222",
+            "dd-pathway-ctx": "test-pathway-ctx",
+        }
+        binary_data = base64.b64encode(
+            json.dumps(trace_context).encode("utf-8")
+        ).decode("utf-8")
+
+        sns_lambda_record = {
+            "EventSource": "aws:sns",
+            "EventSubscriptionArn": (
+                "arn:aws:sns:us-east-1:123456789012:sns-topic:12345678-1234-1234-1234-123456789012"
+            ),
+            "Sns": {
+                "Type": "Notification",
+                "MessageId": "95df01b4-ee98-5cb9-9903-4c221d41eb5e",
+                "TopicArn": "arn:aws:sns:us-east-1:123456789012:sns-topic",
+                "Subject": "Test Subject",
+                "Message": "Hello from SNS!",
+                "Timestamp": "2023-01-01T12:00:00.000Z",
+                "MessageAttributes": {
+                    "_datadog": {"Type": "Binary", "Value": binary_data}
+                },
+            },
+        }
+
+        result = _get_dsm_context_from_lambda(sns_lambda_record)
+
+        assert result is not None
+        assert result == trace_context
+        assert result["x-datadog-trace-id"] == "111111111"
+        assert result["x-datadog-parent-id"] == "222222222"
+        assert result["dd-pathway-ctx"] == "test-pathway-ctx"
+
     def test_no_message_attributes(self):
         """Test message without MessageAttributes returns None."""
         message = {