Add batch processing tests, SNS batch processing not possible

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -68,28 +68,32 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
-def _dsm_set_context_sqs_or_sns_event(event, event_type):
+def _dsm_set_context_sqs_event(event):
     for record in event.get("Records", [])[1:]:
-        if arn := _get_source_arn_sqs_or_sns(record, event_type):
-            context = _extract_context_from_sqs_or_sns_record(record)
-            _dsm_set_checkpoint(context, event_type, arn)
+        if arn := record.get("eventSourceARN"):
+            try:
+                context = _extract_context_from_sqs_or_sns_record(record)
+            except Exception as e:
+                logger.debug(
+                    f"DSM: Failed to extract context with error {e}, will still set checkpoint"
+                )
+                context = None
+            _dsm_set_checkpoint(context, "sqs", arn)
 
 
 def _dsm_set_context_kinesis_event(event):
     for record in event.get("Records", [])[1:]:
         if (arn := record.get("eventSourceARN")) and (kinesis := record.get("kinesis")):
-            context = _extract_context_from_kinesis_record(kinesis)
+            try:
+                context = _extract_context_from_kinesis_record(kinesis)
+            except Exception as e:
+                logger.debug(
+                    f"DSM: Failed to extract context with error {e}, will still set checkpoint"
+                )
+                context = None
             _dsm_set_checkpoint(context, "kinesis", arn)
 
 
-def _get_source_arn_sqs_or_sns(record, event_type):
-    return (
-        record.get("Sns", {}).get("TopicArn")
-        if event_type == "sns"
-        else record.get("eventSourceARN")
-    )
-
-
 def _dsm_set_checkpoint(context_json, event_type, arn):
     if not config.data_streams_enabled:
         return
@@ -271,7 +275,11 @@ def extract_context_from_sqs_or_sns_event_or_context(
 
     try:
         first_record = event.get("Records")[0]
-        source_arn = _get_source_arn_sqs_or_sns(first_record, event_type)
+        source_arn = (
+            first_record.get("Sns", {}).get("TopicArn")
+            if event_type == "sns"
+            else first_record.get("eventSourceARN")
+        )
         dd_data = _extract_context_from_sqs_or_sns_record(first_record)
         if dd_data:
             if is_step_function_event(dd_data):
@@ -283,8 +291,9 @@ def extract_context_from_sqs_or_sns_event_or_context(
                     )
             context = propagator.extract(dd_data)
             _dsm_set_checkpoint(dd_data, event_type, source_arn)
-            if config.data_streams_enabled:
-                _dsm_set_context_sqs_or_sns_event(event, event_type)
+            # Batch Processing does not occur for SNS events
+            if config.data_streams_enabled and event_type == "sqs":
+                _dsm_set_context_sqs_event(event)
             return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
@@ -310,15 +319,17 @@ def extract_context_from_sqs_or_sns_event_or_context(
                         )
         # Still want to set a DSM checkpoint even if DSM context not propagated
         _dsm_set_checkpoint(None, event_type, source_arn)
-        if config.data_streams_enabled:
-            _dsm_set_context_sqs_or_sns_event(event, event_type)
+        # Batch Processing does not occur for SNS events
+        if config.data_streams_enabled and event_type == "sqs":
+            _dsm_set_context_sqs_event(event)
         return extract_context_from_lambda_context(lambda_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
         # Still want to set a DSM checkpoint even if DSM context not propagated
         _dsm_set_checkpoint(None, event_type, source_arn)
-        if config.data_streams_enabled:
-            _dsm_set_context_sqs_or_sns_event(event, event_type)
+        # Batch Processing does not occur for SNS events
+        if config.data_streams_enabled and event_type == "sqs":
+            _dsm_set_context_sqs_event(event)
         return extract_context_from_lambda_context(lambda_context)
 
 

tests/test_tracing.py

@@ -2743,6 +2743,71 @@ def test_sqs_batch_processing(self):
         carrier_get_2 = args_2[2]
         self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), "record2")
 
