moved arn logic to the wrapper

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -217,20 +217,16 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     """
 
     # EventBridge => SQS
-    is_sqs = False
     try:
         context = _extract_context_from_eventbridge_sqs_event(event)
         if _is_context_complete(context):
-            return context, None, None
+            return context, None
     except Exception:
         logger.debug("Failed extracting context as EventBridge to SQS.")
 
     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:
             body_str = first_record.get("body")
@@ -245,8 +241,6 @@ 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:
-                arn = sns_record.get("TopicArn", "")
             msg_attributes = sns_record.get("MessageAttributes") or {}
         dd_payload = msg_attributes.get("_datadog")
         if dd_payload:
@@ -277,13 +271,12 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                         return (
                             extract_context_from_step_functions(dd_data, None),
                             None,
-                            None,
                         )
                     except Exception:
                         logger.debug(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
-                return propagator.extract(dd_data), dd_data, arn
+                return propagator.extract(dd_data), dd_data
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
             # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -308,12 +301,11 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                                 sampling_priority=float(x_ray_context["sampled"]),
                             ),
                             None,
-                            None,
                         )
-        return extract_context_from_lambda_context(lambda_context), None, None
+        return extract_context_from_lambda_context(lambda_context), None
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
-        return extract_context_from_lambda_context(lambda_context), None, None
+        return extract_context_from_lambda_context(lambda_context), None
 
 
 def _extract_context_from_eventbridge_sqs_event(event):
@@ -374,7 +366,6 @@ def extract_context_from_kinesis_event(event, lambda_context):
     """
     try:
         record = get_first_record(event)
-        arn = record.get("eventSourceARN", "")
         kinesis = record.get("kinesis")
         if not kinesis:
             return extract_context_from_lambda_context(lambda_context)
@@ -388,11 +379,11 @@ 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 propagator.extract(dd_ctx), dd_ctx, arn
+                return propagator.extract(dd_ctx), dd_ctx
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
 
-    return extract_context_from_lambda_context(lambda_context), None, None
+    return extract_context_from_lambda_context(lambda_context), None
 
 
 def _deterministic_sha256_hash(s: str, part: str) -> int:
@@ -602,7 +593,6 @@ def extract_dd_trace_context(
     event_source = parse_event_source(event)
     context = None
     dd_json_data = None
-    arn = None
 
     if extractor is not None:
         context = extract_context_custom_extractor(extractor, event, lambda_context)
@@ -611,13 +601,13 @@ def extract_dd_trace_context(
             event, lambda_context, event_source, decode_authorizer_context
         )
     elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS):
