fixes

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -68,7 +68,6 @@ def _dsm_set_kinesis_context(event):
         if not context_json:
             logger.debug("DataStreams skipped lambda message: %r", record)
             return None
-
         carrier_get = _create_carrier_get(context_json)
         set_consume_checkpoint("kinesis", arn, carrier_get)
 
@@ -78,10 +77,21 @@ 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):
+            logger.debug("Unable to parse kinesis data for lambda message")
+            return None
+
     if "Sns" in message:
         message_body = message["Sns"]
 

tests/test_dsm.py

@@ -442,6 +442,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 = {