refactorings

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -0,0 +1,37 @@
+from datadog_lambda import logger
+from datadog_lambda.trigger import EventTypes
+
+
+def set_dsm_context(event, event_source):
+
+    if event_source.equals(EventTypes.SQS):
+        _dsm_set_sqs_context(event)
+
+
+def _dsm_set_sqs_context(event):
+    from datadog_lambda.wrapper import format_err_with_traceback
+
+    from ddtrace.internal.datastreams.processor import (
+        DataStreamsProcessor as processor,
+        DsmPathwayCodec,
+    )
+    from ddtrace.internal.datastreams.botocore import (
+        get_datastreams_context,
+        calculate_sqs_payload_size,
+    )
+
+    records = event.get("Records", [])
+    for record in records:
+        try:
+            queue_arn = record.get("eventSourceARN", "")
+
+            contextjson = get_datastreams_context(record)
+            payload_size = calculate_sqs_payload_size(record)
+
+            ctx = DsmPathwayCodec.decode(contextjson, processor())
+            ctx.set_checkpoint(
+                ["direction:in", "topic:" + queue_arn, "type:sqs"],
+                payload_size=payload_size,
+            )
+        except Exception as e:
+            logger.error(format_err_with_traceback(e))

datadog_lambda/wrapper.py

@@ -9,6 +9,7 @@
 from importlib import import_module
 from time import time_ns
 
+from datadog_lambda.dsm import set_dsm_context
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -291,39 +292,6 @@ def _inject_authorizer_span_headers(self, request_id):
         self.response["context"]["_datadog"] = datadog_data
 
     def _before(self, event, context):
-        def _dsm_set_sqs_context(record):
-            try:
-                queue_arn = record.get("eventSourceARN", "")
-
-                contextjson = get_datastreams_context(record)
-                payload_size = calculate_sqs_payload_size(record)
-
-                ctx = DsmPathwayCodec.decode(contextjson, processor())
-                ctx.set_checkpoint(
-                    ["direction:in", "topic:" + queue_arn, "type:sqs"],
-                    payload_size=payload_size,
-                )
-
-            except Exception as e:
-                logger.error(format_err_with_traceback(e))
-
-        if self.data_streams_enabled:
-            from ddtrace.internal.datastreams.processor import (
-                DataStreamsProcessor as processor,
-                DsmPathwayCodec,
-            )
-            from ddtrace.internal.datastreams.botocore import (
-                get_datastreams_context,
-                calculate_sqs_payload_size,
-            )
-
-            if isinstance(event, dict) and "Records" in event and event["Records"]:
-                sqs_records = [
-                    r for r in event["Records"] if r.get("eventSource") == "aws:sqs"
-                ]
-                if sqs_records:
-                    for record in sqs_records:
-                        _dsm_set_sqs_context(record)
 
         try:
             self.response = None
@@ -360,6 +328,8 @@ def _dsm_set_sqs_context(record):
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, self.decode_authorizer_context
                     )
+                if self.data_streams_enabled:
+                    set_dsm_context(event, event_source)
                 self.span = create_function_execution_span(
                     context=context,
                     function_name=self.function_name,

tests/test_wrapper.py

@@ -565,6 +565,11 @@ def return_type_test(event, context):
 
     @patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"})
     def test_datadog_lambda_wrapper_dsm_sqs_context_pathway_verification(self):
+        from datadog_lambda.trigger import _EventSource, EventTypes
+
+        sqs_event_source = _EventSource(EventTypes.SQS)
+        self.mock_extract_dd_trace_context.return_value = ({}, None, sqs_event_source)
+
         with patch(
             "ddtrace.internal.datastreams.processor.get_connection"
         ) as mock_get_connection:
@@ -729,6 +734,33 @@ def lambda_handler(event, context):
 
                 processor_instance.shutdown(timeout=0.1)
 
+    @patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"})
+    @patch("datadog_lambda.wrapper.set_dsm_context")
+    def test_set_dsm_context_called_when_enabled(self, mock_set_dsm_context):
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            return {"statusCode": 200, "body": "processed"}
+
+        lambda_event = {}
+        lambda_handler(lambda_event, get_mock_context())
+
+        mock_set_dsm_context.assert_called_once()
+
+    @patch("datadog_lambda.wrapper.set_dsm_context")
+    def test_set_dsm_context_not_called_when_disabled(self, mock_set_dsm_context):
+        # Ensure DD_DATA_STREAMS_ENABLED is not in environment
+        if "DD_DATA_STREAMS_ENABLED" in os.environ:
+            del os.environ["DD_DATA_STREAMS_ENABLED"]
+
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            return {"statusCode": 200, "body": "processed"}
+
+        lambda_event = {}
+        lambda_handler(lambda_event, get_mock_context())
+
+        mock_set_dsm_context.assert_not_called()
+
 
 class TestLambdaDecoratorSettings(unittest.TestCase):
     def test_some_envs_should_depend_on_dd_tracing_enabled(self):