simplify PR

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -1,5 +1,4 @@
 import json
-import base64
 
 from ddtrace.internal.logger import get_logger
 from datadog_lambda import logger
@@ -65,36 +64,9 @@ 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.messageAttributes._datadog.binaryValue.decode() (SNS -> SQS -> lambda, raw)
-        - message.body.MessageAttributes._datadog.Value.decode() (SNS -> SQS -> 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
-    elif "Sns" in message:
-        message_body = message["Sns"]
-    else:
-        try:
-            body = message.get("body")
-            if body:
-                message_body = json.loads(body)
-        except (ValueError, TypeError):
-            log.debug("Unable to parse lambda message body as JSON, treat as non-json")
-
-    message_attributes = message_body.get("MessageAttributes") or message_body.get(
-        "messageAttributes"
-    )
+    message_attributes = message.get("messageAttributes")
     if not message_attributes:
         log.debug("DataStreams skipped lambda message: %r", message)
         return None
@@ -105,18 +77,9 @@ def _get_dsm_context_from_lambda(message):
 
     datadog_attr = message_attributes["_datadog"]
 
-    if message_body.get("Type") == "Notification":
-        # SNS -> lambda notification
-        if datadog_attr.get("Type") == "Binary":
-            context_json = json.loads(base64.b64decode(datadog_attr["Value"]).decode())
-    elif "stringValue" in datadog_attr:
+    if "stringValue" in datadog_attr:
         # SQS -> lambda
         context_json = json.loads(datadog_attr["stringValue"])
-    elif "binaryValue" in datadog_attr:
-        # SNS -> SQS -> lambda, raw message delivery
-        context_json = json.loads(
-            base64.b64decode(datadog_attr["binaryValue"]).decode()
-        )
     else:
         log.debug("DataStreams did not handle lambda message: %r", message)
 

tests/test_dsm.py

@@ -1,6 +1,5 @@
 import unittest
 import json
-import base64
 from unittest.mock import patch, MagicMock
 
 from datadog_lambda.dsm import (
@@ -165,150 +164,6 @@ def test_sqs_to_lambda_string_value_format(self):
         assert result["x-datadog-parent-id"] == "321987654"
         assert result["dd-pathway-ctx"] == "test-pathway-ctx"
 
-    def test_sns_to_lambda_format(self):
-        """Test format: message.Sns.MessageAttributes._datadog.Value.decode() (SNS -> lambda)"""
-        trace_context = {
-            "x-datadog-trace-id": "111111111",
-            "x-datadog-parent-id": "222222222",
-            "dd-pathway-ctx": "test-pathway-ctx",
-        }
-        binary_data = base64.b64encode(
-            json.dumps(trace_context).encode("utf-8")
-        ).decode("utf-8")
-
-        sns_lambda_record = {
-            "EventSource": "aws:sns",
-            "EventSubscriptionArn": (
-                "arn:aws:sns:us-east-1:123456789012:sns-topic:12345678-1234-1234-1234-123456789012"
-            ),
-            "Sns": {
-                "Type": "Notification",
-                "MessageId": "95df01b4-ee98-5cb9-9903-4c221d41eb5e",
-                "TopicArn": "arn:aws:sns:us-east-1:123456789012:sns-topic",
-                "Subject": "Test Subject",
-                "Message": "Hello from SNS!",
-                "Timestamp": "2023-01-01T12:00:00.000Z",
-                "MessageAttributes": {
-                    "_datadog": {"Type": "Binary", "Value": binary_data}
-                },
-            },
-        }
-
-        result = _get_dsm_context_from_lambda(sns_lambda_record)
-
-        assert result is not None
-        assert result == trace_context
-        assert result["x-datadog-trace-id"] == "111111111"
-        assert result["x-datadog-parent-id"] == "222222222"
-        assert result["dd-pathway-ctx"] == "test-pathway-ctx"
-
-    def test_sns_to_sqs_to_lambda_binary_value_format(self):
-        """Test format: message.messageAttributes._datadog.binaryValue.decode() (SNS -> SQS -> lambda, raw)"""
-        trace_context = {
-            "x-datadog-trace-id": "777666555",
-            "x-datadog-parent-id": "444333222",
-            "dd-pathway-ctx": "test-pathway-ctx",
-        }
-        binary_data = base64.b64encode(
-            json.dumps(trace_context).encode("utf-8")
-        ).decode("utf-8")
-
-        lambda_record = {
-            "messageId": "test-message-id",
-            "receiptHandle": "test-receipt-handle",
-            "body": "Test message body",
-            "messageAttributes": {
-                "_datadog": {"binaryValue": binary_data, "dataType": "Binary"}
-            },
-            "eventSource": "aws:sqs",
-            "eventSourceARN": "arn:aws:sqs:us-west-2:123456789012:test-queue",
-        }
-
-        result = _get_dsm_context_from_lambda(lambda_record)
-
-        assert result is not None
-        assert result == trace_context
-        assert result["x-datadog-trace-id"] == "777666555"
-        assert result["x-datadog-parent-id"] == "444333222"
-        assert result["dd-pathway-ctx"] == "test-pathway-ctx"
-
-    def test_sns_to_sqs_to_lambda_body_format(self):
-        """Test format: message.body.MessageAttributes._datadog.Value.decode() (SNS -> SQS -> lambda)"""
-        trace_context = {
-            "x-datadog-trace-id": "123987456",
-            "x-datadog-parent-id": "654321987",
-            "x-datadog-sampling-priority": "1",
-            "dd-pathway-ctx": "test-pathway-ctx",
-        }
-
-        message_body = {
-            "Type": "Notification",
-            "MessageId": "test-message-id",
-            "Message": "Test message from SNS",
-            "MessageAttributes": {
-                "_datadog": {
-                    "Type": "Binary",
-                    "Value": base64.b64encode(
-                        json.dumps(trace_context).encode("utf-8")
-                    ).decode("utf-8"),
-                }
-            },
-        }
-
-        lambda_record = {
-            "messageId": "lambda-message-id",
-            "body": json.dumps(message_body),
-            "eventSource": "aws:sqs",
-            "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:sns-to-sqs-queue",
-        }
-
-        result = _get_dsm_context_from_lambda(lambda_record)
-
-        assert result is not None
-        assert result == trace_context
-        assert result["x-datadog-trace-id"] == "123987456"
-        assert result["x-datadog-parent-id"] == "654321987"
-        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 = {