-        context, dd_json_data, arn = extract_context_from_sqs_or_sns_event_or_context(
+        context, dd_json_data = extract_context_from_sqs_or_sns_event_or_context(
             event, lambda_context
         )
     elif event_source.equals(EventTypes.EVENTBRIDGE):
         context = extract_context_from_eventbridge_event(event, lambda_context)
     elif event_source.equals(EventTypes.KINESIS):
-        context, dd_json_data, arn = extract_context_from_kinesis_event(
+        context, dd_json_data = extract_context_from_kinesis_event(
             event, lambda_context
         )
     elif event_source.equals(EventTypes.STEPFUNCTIONS):
@@ -636,7 +626,7 @@ def extract_dd_trace_context(
         if dd_trace_context:
             trace_context_source = TraceContextSource.XRAY
     logger.debug("extracted dd trace context %s", dd_trace_context)
-    return dd_trace_context, trace_context_source, event_source, dd_json_data, arn
+    return dd_trace_context, trace_context_source, event_source, dd_json_data
 
 
 def get_dd_trace_context_obj():

datadog_lambda/wrapper.py

@@ -44,6 +44,7 @@
 from datadog_lambda.trigger import (
     extract_trigger_tags,
     extract_http_status_code_tag,
+    EventTypes,
 )
 
 if config.profiling_enabled:
@@ -220,7 +221,6 @@ def _before(self, event, context):
                 trace_context_source,
                 event_source,
                 dd_json_data,
-                arn,
             ) = extract_dd_trace_context(
                 event,
                 context,
@@ -247,7 +247,10 @@ def _before(self, event, context):
                     )
                 if config.data_streams_enabled:
                     if dd_json_data:
-                        set_dsm_checkpoint(dd_json_data, event_source.to_string(), arn)
+                        source_arn = extract_source_arn(event, event_source)
+                        set_dsm_checkpoint(
+                            dd_json_data, event_source.to_string(), source_arn
+                        )
 
                 self.span = create_function_execution_span(
                     context=context,
@@ -379,4 +382,13 @@ def carrier_get(key):
         logger.debug(f"Failed to set DSM checkpoint: {e}")
 
 
+def extract_source_arn(event, event_source):
+    if event_source.equals(EventTypes.SQS) or event_source.equals(EventTypes.KINESIS):
+        return event.get("Records", [{}])[0].get("eventSourceARN")
+    elif event_source.equals(EventTypes.SNS):
+        return event.get("Records", [{}])[0].get("Sns", {}).get("TopicArn")
+    else:
+        return None
+
+
 datadog_lambda_wrapper = _LambdaDecorator

tests/test_tracing.py

@@ -245,7 +245,7 @@ def test_extract_dd_trace_context(event, expect):
             event = json.load(f)
     ctx = get_mock_context()
 
-    actual, _, _, _, _ = extract_dd_trace_context(event, ctx)
+    actual, _, _, _ = extract_dd_trace_context(event, ctx)
     assert (expect is None) is (actual is None)
     assert (expect is None) or actual.trace_id == expect.trace_id
     assert (expect is None) or actual.span_id == expect.span_id
@@ -269,7 +269,7 @@ def tearDown(self):
     @with_trace_propagation_style("datadog")
     def test_without_datadog_trace_headers(self):
         lambda_ctx = get_mock_context()
-        ctx, source, _, _, _ = extract_dd_trace_context({}, lambda_ctx)
+        ctx, source, _, _ = extract_dd_trace_context({}, lambda_ctx)
         self.assertEqual(source, "xray")
         self.assertEqual(
             ctx,
@@ -292,7 +292,7 @@ def test_without_datadog_trace_headers(self):
     @with_trace_propagation_style("datadog")
     def test_with_non_object_event(self):
         lambda_ctx = get_mock_context()
-        ctx, source, _, _, _ = extract_dd_trace_context(b"", lambda_ctx)
+        ctx, source, _, _ = extract_dd_trace_context(b"", lambda_ctx)
         self.assertEqual(source, "xray")
         self.assertEqual(
             ctx,
@@ -315,7 +315,7 @@ def test_with_non_object_event(self):
     @with_trace_propagation_style("datadog")
     def test_with_incomplete_datadog_trace_headers(self):
         lambda_ctx = get_mock_context()
-        ctx, source, _, _, _ = extract_dd_trace_context(
+        ctx, source, _, _ = extract_dd_trace_context(
             {"headers": {TraceHeader.TRACE_ID: "123"}},
             lambda_ctx,
         )
@@ -340,7 +340,7 @@ def test_with_incomplete_datadog_trace_headers(self):
     def common_tests_with_trace_context_extraction_injection(
         self, headers, event_containing_headers, lambda_context=get_mock_context()
     ):
-        ctx, source, _, _, _ = extract_dd_trace_context(
+        ctx, source, _, _ = extract_dd_trace_context(
             event_containing_headers,
             lambda_context,
         )
@@ -393,7 +393,7 @@ def extractor_foo(event, context):
             return trace_id, parent_id, sampling_priority
 
         lambda_ctx = get_mock_context()
-        ctx, ctx_source, _, _, _ = extract_dd_trace_context(
+        ctx, ctx_source, _, _ = extract_dd_trace_context(
             {
                 "foo": {
                     TraceHeader.TRACE_ID: "123",
@@ -428,7 +428,7 @@ def extractor_raiser(event, context):
             raise Exception("kreator")
 
         lambda_ctx = get_mock_context()
-        ctx, ctx_source, _, _, _ = extract_dd_trace_context(
+        ctx, ctx_source, _, _ = extract_dd_trace_context(
             {
                 "foo": {
                     TraceHeader.TRACE_ID: "123",
@@ -627,7 +627,7 @@ def _test_step_function_trace_data_common(
             TraceHeader.TAGS: f"_dd.p.tid={expected_tid}",
         }
 
-        ctx, source, _, _, _ = extract_dd_trace_context(event, lambda_ctx)
+        ctx, source, _, _ = extract_dd_trace_context(event, lambda_ctx)
 
         self.assertEqual(source, "event")
         self.assertEqual(ctx, expected_context)
@@ -1148,7 +1148,7 @@ def test_mixed_parent_context_when_merging(self):
         # use the dd-trace trace-id and the x-ray parent-id
         # This allows parenting relationships like dd-trace -> x-ray -> dd-trace
         lambda_ctx = get_mock_context()
-        ctx, source, _, _, _ = extract_dd_trace_context(
+        ctx, source, _, _ = extract_dd_trace_context(
             {
                 "headers": {
                     TraceHeader.TRACE_ID: "123",
@@ -1174,7 +1174,7 @@ def test_set_dd_trace_py_root_no_span_id(self):
         os.environ["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320"
 
         lambda_ctx = get_mock_context()
-        ctx, source, _, _, _ = extract_dd_trace_context(
+        ctx, source, _, _ = extract_dd_trace_context(
             {
                 "headers": {
                     TraceHeader.TRACE_ID: "123",
@@ -2467,15 +2467,14 @@ def test_extract_context_from_sqs_event_with_datadog_context(self):
             ]
         }
 
-        (context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
+        (context, dd_json_data) = extract_context_from_sqs_or_sns_event_or_context(
             sqs_event, lambda_ctx
         )
 
         self.assertEqual(context.trace_id, 456)
         self.assertEqual(context.span_id, 789)
         self.assertEqual(context.sampling_priority, 1)
         self.assertEqual(dd_json_data, expected_trace_data)
-        self.assertEqual(arn, "arn:aws:sqs:us-east-1:123456789012:test-queue")
 
     def test_extract_context_from_sns_event_with_datadog_context(self):
         """Test SNS event extraction with datadog context returns expected values."""
@@ -2501,15 +2500,14 @@ def test_extract_context_from_sns_event_with_datadog_context(self):
             ]
         }
 
-        (context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
+        (context, dd_json_data) = extract_context_from_sqs_or_sns_event_or_context(
             sns_event, lambda_ctx
         )
 
         self.assertEqual(context.trace_id, 111)
         self.assertEqual(context.span_id, 222)
         self.assertEqual(context.sampling_priority, 2)
         self.assertEqual(dd_json_data, expected_trace_data)
-        self.assertEqual(arn, "arn:aws:sns:us-east-1:123456789012:test-topic")
 
     def test_extract_context_from_kinesis_event_with_datadog_context(self):
         """Test Kinesis event extraction with datadog context returns expected values."""
@@ -2531,16 +2529,13 @@ def test_extract_context_from_kinesis_event_with_datadog_context(self):
             ]
         }
 
-        context, dd_json_data, arn = extract_context_from_kinesis_event(
+        context, dd_json_data = extract_context_from_kinesis_event(
             kinesis_event, lambda_ctx
         )
 
         self.assertEqual(context.trace_id, 333)
         self.assertEqual(context.span_id, 444)
         self.assertEqual(dd_json_data, expected_trace_data)
-        self.assertEqual(
-            arn, "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
-        )
 
     def test_step_function_context_in_sqs_calls_step_function_extractor(self):
         """Test that Step Function context in SQS calls the step function extractor."""
@@ -2564,13 +2559,12 @@ def test_step_function_context_in_sqs_calls_step_function_extractor(self):
             ]
         }
 
-        (context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
+        (context, dd_json_data) = extract_context_from_sqs_or_sns_event_or_context(
             sfn_sqs_event, lambda_ctx
         )
 
         self.assertIsNotNone(context)
         self.assertIsNone(dd_json_data)
-        self.assertIsNone(arn)
 
     def test_extract_context_from_sns_to_sqs_event_with_datadog_context(self):
         """Test SNS -> SQS event extraction with datadog context returns expected values."""
@@ -2605,12 +2599,11 @@ def test_extract_context_from_sns_to_sqs_event_with_datadog_context(self):
             ]
         }
 
-        (context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
+        (context, dd_json_data) = extract_context_from_sqs_or_sns_event_or_context(
             sns_to_sqs_event, lambda_ctx
         )
 
         self.assertEqual(context.trace_id, 555)
         self.assertEqual(context.span_id, 666)
         self.assertEqual(context.sampling_priority, 1)
         self.assertEqual(dd_json_data, expected_trace_data)
-        self.assertEqual(arn, "arn:aws:sqs:us-east-1:123456789012:test-queue")