fix

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -35,16 +35,25 @@ def carrier_get(key):
 
 
 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)
+    from ddtrace.data_streams import set_consume_checkpoint
 
-    def sns_arn_extractor(record):
+    records = event.get("Records")
+    if records is None:
+        return
+    for record in records:
         sns_data = record.get("Sns")
         if not sns_data:
-            return ""
-        return sns_data.get("TopicArn", "")
+            return
+        arn = sns_data.get("TopicArn", "")
+        context_json = _get_dsm_context_from_lambda(record)
+        if not context_json:
+            logger.debug("DataStreams skipped lambda message: %r", record)
+            return None
+
+        def carrier_get(key):
+            return context_json.get(key)
+
+        set_consume_checkpoint("sns", arn, carrier_get)
 
 
 def _get_dsm_context_from_lambda(message):