fixes

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -59,7 +59,7 @@ def carrier_get(key):
 
 
 def _dsm_set_kinesis_context(event):
-    from ddtrace.internal.datastreams.botocore import calculate_kinesis_payload_size
+    from ddtrace.data_streams import set_consume_checkpoint
 
     records = event.get("Records")
     if records is None:
@@ -68,9 +68,13 @@ def _dsm_set_kinesis_context(event):
     for record in records:
         arn = record.get("eventSourceARN", "")
         context_json = _get_dsm_context_from_lambda(record)
-        payload_size = calculate_kinesis_payload_size(record, context_json)
+        if not context_json:
+            logger.debug("DataStreams skipped lambda message: %r", record)
+            return None
 
-        _dsm_set_context_helper("kinesis", arn, payload_size, context_json)
+        def carrier_get(key):
+            return context_json.get(key)
+        set_consume_checkpoint("kinesis", arn, carrier_get)
 
 
 def _get_dsm_context_from_lambda(message):