remove unneccesary work associated with event_source

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -217,7 +217,7 @@ def create_sns_event(message):
     }
 
 
-def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
+def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context, event_source):
     """
     Extract Datadog trace context from an SQS event.
 
@@ -232,7 +232,6 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
     arn = ""
-    event_source = parse_event_source(event)
 
     # EventBridge => SQS
     try:
@@ -649,7 +648,7 @@ def extract_dd_trace_context(
         )
     elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS):
         context = extract_context_from_sqs_or_sns_event_or_context(
-            event, lambda_context
+            event, lambda_context, event_source
         )
     elif event_source.equals(EventTypes.EVENTBRIDGE):
         context = extract_context_from_eventbridge_event(event, lambda_context)

tests/test_tracing.py

@@ -47,6 +47,7 @@
     extract_context_from_sqs_or_sns_event_or_context,
 )
 
+from datadog_lambda.trigger import parse_event_source
 from tests.utils import get_mock_context
 
 
@@ -2534,7 +2535,7 @@ def test_sqs_event_with_datadog_message_attributes(
         mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_extract.assert_called_once_with(dd_data)
@@ -2568,7 +2569,7 @@ def test_sqs_event_with_binary_datadog_message_attributes(
         mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_extract.assert_called_once_with(dd_data)
@@ -2604,7 +2605,7 @@ def test_sns_event_with_datadog_message_attributes(
         mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_extract.assert_called_once_with(dd_data)
@@ -2638,7 +2639,7 @@ def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present(
         mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_extract.assert_called_once_with(dd_data)
@@ -2680,7 +2681,7 @@ def test_sns_to_sqs_event_detection_and_processing(
         mock_extract.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_extract.assert_called_once_with(dd_data)
@@ -2708,7 +2709,7 @@ def test_sqs_event_without_datadog_message_attributes(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_called_once_with(
@@ -2740,7 +2741,7 @@ def test_sqs_event_with_malformed_datadog_message_attributes(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_called_once_with(
@@ -2770,7 +2771,7 @@ def test_sns_event_without_datadog_message_attributes(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_called_once_with(
@@ -2804,7 +2805,7 @@ def test_sns_event_with_malformed_datadog_message_attributes(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_called_once_with(
@@ -2845,7 +2846,7 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_called_once_with(
@@ -2865,7 +2866,7 @@ def test_sqs_sns_event_with_exception_accessing_first_record(
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
-            event, self.lambda_context
+            event, self.lambda_context, parse_event_source(event)
         )
 
         mock_dsm_set_checkpoint.assert_not_called()