fix

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -15,9 +15,7 @@ def set_dsm_context(event, event_source):
         _dsm_set_sns_context(event)
 
 
-def _dsm_set_context_helper(
-    record, service_type, arn, payload_size
-):
+def _dsm_set_context_helper(record, service_type, arn, payload_size):
     """
     Common helper function for setting DSM context.
 
@@ -64,17 +62,22 @@ def _dsm_set_sqs_context(event):
 
 def _dsm_set_sns_context(event):
     from ddtrace.internal.datastreams.botocore import calculate_sns_payload_size
+    from datadog_lambda.wrapper import format_err_with_traceback
 
-    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", "")
+    records = event.get("Records")
+    if records is None:
+        return
 
-    _dsm_set_context_helper(event, "sns", sns_arn_extractor, sns_payload_calculator)
+    for record in records:
+        try:
+            sns_data = record.get("Sns")
+            if not sns_data:
+                return ""
+            arn = sns_data.get("TopicArn", "")
+            payload_size = calculate_sns_payload_size(record)
+            _dsm_set_context_helper(record, "sns", arn, payload_size)
+        except Exception as e:
+            logger.error(format_err_with_traceback(e))
 
 
 def _get_dsm_context_from_lambda(message):