some fixes

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -65,20 +65,12 @@
 DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000"
 HIGHER_64_BITS = "HIGHER_64_BITS"
 LOWER_64_BITS = "LOWER_64_BITS"
-DSM_PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64"
 
 
 def _dsm_set_checkpoint(context_json, event_type, arn):
     if not config.data_streams_enabled:
         return
 
-    if context_json is not None:
-        if (
-            not isinstance(context_json, dict)
-            or DSM_PROPAGATION_KEY_BASE_64 not in context_json
-        ):
-            return
-
     try:
         from ddtrace.data_streams import set_consume_checkpoint
 
@@ -239,8 +231,8 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     Falls back to lambda context if no trace data is found in the SQS message attributes.
     Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
-    is_sqs = False
     arn = ""
+    event_source = parse_event_source(event)
 
     # EventBridge => SQS
     try:
@@ -253,8 +245,6 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     try:
         first_record = event.get("Records")[0]
         arn = first_record.get("eventSourceARN", "")
-        if arn:
-            is_sqs = True
 
         # logic to deal with SNS => SQS event
         if "body" in first_record:
@@ -270,7 +260,8 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         msg_attributes = first_record.get("messageAttributes")
         if msg_attributes is None:
             sns_record = first_record.get("Sns") or {}
-            if not is_sqs:
+            # SNS->SQS event would extract SNS arn without this check
+            if event_source.equals(EventTypes.SNS):
                 arn = sns_record.get("TopicArn", "")
             msg_attributes = sns_record.get("MessageAttributes") or {}
         dd_payload = msg_attributes.get("_datadog")
@@ -304,7 +295,14 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
                 context = propagator.extract(dd_data)
-                _dsm_set_checkpoint(dd_data, "sqs" if is_sqs else "sns", arn)
+                # Do not want to set checkpoint with "" arn
+                if arn:
+                    _dsm_set_checkpoint(
+                        dd_data,
+                        # In this function only recieves SQS and SNS events, if not SQS must be SNS
+                        "sqs" if event_source.equals(EventTypes.SQS) else "sns",
+                        arn,
+                    )
                 return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
@@ -329,12 +327,23 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                             sampling_priority=float(x_ray_context["sampled"]),
                         )
         # Still want to set a DSM checkpoint even if DSM context not propagated
-        _dsm_set_checkpoint(None, "sqs" if is_sqs else "sns", arn)
+        # In this function only recieves SQS and SNS events, if not SQS must be SNS
+        # Do not want to set checkpoint with "" arn
+        if arn:
+            _dsm_set_checkpoint(
+                None, "sqs" if event_source.equals(EventTypes.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(None, "sqs" if is_sqs else "sns", arn)
+        # Do not want to set checkpoint with "" arn
+        if arn:
+            _dsm_set_checkpoint(
+                None,
+                "sqs" if event_source.equals(EventTypes.SQS) else "sns",
+                arn,
+            )
         return extract_context_from_lambda_context(lambda_context)
 
 
@@ -418,7 +427,9 @@ def extract_context_from_kinesis_event(event, 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, "kinesis", arn)
+    # Do not want to set checkpoint with "" arn
+    if arn:
+        _dsm_set_checkpoint(None, "kinesis", arn)
     return extract_context_from_lambda_context(lambda_context)
 
 

tests/test_tracing.py

@@ -2494,13 +2494,6 @@ 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(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_DSM_PROPAGATION_KEY_BASE_64_not_present(
         self, mock_checkpoint
@@ -2532,6 +2525,7 @@ def test_sqs_event_with_datadog_message_attributes(
                     "messageAttributes": {
                         "_datadog": {"dataType": "String", "stringValue": dd_json_data}
                     },
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2565,6 +2559,7 @@ def test_sqs_event_with_binary_datadog_message_attributes(
                     "messageAttributes": {
                         "_datadog": {"dataType": "Binary", "binaryValue": encoded_data}
                     },
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2600,6 +2595,7 @@ def test_sns_event_with_datadog_message_attributes(
                             "_datadog": {"Type": "String", "Value": dd_json_data}
                         },
                     },
+                    "eventSource": "aws:sns",
                 }
             ]
         }
@@ -2633,6 +2629,7 @@ def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present(
                     "messageAttributes": {
                         "_datadog": {"dataType": "String", "stringValue": dd_json_data}
                     },
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2674,6 +2671,7 @@ def test_sns_to_sqs_event_detection_and_processing(
                     "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
                     "body": json.dumps(sns_notification),
                     "messageAttributes": {},
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2701,6 +2699,7 @@ def test_sqs_event_without_datadog_message_attributes(
                 {
                     "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
                     "messageAttributes": {},
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2732,6 +2731,7 @@ def test_sqs_event_with_malformed_datadog_message_attributes(
                     "messageAttributes": {
                         "_datadog": {"dataType": "String", "stringValue": dd_json_data}
                     },
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2760,7 +2760,8 @@ def test_sns_event_without_datadog_message_attributes(
                     "Sns": {
                         "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
                         "MessageAttributes": {},
-                    }
+                    },
+                    "eventSource": "aws:sns",
                 }
             ]
         }
@@ -2793,7 +2794,8 @@ def test_sns_event_with_malformed_datadog_message_attributes(
                         "MessageAttributes": {
                             "_datadog": {"Type": "String", "Value": dd_json_data}
                         },
-                    }
+                    },
+                    "eventSource": "aws:sns",
                 }
             ]
         }
@@ -2834,6 +2836,7 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes(
                     "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
                     "body": json.dumps(sns_notification),
                     "messageAttributes": {},
+                    "eventSource": "aws:sqs",
                 }
             ]
         }
@@ -2866,7 +2869,7 @@ def test_sqs_sns_event_with_exception_accessing_first_record(
             event, self.lambda_context
         )
 
-        mock_dsm_set_checkpoint.assert_called_once_with(None, "sns", "")
+        mock_dsm_set_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
@@ -2982,6 +2985,6 @@ def test_kinesis_event_with_exception_accessing_first_record(
 
         # Verify that _dsm_set_checkpoint is called with empty string for arn
         # even when exception occurs
-        mock_dsm_set_checkpoint.assert_called_once_with(None, "kinesis", "")
+        mock_dsm_set_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)