fixes, kinesis support, and renaming some functions

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -224,7 +224,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         logger.debug("Failed extracting context as EventBridge to SQS.")
 
     try:
-        dd_json_data, _ = extract_dd_json_data_from_message_attributes(event)
+        dd_json_data, _ = extract_dd_context_from_sqs_or_sns_event(event)
         if dd_json_data:
             dd_data = json.loads(dd_json_data)
 
@@ -264,7 +264,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
-def extract_dd_json_data_from_message_attributes(event):
+def extract_dd_context_from_sqs_or_sns_event(event):
     first_record = event.get("Records")[0]
     arn = first_record.get("eventSourceARN", "")
     is_sqs = bool(arn)
@@ -366,28 +366,34 @@ def extract_context_from_kinesis_event(event, lambda_context):
     Extract datadog trace context from a Kinesis Stream's base64 encoded data string
     """
     try:
-        record = get_first_record(event)
-        kinesis = record.get("kinesis")
-        if not kinesis:
-            return extract_context_from_lambda_context(lambda_context)
-        data = kinesis.get("data")
-        if data:
-            import base64
-
-            b64_bytes = data.encode("ascii")
-            str_bytes = base64.b64decode(b64_bytes)
-            data_str = str_bytes.decode("ascii")
-            data_obj = json.loads(data_str)
-            dd_ctx = data_obj.get("_datadog")
-            if dd_ctx:
-                context = propagator.extract(dd_ctx)
-                return context
+        dd_ctx, _ = extract_dd_context_from_kinesis_event(event, lambda_context)
+        if dd_ctx:
+            return propagator.extract(dd_ctx)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
 
     return extract_context_from_lambda_context(lambda_context)
 
 
+def extract_dd_context_from_kinesis_event(event, lambda_context):
+    record = get_first_record(event)
+    arn = record.get("eventSourceARN", "")
+    kinesis = record.get("kinesis")
+    if not kinesis:
+        return extract_context_from_lambda_context(lambda_context)
+    data = kinesis.get("data")
+    if data:
+        import base64
+
+        b64_bytes = data.encode("ascii")
+        str_bytes = base64.b64decode(b64_bytes)
+        data_str = str_bytes.decode("ascii")
+        data_obj = json.loads(data_str)
+        dd_ctx = data_obj.get("_datadog")
+        return dd_ctx, arn
+    return None, None
+
+
 def _deterministic_sha256_hash(s: str, part: str) -> int:
     import hashlib
 

datadog_lambda/wrapper.py

@@ -240,20 +240,33 @@ def _before(self, event, context):
                         event, context, event_source, config.decode_authorizer_context
                     )
                 if config.data_streams_enabled:
+                    from datadog_lambda.trigger import EventTypes
+
                     from datadog_lambda.tracing import (
-                        extract_dd_json_data_from_message_attributes,
+                        extract_dd_context_from_sqs_or_sns_event,
+                        extract_dd_context_from_kinesis_event,
                     )
 
+                    dd_json_data = None
+                    arn = None
+
                     try:
-                        (
-                            dd_json_data,
-                            arn,
-                        ) = extract_dd_json_data_from_message_attributes(event)
+                        if event_source.equals(EventTypes.SQS) or event_source.equals(
+                            EventTypes.SNS
+                        ):
+                            (
+                                dd_json_data,
+                                arn,
+                            ) = extract_dd_context_from_sqs_or_sns_event(event)
+                        elif event_source.equals(EventTypes.KINESIS):
+                            dd_json_data, arn = extract_dd_context_from_kinesis_event(
+                                event, context
+                            )
                     except Exception as e:
                         logger.debug(f"Failed to extract DSM checkpoint: {e}")
 
                     if dd_json_data:
-                        set_dsm_checkpoint(dd_json_data, event_source, arn)
+                        set_dsm_checkpoint(dd_json_data, event_source.to_string(), arn)
 
                 self.span = create_function_execution_span(
                     context=context,
@@ -363,7 +376,8 @@ def set_dsm_checkpoint(dd_json_data, event_source, arn):
     from ddtrace.data_streams import set_consume_checkpoint
 
     try:
-        dd_json_data = json.loads(dd_json_data)
+        if type(dd_json_data) is not dict:
+            dd_json_data = json.loads(dd_json_data)
         if "dd-pathway-ctx-base64" not in dd_json_data:
             return
 
@@ -374,9 +388,7 @@ def carrier_get(key):
             return carrier_get
 
         carrier_get = create_carrier_get(dd_json_data)
-        set_consume_checkpoint(
-            event_source.to_string(), arn, carrier_get, manual_checkpoint=False
-        )
+        set_consume_checkpoint(event_source, arn, carrier_get, manual_checkpoint=False)
     except Exception as e:
         logger.debug(f"Failed to set DSM checkpoint: {e}")
 

tests/test_tracing.py

@@ -1,4 +1,3 @@
-import base64
 import copy
 import functools
 import json