reworked to remove ddtrace dependencies

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -1,6 +1,5 @@
-import json
 import logging
-
+import json
 from datadog_lambda.trigger import EventTypes
 
 logger = logging.getLogger(__name__)
@@ -11,33 +10,8 @@ def set_dsm_context(event, event_source):
         _dsm_set_sqs_context(event)
 
 
-def _dsm_set_context_helper(service_type, arn, payload_size, context_json):
-    """
-    Common helper function for setting DSM context.
-
-    Args:
-        service_type: The service type string (example: sqs', 'sns')
-        arn:  ARN from the record
-        payload_size: payload size of the record
-        context_json: Datadog context for the record
-    """
-    from ddtrace.internal.datastreams import data_streams_processor
-    from ddtrace.internal.datastreams.processor import DsmPathwayCodec
-
-    processor = data_streams_processor()
-
-    try:
-        ctx = DsmPathwayCodec.decode(context_json, processor)
-        ctx.set_checkpoint(
-            ["direction:in", f"topic:{arn}", f"type:{service_type}"],
-            payload_size=payload_size,
-        )
-    except Exception as e:
-        logger.error(f"Unable to set dsm context: {e}")
-
-
 def _dsm_set_sqs_context(event):
-    from ddtrace.internal.datastreams.botocore import calculate_sqs_payload_size
+    from ddtrace.data_streams import set_consume_checkpoint
 
     records = event.get("Records")
     if records is None:
@@ -46,9 +20,14 @@ def _dsm_set_sqs_context(event):
     for record in records:
         arn = record.get("eventSourceARN", "")
         context_json = _get_dsm_context_from_lambda(record)
-        payload_size = calculate_sqs_payload_size(record, context_json)
+        if not context_json:
+            logger.debug("DataStreams skipped lambda message: %r", record)
+            return None
+
+        def carrier_get(key):
+            return context_json.get(key)
 
-        _dsm_set_context_helper("sqs", arn, payload_size, context_json)
+        set_consume_checkpoint("sqs", arn, carrier_get)
 
 
 def _get_dsm_context_from_lambda(message):