remove unneccessary checks, still set checkpoints even when dsm context fails to be propagated

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -69,18 +69,17 @@
 
 def _dsm_set_checkpoint(context_json, event_type, arn):
     from ddtrace.data_streams import set_consume_checkpoint
-    from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
 
     """
     Extracts the context from a JSON carrier and optionally sets a dsm consume checkpoint
     if the context is complete and data streams are enabled.
     """
-
-    if PROPAGATION_KEY_BASE_64 not in context_json:
+    if not isinstance(context_json, dict):
         return
 
     if not config.data_streams_enabled:
         return
+
     try:
         carrier_get = _create_carrier_get(context_json)
         set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False)
@@ -336,9 +335,13 @@ def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled
                             span_id=int(x_ray_context["parent_id"], 16),
                             sampling_priority=float(x_ray_context["sampled"]),
                         )
+        # Still want to set a DSM checkpoint even if DSM context not propagated
+        _dsm_set_checkpoint({}, "sqs" if is_sqs else "sns", arn)
         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({}, "sqs" if is_sqs else "sns", arn)
         return extract_context_from_lambda_context(lambda_context)
 
 
@@ -422,7 +425,8 @@ def extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled(
                 return 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({}, "kinesis", arn)
     return extract_context_from_lambda_context(lambda_context)
 
 

tests/test_tracing.py

@@ -2480,16 +2480,6 @@ def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self):
         self.assertTrue(callable(args[2]))
         self.assertEqual(kwargs["manual_checkpoint"], False)
 
-    @patch("datadog_lambda.config.Config.data_streams_enabled", True)
-    def test_dsm_set_checkpoint_data_streams_enabled_invalid_context(self):
-        context_json = {"something-malformed": "12345"}
-        event_type = "sqs"
-        arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
-
-        _dsm_set_checkpoint(context_json, event_type, arn)
-
-        self.mock_checkpoint.assert_not_called()
-
     @patch("datadog_lambda.config.Config.data_streams_enabled", True)
     def test_dsm_set_checkpoint_exception_path(self):
         context_json = {"dd-pathway-ctx-base64": "12345"}
@@ -2504,6 +2494,34 @@ def test_dsm_set_checkpoint_exception_path(self):
         self.mock_checkpoint.assert_called_once()
         self.mock_logger.debug.assert_called_once()
 
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
+    def test_dsm_set_checkpoint_non_dict_context_sqs(self, mock_checkpoint):
+        _dsm_set_checkpoint(
+            "not_a_dict", "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
+        )
+        mock_checkpoint.assert_not_called()
+
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
+    def test_dsm_set_checkpoint_non_dict_context_sns_to_sqs(self, mock_checkpoint):
+        _dsm_set_checkpoint(
+            ["not", "a", "dict"], "sqs", "arn:aws:sqs:us-east-1:123456789012:test"
+        )
+        mock_checkpoint.assert_not_called()
+
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
+    def test_dsm_set_checkpoint_non_dict_context_kinesis(self, mock_checkpoint):
+        _dsm_set_checkpoint(
+            12345, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test"
+        )
+        mock_checkpoint.assert_not_called()
+
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
+    def test_dsm_set_checkpoint_non_dict_context_sns(self, mock_checkpoint):
+        _dsm_set_checkpoint(
+            None, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
+        )
+        mock_checkpoint.assert_not_called()
+
 
 class TestCreateCarrierGet(unittest.TestCase):
     def test_create_carrier_get_with_valid_data(self):
@@ -2538,7 +2556,7 @@ def test_create_carrier_get_with_empty_context(self):
         self.assertIsNone(carrier_get("any-key"))
 
 
-class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase):
+class TestExtractContextFromSqsOrSnsEventWithDSMLogic(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
@@ -2716,8 +2734,168 @@ def test_sns_to_sqs_event_detection_and_processing(
         )
         self.assertEqual(result, mock_context)
 
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_sqs_event_without_datadog_message_attributes(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "messageAttributes": {},
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_sqs_event_with_malformed_datadog_message_attributes(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        dd_json_data = "not-json"
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "messageAttributes": {
+                        "_datadog": {"dataType": "String", "stringValue": dd_json_data}
+                    },
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_sns_event_without_datadog_message_attributes(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        event = {
+            "Records": [
+                {
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+                        "MessageAttributes": {},
+                    }
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_sns_event_with_malformed_datadog_message_attributes(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        dd_json_data = "not-json"
+
+        event = {
+            "Records": [
+                {
+                    "Sns": {
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+                        "MessageAttributes": {
+                            "_datadog": {"Type": "String", "Value": dd_json_data}
+                        },
+                    }
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_sns_to_sqs_event_with_malformed_datadog_message_attributes(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        """Test SNS->SQS case where SQS body contains SNS notification with malformed datadog ctx"""
+        dd_json_data = "not-json"
+
+        sns_notification = {
+            "Type": "Notification",
+            "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
+            "MessageAttributes": {
+                "_datadog": {"Type": "String", "Value": dd_json_data}
+            },
+            "Message": "test message",
+        }
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
+                    "body": json.dumps(sns_notification),
+                    "messageAttributes": {},
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
 
-class TestExtractContextFromKinesisEvent(unittest.TestCase):
+class TestExtractContextFromKinesisEventWithDSMLogic(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
@@ -2754,3 +2932,67 @@ def test_kinesis_event_with_datadog_data(
             "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
         )
         self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_kinesis_event_without_datadog_data(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        kinesis_data = {"message": "test"}
+        kinesis_data_str = json.dumps(kinesis_data)
+        encoded_data = base64.b64encode(kinesis_data_str.encode()).decode()
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+                    "kinesis": {"data": encoded_data},
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {},
+            "kinesis",
+            "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)
+
+    @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
+    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    def test_kinesis_event_with_malformed_data(
+        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+    ):
+        encoded_data = "not-base64-or-json"
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+                    "kinesis": {"data": encoded_data},
+                }
+            ]
+        }
+
+        mock_context = Context(trace_id=123, span_id=456)
+        mock_extract_from_lambda_context.return_value = mock_context
+
+        result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled(
+            event, self.lambda_context
+        )
+
+        mock_dsm_set_checkpoint.assert_called_once_with(
+            {},
+            "kinesis",
+            "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+        )
+        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        self.assertEqual(result, mock_context)