extract side effects into own function

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -66,28 +66,22 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
-def _extract_context_with_data_streams(context_json, event_type, arn):
-    from ddtrace.data_streams import set_consume_checkpoint
-
-    """
-    Extracts the context from a JSON carrier and optionally sets a dsm consume checkpoint
-    if the context is complete and data streams are enabled.
-    """
-    context = propagator.extract(context_json)
-
+def _set_data_streams_checkpoint(context_json, event_type, arn):
     if not config.data_streams_enabled:
-        return context
+        return
 
     if "dd-pathway-ctx-base64" not in context_json:
-        return context
+        return
+
     try:
+        from ddtrace.data_streams import set_consume_checkpoint
+
         carrier_get = _create_carrier_get(context_json)
         set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False)
     except Exception as e:
         logger.debug(
             f"DSM:Failed to set consume checkpoint for {event_type} {arn}: {e}"
         )
-    return context
 
 
 def _create_carrier_get(context_json):
@@ -308,9 +302,10 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                         logger.debug(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
-                return _extract_context_with_data_streams(
-                    dd_data, "sqs" if is_sqs else "sns", arn
-                )
+                event_type = "sqs" if is_sqs else "sns"
+                context = propagator.extract(dd_data)
+                _set_data_streams_checkpoint(dd_data, event_type, arn)
+                return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
             # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -411,7 +406,9 @@ def extract_context_from_kinesis_event(event, lambda_context):
             data_obj = json.loads(data_str)
             dd_ctx = data_obj.get("_datadog")
             if dd_ctx:
-                return _extract_context_with_data_streams(dd_ctx, "kinesis", arn)
+                context = propagator.extract(dd_ctx)
+                _set_data_streams_checkpoint(dd_ctx, "kinesis", arn)
+                return context
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
 

tests/test_tracing.py

@@ -42,7 +42,6 @@
     service_mapping as global_service_mapping,
     propagator,
     emit_telemetry_on_exception_outside_of_handler,
-    _extract_context_with_data_streams,
     _create_carrier_get,
     extract_context_from_sqs_or_sns_event_or_context,
     extract_context_from_kinesis_event,
@@ -2445,12 +2444,8 @@ def test_exception_outside_handler_tracing_disabled(
         mock_trace.assert_not_called()
 
 
-class TestExtractContext(unittest.TestCase):
+class TestSetDataStreamsCheckpoint(unittest.TestCase):
     def setUp(self):
-        patcher = patch("datadog_lambda.tracing.propagator.extract")
-        self.mock_extract = patcher.start()
-        self.addCleanup(patcher.stop)
-
         checkpoint_patcher = patch("ddtrace.data_streams.set_consume_checkpoint")
         self.mock_checkpoint = checkpoint_patcher.start()
         self.addCleanup(checkpoint_patcher.stop)
@@ -2460,73 +2455,61 @@ def setUp(self):
         self.addCleanup(logger_patcher.stop)
 
     @patch("datadog_lambda.config.Config.data_streams_enabled", False)
-    def test_extract_context_data_streams_disabled(self):
+    def test_set_checkpoint_data_streams_disabled(self):
         context_json = {"dd-pathway-ctx-base64": "12345"}
         event_type = "sqs"
         arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
 
-        mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        self.mock_extract.return_value = mock_context
+        from datadog_lambda.tracing import _set_data_streams_checkpoint
 
-        result = _extract_context_with_data_streams(context_json, event_type, arn)
+        _set_data_streams_checkpoint(context_json, event_type, arn)
 
-        self.mock_extract.assert_called_once_with(context_json)
         self.mock_checkpoint.assert_not_called()
-        self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.config.Config.data_streams_enabled", True)
-    def test_extract_context_data_streams_enabled_complete_context(self):
+    def test_set_checkpoint_data_streams_enabled_complete_context(self):
         context_json = {"dd-pathway-ctx-base64": "12345"}
         event_type = "sqs"
         arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
 
-        mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        self.mock_extract.return_value = mock_context
+        from datadog_lambda.tracing import _set_data_streams_checkpoint
 
-        result = _extract_context_with_data_streams(context_json, event_type, arn)
+        _set_data_streams_checkpoint(context_json, event_type, arn)
 
-        self.mock_extract.assert_called_once_with(context_json)
         self.mock_checkpoint.assert_called_once()
         args, kwargs = self.mock_checkpoint.call_args
         self.assertEqual(args[0], event_type)
         self.assertEqual(args[1], arn)
         self.assertTrue(callable(args[2]))
         self.assertEqual(kwargs["manual_checkpoint"], False)
-        self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.config.Config.data_streams_enabled", True)
-    def test_extract_context_data_streams_enabled_invalid_context(self):
+    def test_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"
 
-        mock_context = Context(trace_id=12345, span_id=12345, sampling_priority=1)
-        self.mock_extract.return_value = mock_context
+        from datadog_lambda.tracing import _set_data_streams_checkpoint
 
-        result = _extract_context_with_data_streams(context_json, event_type, arn)
+        _set_data_streams_checkpoint(context_json, event_type, arn)
 
-        self.mock_extract.assert_called_once_with(context_json)
         self.mock_checkpoint.assert_not_called()
-        self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.config.Config.data_streams_enabled", True)
-    def test_extract_context_exception_path(self):
+    def test_set_checkpoint_exception_path(self):
         context_json = {"dd-pathway-ctx-base64": "12345"}
         event_type = "sqs"
         arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
 
-        mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        self.mock_extract.return_value = mock_context
-
         test_exception = Exception("Test exception")
         self.mock_checkpoint.side_effect = test_exception
 
-        result = _extract_context_with_data_streams(context_json, event_type, arn)
+        from datadog_lambda.tracing import _set_data_streams_checkpoint
+
+        _set_data_streams_checkpoint(context_json, event_type, arn)
 
-        self.mock_extract.assert_called_once_with(context_json)
         self.mock_checkpoint.assert_called_once()
         self.mock_logger.debug.assert_called_once()
-        self.assertEqual(result, mock_context)
 
 
 class TestCreateCarrierGet(unittest.TestCase):
@@ -2566,9 +2549,10 @@ class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
     def test_sqs_event_with_datadog_message_attributes(
-        self, mock_extract_context_with_data_streams
+        self, mock_extract, mock_set_checkpoint
     ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
@@ -2585,20 +2569,22 @@ def test_sqs_event_with_datadog_message_attributes(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
     def test_sqs_event_with_binary_datadog_message_attributes(
-        self, mock_extract_context_with_data_streams
+        self, mock_extract, mock_set_checkpoint
     ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
@@ -2616,20 +2602,22 @@ def test_sqs_event_with_binary_datadog_message_attributes(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
     def test_sns_event_with_datadog_message_attributes(
-        self, mock_extract_context_with_data_streams
+        self, mock_extract, mock_set_checkpoint
     ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
@@ -2649,20 +2637,22 @@ def test_sns_event_with_datadog_message_attributes(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
     def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present(
-        self, mock_extract_context_with_data_streams
+        self, mock_extract, mock_set_checkpoint
     ):
         """Test that is_sqs = True when eventSourceARN is present in first record"""
         dd_data = {"dd-pathway-ctx-base64": "12345"}
@@ -2680,20 +2670,22 @@ def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
     def test_sns_to_sqs_event_detection_and_processing(
-        self, mock_extract_context_with_data_streams
+        self, mock_extract, mock_set_checkpoint
     ):
         """Test SNS->SQS case where SQS body contains SNS notification"""
         dd_data = {"dd-pathway-ctx-base64": "12345"}
@@ -2719,13 +2711,14 @@ def test_sns_to_sqs_event_detection_and_processing(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
@@ -2735,10 +2728,9 @@ class TestExtractContextFromKinesisEvent(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
-    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
-    def test_kinesis_event_with_datadog_data(
-        self, mock_extract_context_with_data_streams
-    ):
+    @patch("datadog_lambda.tracing._set_data_streams_checkpoint")
+    @patch("datadog_lambda.tracing.propagator.extract")
+    def test_kinesis_event_with_datadog_data(self, mock_extract, mock_set_checkpoint):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         kinesis_data = {"_datadog": dd_data, "message": "test"}
         kinesis_data_str = json.dumps(kinesis_data)
@@ -2754,11 +2746,12 @@ def test_kinesis_event_with_datadog_data(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context_with_data_streams.return_value = mock_context
+        mock_extract.return_value = mock_context
 
         result = extract_context_from_kinesis_event(event, self.lambda_context)
 
-        mock_extract_context_with_data_streams.assert_called_once_with(
+        mock_extract.assert_called_once_with(dd_data)
+        mock_set_checkpoint.assert_called_once_with(
             dd_data,
             "kinesis",
             "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",