+    def test_sqs_batch_processing_with_invalid_record(self):
+        dd_data_1 = {"dd-pathway-ctx-base64": "record1"}
+        dd_data_3 = {"dd-pathway-ctx-base64": "record3"}
+        dd_json_data_1 = json.dumps(dd_data_1)
+        dd_json_data_3 = json.dumps(dd_data_3)
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "messageAttributes": {
+                        "_datadog": {
+                            "dataType": "String",
+                            "stringValue": dd_json_data_1,
+                        }
+                    },
+                    "eventSource": "aws:sqs",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "messageAttributes": {
+                        "_datadog": {
+                            "dataType": "String",
+                            "stringValue": "invalid json",  # This will cause extraction to fail
+                        }
+                    },
+                    "eventSource": "aws:sqs",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "messageAttributes": {
+                        "_datadog": {
+                            "dataType": "String",
+                            "stringValue": dd_json_data_3,
+                        }
+                    },
+                    "eventSource": "aws:sqs",
+                },
+            ]
+        }
+
+        extract_context_from_sqs_or_sns_event_or_context(
+            event, self.lambda_context, parse_event_source(event)
+        )
+
+        self.assertEqual(self.mock_checkpoint.call_count, 3)
+
+        args_1, _ = self.mock_checkpoint.call_args_list[0]
+        self.assertEqual(args_1[0], "sqs")
+        self.assertEqual(args_1[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_1 = args_1[2]
+        self.assertEqual(carrier_get_1("dd-pathway-ctx-base64"), "record1")
+
+        args_2, _ = self.mock_checkpoint.call_args_list[1]
+        self.assertEqual(args_2[0], "sqs")
+        self.assertEqual(args_2[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_2 = args_2[2]
+        self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), None)
+
+        args_3, _ = self.mock_checkpoint.call_args_list[2]
+        self.assertEqual(args_3[0], "sqs")
+        self.assertEqual(args_3[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_3 = args_3[2]
+        self.assertEqual(carrier_get_3("dd-pathway-ctx-base64"), "record3")
+
     def test_sqs_source_arn_not_found(self):
         event = {
             "Records": [
@@ -3019,53 +3084,6 @@ def test_sns_data_streams_disabled(self):
 
         self.mock_checkpoint.assert_not_called()
 
-    def test_sns_batch_processing(self):
-        dd_data_1 = {"dd-pathway-ctx-base64": "record1"}
-        dd_data_2 = {"dd-pathway-ctx-base64": "record2"}
-        dd_json_data_1 = json.dumps(dd_data_1)
-        dd_json_data_2 = json.dumps(dd_data_2)
-
-        event = {
-            "Records": [
-                {
-                    "Sns": {
-                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
-                        "MessageAttributes": {
-                            "_datadog": {"Type": "String", "Value": dd_json_data_1}
-                        },
-                    },
-                    "eventSource": "aws:sns",
-                },
-                {
-                    "Sns": {
-                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
-                        "MessageAttributes": {
-                            "_datadog": {"Type": "String", "Value": dd_json_data_2}
-                        },
-                    },
-                    "eventSource": "aws:sns",
-                },
-            ]
-        }
-
-        extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context, parse_event_source(event)
-        )
-
-        self.assertEqual(self.mock_checkpoint.call_count, 2)
-
-        args_1, _ = self.mock_checkpoint.call_args_list[0]
-        self.assertEqual(args_1[0], "sns")
-        self.assertEqual(args_1[1], "arn:aws:sns:us-east-1:123456789012:test-topic")
-        carrier_get_1 = args_1[2]
-        self.assertEqual(carrier_get_1("dd-pathway-ctx-base64"), "record1")
-
-        args_2, _ = self.mock_checkpoint.call_args_list[1]
-        self.assertEqual(args_2[0], "sns")
-        self.assertEqual(args_2[1], "arn:aws:sns:us-east-1:123456789012:test-topic")
-        carrier_get_2 = args_2[2]
-        self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), "record2")
-
     # SNS -> SQS TESTS
 
     def test_sns_to_sqs_context_propagated_string_value(self):
@@ -3362,6 +3380,81 @@ def test_sns_to_sqs_batch_processing(self):
         carrier_get_2 = args_2[2]
         self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), "record2")
 
+    def test_sns_to_sqs_batch_processing_with_invalid_record(self):
+        dd_data_1 = {"dd-pathway-ctx-base64": "record1"}
+        dd_data_3 = {"dd-pathway-ctx-base64": "record3"}
+        dd_json_data_1 = json.dumps(dd_data_1)
+        dd_json_data_3 = json.dumps(dd_data_3)
+
+        sns_message_1 = {
+            "Type": "Notification",
+            "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+            "MessageAttributes": {
+                "_datadog": {"Type": "String", "Value": dd_json_data_1}
+            },
+        }
+        sns_message_2 = {
+            "Type": "Notification",
+            "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+            "MessageAttributes": {
+                "_datadog": {
+                    "Type": "String",
+                    "Value": "invalid json",
+                }  # This will cause extraction to fail
+            },
+        }
+        sns_message_3 = {
+            "Type": "Notification",
+            "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+            "MessageAttributes": {
+                "_datadog": {"Type": "String", "Value": dd_json_data_3}
+            },
+        }
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "body": json.dumps(sns_message_1),
+                    "eventSource": "aws:sqs",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "body": json.dumps(sns_message_2),
+                    "eventSource": "aws:sqs",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "body": json.dumps(sns_message_3),
+                    "eventSource": "aws:sqs",
+                },
+            ]
+        }
+
+        extract_context_from_sqs_or_sns_event_or_context(
+            event, self.lambda_context, parse_event_source(event)
+        )
+
+        self.assertEqual(self.mock_checkpoint.call_count, 3)
+
+        args_1, _ = self.mock_checkpoint.call_args_list[0]
+        self.assertEqual(args_1[0], "sqs")
+        self.assertEqual(args_1[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_1 = args_1[2]
+        self.assertEqual(carrier_get_1("dd-pathway-ctx-base64"), "record1")
+
+        args_2, _ = self.mock_checkpoint.call_args_list[1]
+        self.assertEqual(args_2[0], "sqs")
+        self.assertEqual(args_2[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_2 = args_2[2]
+        self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), None)
+
+        args_3, _ = self.mock_checkpoint.call_args_list[2]
+        self.assertEqual(args_3[0], "sqs")
+        self.assertEqual(args_3[1], "arn:aws:sqs:us-east-1:123456789012:test-queue")
+        carrier_get_3 = args_3[2]
+        self.assertEqual(carrier_get_3("dd-pathway-ctx-base64"), "record3")
+
     def test_sns_to_sqs_source_arn_not_found(self):
         sns_notification = {
             "Type": "Notification",
@@ -3574,6 +3667,63 @@ def test_kinesis_batch_processing(self):
         carrier_get_2 = args_2[2]
         self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), "record2")
 
+    def test_kinesis_batch_processing_with_invalid_record(self):
+        dd_data_1 = {"dd-pathway-ctx-base64": "record1"}
+        dd_data_3 = {"dd-pathway-ctx-base64": "record3"}
+
+        kinesis_data_1 = {"_datadog": dd_data_1, "message": "test1"}
+        kinesis_data_2 = {"invalid": "data"}
+        kinesis_data_3 = {"_datadog": dd_data_3, "message": "test3"}
+
+        encoded_data_1 = base64.b64encode(json.dumps(kinesis_data_1).encode()).decode()
+        encoded_data_2 = base64.b64encode(json.dumps(kinesis_data_2).encode()).decode()
+        encoded_data_3 = base64.b64encode(json.dumps(kinesis_data_3).encode()).decode()
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+                    "kinesis": {"data": encoded_data_1},
+                },
+                {
+                    "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+                    "kinesis": {"data": encoded_data_2},
+                },
+                {
+                    "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+                    "kinesis": {"data": encoded_data_3},
+                },
+            ]
+        }
+
+        extract_context_from_kinesis_event(event, self.lambda_context)
+
+        self.assertEqual(self.mock_checkpoint.call_count, 3)
+
+        args_1, _ = self.mock_checkpoint.call_args_list[0]
+        self.assertEqual(args_1[0], "kinesis")
+        self.assertEqual(
+            args_1[1], "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
+        )
+        carrier_get_1 = args_1[2]
+        self.assertEqual(carrier_get_1("dd-pathway-ctx-base64"), "record1")
+
+        args_2, _ = self.mock_checkpoint.call_args_list[1]
+        self.assertEqual(args_2[0], "kinesis")
+        self.assertEqual(
+            args_2[1], "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
+        )
+        carrier_get_2 = args_2[2]
+        self.assertEqual(carrier_get_2("dd-pathway-ctx-base64"), None)
+
+        args_3, _ = self.mock_checkpoint.call_args_list[2]
+        self.assertEqual(args_3[0], "kinesis")
+        self.assertEqual(
+            args_3[1], "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
+        )
+        carrier_get_3 = args_3[2]
+        self.assertEqual(carrier_get_3("dd-pathway-ctx-base64"), "record3")
+
     def test_kinesis_source_arn_not_found(self):
         kinesis_data = {"message": "test"}
         kinesis_data_str = json.dumps(kinesis_data)