fixes

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -81,29 +81,36 @@ def _dsm_set_sns_context(event):
 def _dsm_set_kinesis_context(event):
     from ddtrace.internal.datastreams.botocore import calculate_kinesis_payload_size
 
-    def kinesis_payload_calculator(record, context_json):
-        return calculate_kinesis_payload_size(record)
-
-    def kinesis_arn_extractor(record):
-        arn = record.get("eventSourceARN")
-        if arn is None:
-            return ""
-        return arn
+    records = event.get("Records")
+    if records is None:
+        return
 
-    _dsm_set_context_helper(
-        event, "kinesis", kinesis_arn_extractor, kinesis_payload_calculator
-    )
+    for record in records:
+        arn = record.get("eventSourceARN", "")
+        payload_size = calculate_kinesis_payload_size(record)
+        _dsm_set_context_helper(record, "kinesis", arn, payload_size)
 
 
 def _get_dsm_context_from_lambda(message):
     """
     Lambda-specific message formats:
         - message.messageAttributes._datadog.stringValue (SQS -> lambda)
         - message.Sns.MessageAttributes._datadog.Value.decode() (SNS -> lambda)
+        - message.kinesis.data.decode()._datadog (Kinesis -> lambda)
     """
     context_json = None
     message_body = message
 
+    if "kinesis" in message:
+        try:
+            kinesis_data = json.loads(
+                base64.b64decode(message["kinesis"]["data"]).decode()
+            )
+            return kinesis_data.get("_datadog")
+        except (ValueError, TypeError, KeyError):
+            log.debug("Unable to parse kinesis data for lambda message")
+            return None
+
     if "Sns" in message:
         message_body = message["Sns"]
 

tests/test_dsm.py

@@ -386,6 +386,45 @@ def test_sns_to_lambda_format(self):
         assert result["x-datadog-parent-id"] == "222222222"
         assert result["dd-pathway-ctx"] == "test-pathway-ctx"
 
+    def test_kinesis_to_lambda_format(self):
+        """Test format: message.kinesis.data.decode()._datadog (Kinesis -> lambda)"""
+        trace_context = {
+            "x-datadog-trace-id": "555444333",
+            "x-datadog-parent-id": "888777666",
+            "dd-pathway-ctx": "test-pathway-ctx",
+        }
+
+        # Create the kinesis data payload
+        kinesis_payload = {
+            "_datadog": trace_context,
+            "actualData": "some business data",
+        }
+        encoded_kinesis_data = base64.b64encode(
+            json.dumps(kinesis_payload).encode("utf-8")
+        ).decode("utf-8")
+
+        kinesis_lambda_record = {
+            "eventSource": "aws:kinesis",
+            "eventSourceARN": (
+                "arn:aws:kinesis:us-east-1:123456789012:stream/my-stream"
+            ),
+            "kinesis": {
+                "data": encoded_kinesis_data,
+                "partitionKey": "partition-key-1",
+                "sequenceNumber": (
+                    "49590338271490256608559692538361571095921575989136588898"
+                ),
+            },
+        }
+
+        result = _get_dsm_context_from_lambda(kinesis_lambda_record)
+
+        assert result is not None
+        assert result == trace_context
+        assert result["x-datadog-trace-id"] == "555444333"
+        assert result["x-datadog-parent-id"] == "888777666"
+        assert result["dd-pathway-ctx"] == "test-pathway-ctx"
+
     def test_no_message_attributes(self):
         """Test message without MessageAttributes returns None."""
         message = {