From 18bcf3c1034c545a710deeb5354fbd787367de6d Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 19 Jun 2025 11:27:25 -0400 Subject: [PATCH 01/42] removed current dsm implementation --- datadog_lambda/dsm.py | 38 ------------- datadog_lambda/wrapper.py | 5 +- tests/test_dsm.py | 112 -------------------------------------- 3 files changed, 2 insertions(+), 153 deletions(-) delete mode 100644 datadog_lambda/dsm.py diff --git a/datadog_lambda/dsm.py b/datadog_lambda/dsm.py deleted file mode 100644 index 427f5e47..00000000 --- a/datadog_lambda/dsm.py +++ /dev/null @@ -1,38 +0,0 @@ -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 import data_streams_processor - from ddtrace.internal.datastreams.processor import DsmPathwayCodec - from ddtrace.internal.datastreams.botocore import ( - get_datastreams_context, - calculate_sqs_payload_size, - ) - - records = event.get("Records") - if records is None: - return - processor = data_streams_processor() - - 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", f"topic:{queue_arn}", "type:sqs"], - payload_size=payload_size, - ) - except Exception as e: - logger.error(format_err_with_traceback(e)) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 98a1e87e..eeb1bdf0 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -10,7 +10,8 @@ from time import time_ns from datadog_lambda.asm import asm_start_response, asm_start_request -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, @@ -240,8 +241,6 @@ def _before(self, event, context): self.inferred_span = create_inferred_span( event, context, event_source, config.decode_authorizer_context ) - if config.data_streams_enabled: - set_dsm_context(event, event_source) self.span = create_function_execution_span( context=context, function_name=config.function_name, diff --git a/tests/test_dsm.py b/tests/test_dsm.py index 544212d8..e69de29b 100644 --- a/tests/test_dsm.py +++ b/tests/test_dsm.py @@ -1,112 +0,0 @@ -import unittest -from unittest.mock import patch, MagicMock - -from datadog_lambda.dsm import set_dsm_context, _dsm_set_sqs_context -from datadog_lambda.trigger import EventTypes, _EventSource - - -class TestDsmSQSContext(unittest.TestCase): - def setUp(self): - patcher = patch("datadog_lambda.dsm._dsm_set_sqs_context") - self.mock_dsm_set_sqs_context = patcher.start() - self.addCleanup(patcher.stop) - - patcher = patch("ddtrace.internal.datastreams.data_streams_processor") - self.mock_data_streams_processor = patcher.start() - self.addCleanup(patcher.stop) - - patcher = patch("ddtrace.internal.datastreams.botocore.get_datastreams_context") - self.mock_get_datastreams_context = patcher.start() - self.mock_get_datastreams_context.return_value = {} - self.addCleanup(patcher.stop) - - patcher = patch( - "ddtrace.internal.datastreams.botocore.calculate_sqs_payload_size" - ) - self.mock_calculate_sqs_payload_size = patcher.start() - self.mock_calculate_sqs_payload_size.return_value = 100 - self.addCleanup(patcher.stop) - - patcher = patch("ddtrace.internal.datastreams.processor.DsmPathwayCodec.decode") - self.mock_dsm_pathway_codec_decode = patcher.start() - self.addCleanup(patcher.stop) - - def test_non_sqs_event_source_does_nothing(self): - """Test that non-SQS event sources don't trigger DSM context setting""" - event = {} - # Use Unknown Event Source - event_source = _EventSource(EventTypes.UNKNOWN) - set_dsm_context(event, event_source) - - # DSM context should not be set for non-SQS events - self.mock_dsm_set_sqs_context.assert_not_called() - - def test_sqs_event_with_no_records_does_nothing(self): - """Test that events where Records is None don't trigger DSM processing""" - events_with_no_records = [ - {}, - {"Records": None}, - {"someOtherField": "value"}, - ] - - for event in events_with_no_records: - _dsm_set_sqs_context(event) - self.mock_data_streams_processor.assert_not_called() - - def test_sqs_event_triggers_dsm_sqs_context(self): - """Test that SQS event sources trigger the SQS-specific DSM context function""" - sqs_event = { - "Records": [ - { - "eventSource": "aws:sqs", - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:my-queue", - "body": "Hello from SQS!", - } - ] - } - - event_source = _EventSource(EventTypes.SQS) - set_dsm_context(sqs_event, event_source) - - self.mock_dsm_set_sqs_context.assert_called_once_with(sqs_event) - - def test_sqs_multiple_records_process_each_record(self): - """Test that each record in an SQS event gets processed individually""" - multi_record_event = { - "Records": [ - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue1", - "body": "Message 1", - }, - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue2", - "body": "Message 2", - }, - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue3", - "body": "Message 3", - }, - ] - } - - mock_context = MagicMock() - self.mock_dsm_pathway_codec_decode.return_value = mock_context - - _dsm_set_sqs_context(multi_record_event) - - self.assertEqual(mock_context.set_checkpoint.call_count, 3) - - calls = mock_context.set_checkpoint.call_args_list - expected_arns = [ - "arn:aws:sqs:us-east-1:123456789012:queue1", - "arn:aws:sqs:us-east-1:123456789012:queue2", - "arn:aws:sqs:us-east-1:123456789012:queue3", - ] - - for i, call in enumerate(calls): - args, kwargs = call - tags = args[0] - self.assertIn("direction:in", tags) - self.assertIn(f"topic:{expected_arns[i]}", tags) - self.assertIn("type:sqs", tags) - self.assertEqual(kwargs["payload_size"], 100) From af174d4d7f0c7503719d7571a6084880ac458a37 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 19 Jun 2025 11:28:25 -0400 Subject: [PATCH 02/42] new dsm lambda implementation --- datadog_lambda/tracing.py | 50 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 89a4126b..e96693df 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -50,6 +50,9 @@ set_tracer_provider(TracerProvider()) +data_streams_enabled = ( + os.environ.get("DD_DATA_STREAMS_ENABLED", "false").lower() == "true" +) logger = logging.getLogger(__name__) @@ -67,6 +70,37 @@ LOWER_64_BITS = "LOWER_64_BITS" +def _extract_context(context_json, event_type, arn): + from ddtrace.data_streams import set_consume_checkpoint + + """ + Extracts the context from a JSON carrier and optionally sets a consume checkpoint + if the context is complete and data streams are enabled. + """ + context = propagator.extract(context_json) + + if not _is_context_complete(context): + return context + + if not data_streams_enabled: + return context + try: + 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): + def carrier_get(key): + return context_json.get(key) + + return carrier_get + + def _convert_xray_trace_id(xray_trace_id): """ Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int). @@ -215,6 +249,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. """ + is_sqs = False + arn = None # EventBridge => SQS try: @@ -226,6 +262,9 @@ 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: @@ -272,8 +311,12 @@ 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 propagator.extract(dd_data) + if is_sqs: + return _extract_context(dd_data, "sqs", arn) + else: + return _extract_context( + dd_data, "sns", sns_record.get(("TopicArn")) + ) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 @@ -360,6 +403,7 @@ 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) @@ -373,7 +417,7 @@ 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) + return _extract_context(dd_ctx, "kinesis", arn) except Exception as e: logger.debug("The trace extractor returned with error %s", e) From eeea762cd855e80034db265d8b9d056c344c49e9 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 19 Jun 2025 17:01:39 -0400 Subject: [PATCH 03/42] check env variable in proper way --- datadog_lambda/tracing.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index e96693df..2fdff2c3 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -50,10 +50,6 @@ set_tracer_provider(TracerProvider()) -data_streams_enabled = ( - os.environ.get("DD_DATA_STREAMS_ENABLED", "false").lower() == "true" -) - logger = logging.getLogger(__name__) dd_trace_context = None @@ -82,7 +78,7 @@ def _extract_context(context_json, event_type, arn): if not _is_context_complete(context): return context - if not data_streams_enabled: + if not config.data_streams_enabled: return context try: carrier_get = _create_carrier_get(context_json) From 1c7807218b8dc77f552b1446fbc3096e1d7bbf53 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 19 Jun 2025 17:01:45 -0400 Subject: [PATCH 04/42] add tests --- tests/test_dsm.py | 253 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 253 insertions(+) diff --git a/tests/test_dsm.py b/tests/test_dsm.py index e69de29b..13d9a134 100644 --- a/tests/test_dsm.py +++ b/tests/test_dsm.py @@ -0,0 +1,253 @@ +import json +import unittest +import base64 +import os +from unittest.mock import patch + +from ddtrace.trace import Context + +from datadog_lambda.tracing import ( + _extract_context, + _create_carrier_get, + extract_context_from_sqs_or_sns_event_or_context, + extract_context_from_kinesis_event, +) +from tests.utils import get_mock_context + + +class TestExtractContext(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) + + logger_patcher = patch("datadog_lambda.tracing.logger") + self.mock_logger = logger_patcher.start() + self.addCleanup(logger_patcher.stop) + + def test_extract_context_data_streams_disabled(self): + with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'false'}): + 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 + + result = _extract_context(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) + + def test_extract_context_data_streams_enabled_complete_context(self): + with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + 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 + + result = _extract_context(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) + + def test_extract_context_data_streams_enabled_incomplete_context(self): + with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + 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=None, sampling_priority=1) + self.mock_extract.return_value = mock_context + + result = _extract_context(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) + + def test_extract_context_exception_path(self): + with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + 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(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): + def test_create_carrier_get_with_valid_data(self): + context_json = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1" + } + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") + self.assertEqual(carrier_get("x-datadog-parent-id"), "67890") + self.assertEqual(carrier_get("x-datadog-sampling-priority"), "1") + + def test_create_carrier_get_with_missing_key(self): + context_json = {"x-datadog-trace-id": "12345"} + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") + self.assertIsNone(carrier_get("x-datadog-parent-id")) + + def test_create_carrier_get_with_empty_context(self): + context_json = {} + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertIsNone(carrier_get("any-key")) + + +class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase): + def setUp(self): + self.lambda_context = get_mock_context() + + @patch("datadog_lambda.tracing._extract_context") + def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [{ + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": { + "dataType": "String", + "stringValue": dd_json_data + } + } + }] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + + mock_extract_context.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") + def test_sqs_event_with_binary_datadog_message_attributes(self, mock_extract_context): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + encoded_data = base64.b64encode(dd_json_data.encode()).decode() + + event = { + "Records": [{ + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": { + "dataType": "Binary", + "binaryValue": encoded_data + } + } + }] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + + mock_extract_context.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") + def test_sns_event_with_datadog_message_attributes(self, mock_extract_context): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [{ + "eventSourceARN": "", + "Sns": { + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": { + "Type": "String", + "Value": dd_json_data + } + } + } + }] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + + mock_extract_context.assert_called_once_with( + dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + ) + self.assertEqual(result, mock_context) + + +class TestExtractContextFromKinesisEvent(unittest.TestCase): + def setUp(self): + self.lambda_context = get_mock_context() + + @patch("datadog_lambda.tracing._extract_context") + def test_kinesis_event_with_datadog_data(self, mock_extract_context): + dd_data = {"dd-pathway-ctx-base64": "12345"} + kinesis_data = {"_datadog": dd_data, "message": "test"} + kinesis_data_str = json.dumps(kinesis_data) + encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() + + event = { + "Records": [{ + "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + "kinesis": { + "data": encoded_data + } + }] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context.return_value = mock_context + + result = extract_context_from_kinesis_event(event, self.lambda_context) + + mock_extract_context.assert_called_once_with( + dd_data, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream" + ) + self.assertEqual(result, mock_context) From 7a392909c4188901990405841c53bb90304c98ac Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 09:44:15 -0400 Subject: [PATCH 05/42] add more detailed comment to function --- datadog_lambda/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 2fdff2c3..e9b45025 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -70,7 +70,7 @@ def _extract_context(context_json, event_type, arn): from ddtrace.data_streams import set_consume_checkpoint """ - Extracts the context from a JSON carrier and optionally sets a 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) From e7d8a72ac2d2cf2d4f0096fa8d93ad5aafa5c0c6 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 09:47:34 -0400 Subject: [PATCH 06/42] fixed lint for tests --- tests/test_dsm.py | 91 +++++++++++++++++++++++++---------------------- 1 file changed, 49 insertions(+), 42 deletions(-) diff --git a/tests/test_dsm.py b/tests/test_dsm.py index 13d9a134..576fdea7 100644 --- a/tests/test_dsm.py +++ b/tests/test_dsm.py @@ -30,7 +30,7 @@ def setUp(self): self.addCleanup(logger_patcher.stop) def test_extract_context_data_streams_disabled(self): - with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'false'}): + with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "false"}): context_json = {"dd-pathway-ctx-base64": "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -45,7 +45,7 @@ def test_extract_context_data_streams_disabled(self): self.assertEqual(result, mock_context) def test_extract_context_data_streams_enabled_complete_context(self): - with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): context_json = {"dd-pathway-ctx-base64": "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -65,7 +65,7 @@ def test_extract_context_data_streams_enabled_complete_context(self): self.assertEqual(result, mock_context) def test_extract_context_data_streams_enabled_incomplete_context(self): - with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): context_json = {"dd-pathway-ctx-base64": "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -80,7 +80,7 @@ def test_extract_context_data_streams_enabled_incomplete_context(self): self.assertEqual(result, mock_context) def test_extract_context_exception_path(self): - with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}): + with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): context_json = {"dd-pathway-ctx-base64": "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -104,7 +104,7 @@ def test_create_carrier_get_with_valid_data(self): context_json = { "x-datadog-trace-id": "12345", "x-datadog-parent-id": "67890", - "x-datadog-sampling-priority": "1" + "x-datadog-sampling-priority": "1", } carrier_get = _create_carrier_get(context_json) @@ -142,21 +142,22 @@ def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context): dd_json_data = json.dumps(dd_data) event = { - "Records": [{ - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", - "messageAttributes": { - "_datadog": { - "dataType": "String", - "stringValue": dd_json_data - } + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "String", "stringValue": dd_json_data} + }, } - }] + ] } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context + ) mock_extract_context.assert_called_once_with( dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -164,27 +165,30 @@ def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context): self.assertEqual(result, mock_context) @patch("datadog_lambda.tracing._extract_context") - def test_sqs_event_with_binary_datadog_message_attributes(self, mock_extract_context): + def test_sqs_event_with_binary_datadog_message_attributes( + self, mock_extract_context + ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) encoded_data = base64.b64encode(dd_json_data.encode()).decode() event = { - "Records": [{ - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", - "messageAttributes": { - "_datadog": { - "dataType": "Binary", - "binaryValue": encoded_data - } + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "Binary", "binaryValue": encoded_data} + }, } - }] + ] } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context + ) mock_extract_context.assert_called_once_with( dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -197,24 +201,25 @@ def test_sns_event_with_datadog_message_attributes(self, mock_extract_context): dd_json_data = json.dumps(dd_data) event = { - "Records": [{ - "eventSourceARN": "", - "Sns": { - "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", - "MessageAttributes": { - "_datadog": { - "Type": "String", - "Value": dd_json_data - } - } + "Records": [ + { + "eventSourceARN": "", + "Sns": { + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": {"Type": "String", "Value": dd_json_data} + }, + }, } - }] + ] } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context) + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context + ) mock_extract_context.assert_called_once_with( dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" @@ -234,12 +239,12 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context): encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() event = { - "Records": [{ - "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", - "kinesis": { - "data": encoded_data + "Records": [ + { + "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + "kinesis": {"data": encoded_data}, } - }] + ] } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) @@ -248,6 +253,8 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context): result = extract_context_from_kinesis_event(event, self.lambda_context) mock_extract_context.assert_called_once_with( - dd_data, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream" + dd_data, + "kinesis", + "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", ) self.assertEqual(result, mock_context) From 6e7e5eb9d6e9c9532b33d515e4cd74c3aaa502af Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 10:00:16 -0400 Subject: [PATCH 07/42] remove not needed tests --- tests/test_wrapper.py | 60 ------------------------------------------- 1 file changed, 60 deletions(-) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index f0240905..57d2eb77 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -77,10 +77,6 @@ def setUp(self): self.mock_dd_lambda_layer_tag = patcher.start() self.addCleanup(patcher.stop) - patcher = patch("datadog_lambda.wrapper.set_dsm_context") - self.mock_set_dsm_context = patcher.start() - self.addCleanup(patcher.stop) - @patch("datadog_lambda.config.Config.trace_enabled", False) def test_datadog_lambda_wrapper(self): @wrapper.datadog_lambda_wrapper @@ -562,62 +558,6 @@ def return_type_test(event, context): self.assertEqual(result, test_result) self.assertFalse(MockPrintExc.called) - def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self): - os.environ["DD_DATA_STREAMS_ENABLED"] = "true" - os.environ["DD_TRACE_ENABLED"] = "true" - - @wrapper.datadog_lambda_wrapper - def lambda_handler(event, context): - return "ok" - - result = lambda_handler({}, get_mock_context()) - self.assertEqual(result, "ok") - self.mock_set_dsm_context.assert_called_once() - - del os.environ["DD_DATA_STREAMS_ENABLED"] - - def test_set_dsm_context_not_called_when_only_DSM_enabled(self): - os.environ["DD_DATA_STREAMS_ENABLED"] = "true" - os.environ["DD_TRACE_ENABLED"] = "false" - - @wrapper.datadog_lambda_wrapper - def lambda_handler(event, context): - return "ok" - - result = lambda_handler({}, get_mock_context()) - self.assertEqual(result, "ok") - self.mock_set_dsm_context.assert_not_called() - - del os.environ["DD_DATA_STREAMS_ENABLED"] - - def test_set_dsm_context_not_called_when_only_tracing_enabled(self): - os.environ["DD_DATA_STREAMS_ENABLED"] = "false" - os.environ["DD_TRACE_ENABLED"] = "true" - - @wrapper.datadog_lambda_wrapper - def lambda_handler(event, context): - return "ok" - - result = lambda_handler({}, get_mock_context()) - self.assertEqual(result, "ok") - self.mock_set_dsm_context.assert_not_called() - - del os.environ["DD_DATA_STREAMS_ENABLED"] - - def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self): - os.environ["DD_DATA_STREAMS_ENABLED"] = "false" - os.environ["DD_TRACE_ENABLED"] = "false" - - @wrapper.datadog_lambda_wrapper - def lambda_handler(event, context): - return "ok" - - result = lambda_handler({}, get_mock_context()) - self.assertEqual(result, "ok") - self.mock_set_dsm_context.assert_not_called() - - del os.environ["DD_DATA_STREAMS_ENABLED"] - class TestLambdaWrapperWithTraceContext(unittest.TestCase): xray_root = "1-5e272390-8c398be037738dc042009320" From 4733bb05deea60f7ac0a38505b8dbc5d70fc9fc3 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 10:21:02 -0400 Subject: [PATCH 08/42] fix --- datadog_lambda/tracing.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index e9b45025..f4da9b82 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -309,10 +309,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): ) if is_sqs: return _extract_context(dd_data, "sqs", arn) - else: - return _extract_context( - dd_data, "sns", sns_record.get(("TopicArn")) - ) + return _extract_context(dd_data, "sns", sns_record.get(("TopicArn"))) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 From fd395e3132d31d2c9821ea2381447b46ea9e96f2 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 10:28:08 -0400 Subject: [PATCH 09/42] error handling --- datadog_lambda/tracing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index f4da9b82..2e84d183 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -309,7 +309,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): ) if is_sqs: return _extract_context(dd_data, "sqs", arn) - return _extract_context(dd_data, "sns", sns_record.get(("TopicArn"))) + return _extract_context( + dd_data, "sns", sns_record.get(("TopicArn"), "") + ) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 From 98355730f486427cc726f4dc17c87d03f30650d4 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 10:29:52 -0400 Subject: [PATCH 10/42] fix --- datadog_lambda/tracing.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 2e84d183..42a2b949 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -309,9 +309,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): ) if is_sqs: return _extract_context(dd_data, "sqs", arn) - return _extract_context( - dd_data, "sns", sns_record.get(("TopicArn"), "") - ) + return _extract_context(dd_data, "sns", sns_record.get("TopicArn", "")) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 From 3bbe48dc92692d7ac8b5480b3757f33c6f277a36 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 15:44:54 -0400 Subject: [PATCH 11/42] some fixes --- datadog_lambda/tracing.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 42a2b949..379c5f4b 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -75,10 +75,10 @@ def _extract_context(context_json, event_type, arn): """ context = propagator.extract(context_json) - if not _is_context_complete(context): + if not config.data_streams_enabled: return context - if not config.data_streams_enabled: + if not _is_context_complete(context): return context try: carrier_get = _create_carrier_get(context_json) @@ -246,7 +246,6 @@ 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. """ is_sqs = False - arn = None # EventBridge => SQS try: @@ -276,6 +275,7 @@ 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 {} + arn = sns_record.get("TopicArn", "") msg_attributes = sns_record.get("MessageAttributes") or {} dd_payload = msg_attributes.get("_datadog") if dd_payload: @@ -307,9 +307,7 @@ 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." ) - if is_sqs: - return _extract_context(dd_data, "sqs", arn) - return _extract_context(dd_data, "sns", sns_record.get("TopicArn", "")) + return _extract_context(dd_data, "sqs" if is_sqs else "sns", arn) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 From f533c0987fc20add7d3dcecf6a86ab085e9afa83 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 16:08:19 -0400 Subject: [PATCH 12/42] renamed to extract_context_with_datastreams --- datadog_lambda/tracing.py | 8 ++++--- tests/test_dsm.py | 48 ++++++++++++++++++++++----------------- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 379c5f4b..79951b1a 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -66,7 +66,7 @@ LOWER_64_BITS = "LOWER_64_BITS" -def _extract_context(context_json, event_type, arn): +def _extract_context_with_data_streams(context_json, event_type, arn): from ddtrace.data_streams import set_consume_checkpoint """ @@ -307,7 +307,9 @@ 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(dd_data, "sqs" if is_sqs else "sns", arn) + return _extract_context_with_data_streams( + dd_data, "sqs" if is_sqs else "sns", arn + ) else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 @@ -408,7 +410,7 @@ 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(dd_ctx, "kinesis", arn) + return _extract_context_with_data_streams(dd_ctx, "kinesis", arn) except Exception as e: logger.debug("The trace extractor returned with error %s", e) diff --git a/tests/test_dsm.py b/tests/test_dsm.py index 576fdea7..c5319a82 100644 --- a/tests/test_dsm.py +++ b/tests/test_dsm.py @@ -7,7 +7,7 @@ from ddtrace.trace import Context from datadog_lambda.tracing import ( - _extract_context, + _extract_context_with_data_streams, _create_carrier_get, extract_context_from_sqs_or_sns_event_or_context, extract_context_from_kinesis_event, @@ -38,7 +38,7 @@ def test_extract_context_data_streams_disabled(self): mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) self.mock_extract.return_value = mock_context - result = _extract_context(context_json, event_type, arn) + result = _extract_context_with_data_streams(context_json, event_type, arn) self.mock_extract.assert_called_once_with(context_json) self.mock_checkpoint.assert_not_called() @@ -53,7 +53,7 @@ def test_extract_context_data_streams_enabled_complete_context(self): mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) self.mock_extract.return_value = mock_context - result = _extract_context(context_json, event_type, arn) + result = _extract_context_with_data_streams(context_json, event_type, arn) self.mock_extract.assert_called_once_with(context_json) self.mock_checkpoint.assert_called_once() @@ -73,7 +73,7 @@ def test_extract_context_data_streams_enabled_incomplete_context(self): mock_context = Context(trace_id=12345, span_id=None, sampling_priority=1) self.mock_extract.return_value = mock_context - result = _extract_context(context_json, event_type, arn) + result = _extract_context_with_data_streams(context_json, event_type, arn) self.mock_extract.assert_called_once_with(context_json) self.mock_checkpoint.assert_not_called() @@ -91,7 +91,7 @@ def test_extract_context_exception_path(self): test_exception = Exception("Test exception") self.mock_checkpoint.side_effect = test_exception - result = _extract_context(context_json, event_type, arn) + result = _extract_context_with_data_streams(context_json, event_type, arn) self.mock_extract.assert_called_once_with(context_json) self.mock_checkpoint.assert_called_once() @@ -136,8 +136,10 @@ class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() - @patch("datadog_lambda.tracing._extract_context") - def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context): + @patch("datadog_lambda.tracing._extract_context_with_data_streams") + def test_sqs_event_with_datadog_message_attributes( + self, mock_extract_context_with_data_streams + ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -153,20 +155,20 @@ def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context): } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context.return_value = mock_context + mock_extract_context_with_data_streams.return_value = mock_context result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) - mock_extract_context.assert_called_once_with( + mock_extract_context_with_data_streams.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") + @patch("datadog_lambda.tracing._extract_context_with_data_streams") def test_sqs_event_with_binary_datadog_message_attributes( - self, mock_extract_context + self, mock_extract_context_with_data_streams ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -184,19 +186,21 @@ def test_sqs_event_with_binary_datadog_message_attributes( } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context.return_value = mock_context + mock_extract_context_with_data_streams.return_value = mock_context result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) - mock_extract_context.assert_called_once_with( + mock_extract_context_with_data_streams.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") - def test_sns_event_with_datadog_message_attributes(self, mock_extract_context): + @patch("datadog_lambda.tracing._extract_context_with_data_streams") + def test_sns_event_with_datadog_message_attributes( + self, mock_extract_context_with_data_streams + ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -215,13 +219,13 @@ def test_sns_event_with_datadog_message_attributes(self, mock_extract_context): } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context.return_value = mock_context + mock_extract_context_with_data_streams.return_value = mock_context result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) - mock_extract_context.assert_called_once_with( + mock_extract_context_with_data_streams.assert_called_once_with( dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" ) self.assertEqual(result, mock_context) @@ -231,8 +235,10 @@ class TestExtractContextFromKinesisEvent(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() - @patch("datadog_lambda.tracing._extract_context") - def test_kinesis_event_with_datadog_data(self, mock_extract_context): + @patch("datadog_lambda.tracing._extract_context_with_data_streams") + def test_kinesis_event_with_datadog_data( + self, mock_extract_context_with_data_streams + ): dd_data = {"dd-pathway-ctx-base64": "12345"} kinesis_data = {"_datadog": dd_data, "message": "test"} kinesis_data_str = json.dumps(kinesis_data) @@ -248,11 +254,11 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context): } mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context.return_value = mock_context + mock_extract_context_with_data_streams.return_value = mock_context result = extract_context_from_kinesis_event(event, self.lambda_context) - mock_extract_context.assert_called_once_with( + mock_extract_context_with_data_streams.assert_called_once_with( dd_data, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", From 86cf5ff989fb7dd521632ed7c33ccdbe1001b628 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 17:15:46 -0400 Subject: [PATCH 13/42] changed to explicit check of dsm context --- datadog_lambda/tracing.py | 2 +- tests/test_dsm.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 79951b1a..f40b2111 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -78,7 +78,7 @@ def _extract_context_with_data_streams(context_json, event_type, arn): if not config.data_streams_enabled: return context - if not _is_context_complete(context): + if "dd-pathway-ctx-base64" not in context_json: return context try: carrier_get = _create_carrier_get(context_json) diff --git a/tests/test_dsm.py b/tests/test_dsm.py index c5319a82..420d39ea 100644 --- a/tests/test_dsm.py +++ b/tests/test_dsm.py @@ -64,13 +64,13 @@ def test_extract_context_data_streams_enabled_complete_context(self): self.assertEqual(kwargs["manual_checkpoint"], False) self.assertEqual(result, mock_context) - def test_extract_context_data_streams_enabled_incomplete_context(self): + def test_extract_context_data_streams_enabled_invalid_context(self): with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): - context_json = {"dd-pathway-ctx-base64": "12345"} + 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=None, sampling_priority=1) + mock_context = Context(trace_id=12345, span_id=12345, sampling_priority=1) self.mock_extract.return_value = mock_context result = _extract_context_with_data_streams(context_json, event_type, arn) From 39547146dc569f4768e6757f09f541c9b455cd5f Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 17:58:30 -0400 Subject: [PATCH 14/42] move dsm tests to test_tracing.py --- tests/test_dsm.py | 266 ------------------------------------------ tests/test_tracing.py | 256 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 266 deletions(-) delete mode 100644 tests/test_dsm.py diff --git a/tests/test_dsm.py b/tests/test_dsm.py deleted file mode 100644 index 420d39ea..00000000 --- a/tests/test_dsm.py +++ /dev/null @@ -1,266 +0,0 @@ -import json -import unittest -import base64 -import os -from unittest.mock import patch - -from ddtrace.trace import Context - -from datadog_lambda.tracing import ( - _extract_context_with_data_streams, - _create_carrier_get, - extract_context_from_sqs_or_sns_event_or_context, - extract_context_from_kinesis_event, -) -from tests.utils import get_mock_context - - -class TestExtractContext(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) - - logger_patcher = patch("datadog_lambda.tracing.logger") - self.mock_logger = logger_patcher.start() - self.addCleanup(logger_patcher.stop) - - def test_extract_context_data_streams_disabled(self): - with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "false"}): - 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 - - result = _extract_context_with_data_streams(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) - - def test_extract_context_data_streams_enabled_complete_context(self): - with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): - 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 - - result = _extract_context_with_data_streams(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) - - def test_extract_context_data_streams_enabled_invalid_context(self): - with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): - 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 - - result = _extract_context_with_data_streams(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) - - def test_extract_context_exception_path(self): - with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}): - 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) - - 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): - def test_create_carrier_get_with_valid_data(self): - context_json = { - "x-datadog-trace-id": "12345", - "x-datadog-parent-id": "67890", - "x-datadog-sampling-priority": "1", - } - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") - self.assertEqual(carrier_get("x-datadog-parent-id"), "67890") - self.assertEqual(carrier_get("x-datadog-sampling-priority"), "1") - - def test_create_carrier_get_with_missing_key(self): - context_json = {"x-datadog-trace-id": "12345"} - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") - self.assertIsNone(carrier_get("x-datadog-parent-id")) - - def test_create_carrier_get_with_empty_context(self): - context_json = {} - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertIsNone(carrier_get("any-key")) - - -class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase): - def setUp(self): - self.lambda_context = get_mock_context() - - @patch("datadog_lambda.tracing._extract_context_with_data_streams") - def test_sqs_event_with_datadog_message_attributes( - self, mock_extract_context_with_data_streams - ): - dd_data = {"dd-pathway-ctx-base64": "12345"} - dd_json_data = json.dumps(dd_data) - - event = { - "Records": [ - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", - "messageAttributes": { - "_datadog": {"dataType": "String", "stringValue": dd_json_data} - }, - } - ] - } - - mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context_with_data_streams.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( - 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") - def test_sqs_event_with_binary_datadog_message_attributes( - self, mock_extract_context_with_data_streams - ): - dd_data = {"dd-pathway-ctx-base64": "12345"} - dd_json_data = json.dumps(dd_data) - encoded_data = base64.b64encode(dd_json_data.encode()).decode() - - event = { - "Records": [ - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", - "messageAttributes": { - "_datadog": {"dataType": "Binary", "binaryValue": encoded_data} - }, - } - ] - } - - mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context_with_data_streams.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( - 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") - def test_sns_event_with_datadog_message_attributes( - self, mock_extract_context_with_data_streams - ): - dd_data = {"dd-pathway-ctx-base64": "12345"} - dd_json_data = json.dumps(dd_data) - - event = { - "Records": [ - { - "eventSourceARN": "", - "Sns": { - "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", - "MessageAttributes": { - "_datadog": {"Type": "String", "Value": dd_json_data} - }, - }, - } - ] - } - - mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context_with_data_streams.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( - dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" - ) - self.assertEqual(result, mock_context) - - -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 - ): - dd_data = {"dd-pathway-ctx-base64": "12345"} - kinesis_data = {"_datadog": dd_data, "message": "test"} - kinesis_data_str = json.dumps(kinesis_data) - encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() - - event = { - "Records": [ - { - "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", - "kinesis": {"data": encoded_data}, - } - ] - } - - mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract_context_with_data_streams.return_value = mock_context - - result = extract_context_from_kinesis_event(event, self.lambda_context) - - mock_extract_context_with_data_streams.assert_called_once_with( - dd_data, - "kinesis", - "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", - ) - self.assertEqual(result, mock_context) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index a629343e..80d0a0db 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1,3 +1,4 @@ +import base64 import copy import functools import json @@ -41,6 +42,10 @@ 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, ) from tests.utils import get_mock_context @@ -2438,3 +2443,254 @@ def test_exception_outside_handler_tracing_disabled( mock_submit_errors_metric.assert_called_once_with(None) mock_trace.assert_not_called() + + +class TestExtractContext(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) + + logger_patcher = patch("datadog_lambda.tracing.logger") + self.mock_logger = logger_patcher.start() + self.addCleanup(logger_patcher.stop) + + @patch("datadog_lambda.config.Config.data_streams_enabled", False) + def test_extract_context_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 + + result = _extract_context_with_data_streams(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): + 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 + + result = _extract_context_with_data_streams(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): + 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 + + result = _extract_context_with_data_streams(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): + 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) + + 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): + def test_create_carrier_get_with_valid_data(self): + context_json = { + "x-datadog-trace-id": "12345", + "x-datadog-parent-id": "67890", + "x-datadog-sampling-priority": "1", + } + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") + self.assertEqual(carrier_get("x-datadog-parent-id"), "67890") + self.assertEqual(carrier_get("x-datadog-sampling-priority"), "1") + + def test_create_carrier_get_with_missing_key(self): + context_json = {"x-datadog-trace-id": "12345"} + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") + self.assertIsNone(carrier_get("x-datadog-parent-id")) + + def test_create_carrier_get_with_empty_context(self): + context_json = {} + + carrier_get = _create_carrier_get(context_json) + + self.assertTrue(callable(carrier_get)) + self.assertIsNone(carrier_get("any-key")) + + +class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase): + def setUp(self): + self.lambda_context = get_mock_context() + + @patch("datadog_lambda.tracing._extract_context_with_data_streams") + def test_sqs_event_with_datadog_message_attributes( + self, mock_extract_context_with_data_streams + ): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "String", "stringValue": dd_json_data} + }, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.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( + 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") + def test_sqs_event_with_binary_datadog_message_attributes( + self, mock_extract_context_with_data_streams + ): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + encoded_data = base64.b64encode(dd_json_data.encode()).decode() + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "Binary", "binaryValue": encoded_data} + }, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.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( + 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") + def test_sns_event_with_datadog_message_attributes( + self, mock_extract_context_with_data_streams + ): + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [ + { + "eventSourceARN": "", + "Sns": { + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": {"Type": "String", "Value": dd_json_data} + }, + }, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.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( + dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + ) + self.assertEqual(result, mock_context) + + +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 + ): + dd_data = {"dd-pathway-ctx-base64": "12345"} + kinesis_data = {"_datadog": dd_data, "message": "test"} + kinesis_data_str = json.dumps(kinesis_data) + encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + "kinesis": {"data": encoded_data}, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.return_value = mock_context + + result = extract_context_from_kinesis_event(event, self.lambda_context) + + mock_extract_context_with_data_streams.assert_called_once_with( + dd_data, + "kinesis", + "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + ) + self.assertEqual(result, mock_context) From 6f9d42c24b7c8126335e234aaf97f1ff19177a9b Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 18:20:27 -0400 Subject: [PATCH 15/42] add wanted tests --- tests/test_tracing.py | 70 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 80d0a0db..c79e1524 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2660,6 +2660,76 @@ def test_sns_event_with_datadog_message_attributes( ) self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing._extract_context_with_data_streams") + def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present( + self, mock_extract_context_with_data_streams + ): + """Test that is_sqs = True when eventSourceARN is present in first record""" + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "String", "stringValue": dd_json_data} + }, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.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( + 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") + def test_sns_to_sqs_event_detection_and_processing( + self, mock_extract_context_with_data_streams + ): + """Test SNS->SQS case where SQS body contains SNS notification""" + dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_json_data = json.dumps(dd_data) + + sns_notification = { + "Type": "Notification", + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": {"Type": "String", "Value": dd_json_data} + }, + "Message": "test message", + } + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "body": json.dumps(sns_notification), + "messageAttributes": {}, + } + ] + } + + mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) + mock_extract_context_with_data_streams.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( + dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + ) + self.assertEqual(result, mock_context) + class TestExtractContextFromKinesisEvent(unittest.TestCase): def setUp(self): From 33cdc6c79dd77161be68049f3b809b8c59a2d694 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Fri, 20 Jun 2025 18:20:39 -0400 Subject: [PATCH 16/42] caught sns -> sqs bug --- datadog_lambda/tracing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index f40b2111..62707783 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -275,7 +275,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 {} - arn = sns_record.get("TopicArn", "") + 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: From 3672f39b2265a251f7baea0b61ea4e62b1824554 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 17:31:52 -0400 Subject: [PATCH 17/42] revert back to original tracing.py implementation --- datadog_lambda/tracing.py | 39 +++++++----- tests/test_tracing.py | 122 +++++++++++++++++--------------------- 2 files changed, 80 insertions(+), 81 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 62707783..f30617b7 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -66,20 +66,20 @@ LOWER_64_BITS = "LOWER_64_BITS" -def _extract_context_with_data_streams(context_json, event_type, arn): +def _dsm_set_checkpoint(context_json, event_type, arn): from ddtrace.data_streams import set_consume_checkpoint + from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 """ 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) - if not config.data_streams_enabled: - return context + if PROPAGATION_KEY_BASE_64 not in context_json: + return - if "dd-pathway-ctx-base64" not in context_json: - return context + if not config.data_streams_enabled: + return try: carrier_get = _create_carrier_get(context_json) set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) @@ -87,7 +87,6 @@ def _extract_context_with_data_streams(context_json, event_type, arn): logger.debug( f"DSM:Failed to set consume checkpoint for {event_type} {arn}: {e}" ) - return context def _create_carrier_get(context_json): @@ -232,7 +231,9 @@ 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_and_set_dsm_checkpoint_if_enabled( + event, lambda_context +): """ Extract Datadog trace context from an SQS event. @@ -244,6 +245,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): 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 @@ -308,9 +310,9 @@ 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 - ) + context = propagator.extract(dd_data) + _dsm_set_checkpoint(dd_data, "sqs" if is_sqs else "sns", arn) + return context else: # Handle case where trace context is injected into attributes.AWSTraceHeader # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 @@ -391,9 +393,12 @@ def extract_context_from_eventbridge_event(event, lambda_context): return extract_context_from_lambda_context(lambda_context) -def extract_context_from_kinesis_event(event, lambda_context): +def extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( + event, lambda_context +): """ Extract datadog trace context from a Kinesis Stream's base64 encoded data string + Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported. """ try: record = get_first_record(event) @@ -411,7 +416,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) + _dsm_set_checkpoint(dd_ctx, "kinesis", arn) + return context except Exception as e: logger.debug("The trace extractor returned with error %s", e) @@ -631,13 +638,15 @@ 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 = extract_context_from_sqs_or_sns_event_or_context( + context = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( 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 = extract_context_from_kinesis_event(event, lambda_context) + context = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( + event, lambda_context + ) elif event_source.equals(EventTypes.STEPFUNCTIONS): context = extract_context_from_step_functions(event, lambda_context) else: diff --git a/tests/test_tracing.py b/tests/test_tracing.py index c79e1524..5e9e8828 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -42,10 +42,10 @@ service_mapping as global_service_mapping, propagator, emit_telemetry_on_exception_outside_of_handler, - _extract_context_with_data_streams, + _dsm_set_checkpoint, _create_carrier_get, - extract_context_from_sqs_or_sns_event_or_context, - extract_context_from_kinesis_event, + extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled, + extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled, ) from tests.utils import get_mock_context @@ -2445,12 +2445,8 @@ def test_exception_outside_handler_tracing_disabled( mock_trace.assert_not_called() -class TestExtractContext(unittest.TestCase): +class TestDsmSetCheckpoint(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 +2456,53 @@ 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_dsm_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 - - result = _extract_context_with_data_streams(context_json, event_type, arn) + _dsm_set_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_dsm_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 - - result = _extract_context_with_data_streams(context_json, event_type, arn) + _dsm_set_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_dsm_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 + _dsm_set_checkpoint(context_json, event_type, arn) - result = _extract_context_with_data_streams(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_dsm_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) + _dsm_set_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 +2542,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._dsm_set_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_dsm_set_checkpoint ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -2585,20 +2562,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( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( event, self.lambda_context ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_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._dsm_set_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_dsm_set_checkpoint ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -2616,20 +2595,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( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( event, self.lambda_context ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_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._dsm_set_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_dsm_set_checkpoint ): dd_data = {"dd-pathway-ctx-base64": "12345"} dd_json_data = json.dumps(dd_data) @@ -2649,20 +2630,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( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( event, self.lambda_context ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_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._dsm_set_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_dsm_set_checkpoint ): """Test that is_sqs = True when eventSourceARN is present in first record""" dd_data = {"dd-pathway-ctx-base64": "12345"} @@ -2680,20 +2663,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( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( event, self.lambda_context ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_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._dsm_set_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_dsm_set_checkpoint ): """Test SNS->SQS case where SQS body contains SNS notification""" dd_data = {"dd-pathway-ctx-base64": "12345"} @@ -2719,13 +2704,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( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( event, self.lambda_context ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_set_checkpoint.assert_called_once_with( dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" ) self.assertEqual(result, mock_context) @@ -2735,9 +2721,10 @@ class TestExtractContextFromKinesisEvent(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() - @patch("datadog_lambda.tracing._extract_context_with_data_streams") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("datadog_lambda.tracing.propagator.extract") def test_kinesis_event_with_datadog_data( - self, mock_extract_context_with_data_streams + self, mock_extract, mock_dsm_set_checkpoint ): dd_data = {"dd-pathway-ctx-base64": "12345"} kinesis_data = {"_datadog": dd_data, "message": "test"} @@ -2754,11 +2741,14 @@ 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) + result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( + event, self.lambda_context + ) - mock_extract_context_with_data_streams.assert_called_once_with( + mock_extract.assert_called_once_with(dd_data) + mock_dsm_set_checkpoint.assert_called_once_with( dd_data, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", From 12deafc6c8f2038563dacce15c92d32be2b6b4ec Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 17:43:21 -0400 Subject: [PATCH 18/42] fix lint --- datadog_lambda/tracing.py | 4 ++-- tests/test_tracing.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index f30617b7..5b49854e 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -231,7 +231,7 @@ def create_sns_event(message): } -def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( +def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, lambda_context ): """ @@ -638,7 +638,7 @@ 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 = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + context = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, lambda_context ) elif event_source.equals(EventTypes.EVENTBRIDGE): diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 5e9e8828..d58834a0 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -45,7 +45,7 @@ _dsm_set_checkpoint, _create_carrier_get, extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled, - extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled, + extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled, ) from tests.utils import get_mock_context @@ -2564,7 +2564,7 @@ def test_sqs_event_with_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, self.lambda_context ) @@ -2597,7 +2597,7 @@ def test_sqs_event_with_binary_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, self.lambda_context ) @@ -2632,7 +2632,7 @@ def test_sns_event_with_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, self.lambda_context ) @@ -2665,7 +2665,7 @@ 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.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, self.lambda_context ) @@ -2706,7 +2706,7 @@ def test_sns_to_sqs_event_detection_and_processing( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_checkpoint_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( event, self.lambda_context ) From d8a4379c623b4b9b9d4efc7822c6b2cfb5d78ecb Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 17:45:47 -0400 Subject: [PATCH 19/42] revert spacing stuff to original --- datadog_lambda/tracing.py | 1 + datadog_lambda/wrapper.py | 2 -- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 5b49854e..b33c60de 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -50,6 +50,7 @@ set_tracer_provider(TracerProvider()) + logger = logging.getLogger(__name__) dd_trace_context = None diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index eeb1bdf0..5d699a40 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -10,8 +10,6 @@ from time import time_ns from datadog_lambda.asm import asm_start_response, asm_start_request - - from datadog_lambda.extension import should_use_extension, flush_extension from datadog_lambda.cold_start import ( set_cold_start, From b5af84f487b57c14ab179a8428ba07b9de3d67c7 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 18:50:25 -0400 Subject: [PATCH 20/42] remove unneccessary checks, still set checkpoints even when dsm context fails to be propagated --- datadog_lambda/tracing.py | 12 +- tests/test_tracing.py | 266 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 262 insertions(+), 16 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b33c60de..4acbe0a6 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -69,18 +69,17 @@ def _dsm_set_checkpoint(context_json, event_type, arn): from ddtrace.data_streams import set_consume_checkpoint - from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 """ Extracts the context from a JSON carrier and optionally sets a dsm consume checkpoint if the context is complete and data streams are enabled. """ - - if PROPAGATION_KEY_BASE_64 not in context_json: + if not isinstance(context_json, dict): return if not config.data_streams_enabled: return + try: carrier_get = _create_carrier_get(context_json) set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) @@ -336,9 +335,13 @@ def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled span_id=int(x_ray_context["parent_id"], 16), sampling_priority=float(x_ray_context["sampled"]), ) + # Still want to set a DSM checkpoint even if DSM context not propagated + _dsm_set_checkpoint({}, "sqs" if is_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({}, "sqs" if is_sqs else "sns", arn) return extract_context_from_lambda_context(lambda_context) @@ -422,7 +425,8 @@ def extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( return 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({}, "kinesis", arn) return extract_context_from_lambda_context(lambda_context) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index d58834a0..9eafa091 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2480,16 +2480,6 @@ def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self): self.assertTrue(callable(args[2])) self.assertEqual(kwargs["manual_checkpoint"], False) - @patch("datadog_lambda.config.Config.data_streams_enabled", True) - def test_dsm_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" - - _dsm_set_checkpoint(context_json, event_type, arn) - - self.mock_checkpoint.assert_not_called() - @patch("datadog_lambda.config.Config.data_streams_enabled", True) def test_dsm_set_checkpoint_exception_path(self): context_json = {"dd-pathway-ctx-base64": "12345"} @@ -2504,6 +2494,34 @@ 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_sqs(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_non_dict_context_sns_to_sqs(self, mock_checkpoint): + _dsm_set_checkpoint( + ["not", "a", "dict"], "sqs", "arn:aws:sqs:us-east-1:123456789012:test" + ) + mock_checkpoint.assert_not_called() + + @patch("ddtrace.data_streams.set_consume_checkpoint") + def test_dsm_set_checkpoint_non_dict_context_kinesis(self, mock_checkpoint): + _dsm_set_checkpoint( + 12345, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test" + ) + mock_checkpoint.assert_not_called() + + @patch("ddtrace.data_streams.set_consume_checkpoint") + def test_dsm_set_checkpoint_non_dict_context_sns(self, mock_checkpoint): + _dsm_set_checkpoint( + None, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + ) + mock_checkpoint.assert_not_called() + class TestCreateCarrierGet(unittest.TestCase): def test_create_carrier_get_with_valid_data(self): @@ -2538,7 +2556,7 @@ def test_create_carrier_get_with_empty_context(self): self.assertIsNone(carrier_get("any-key")) -class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase): +class TestExtractContextFromSqsOrSnsEventWithDSMLogic(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() @@ -2716,8 +2734,168 @@ def test_sns_to_sqs_event_detection_and_processing( ) self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sqs_event_without_datadog_message_attributes( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": {}, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sqs_event_with_malformed_datadog_message_attributes( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + dd_json_data = "not-json" + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "messageAttributes": { + "_datadog": {"dataType": "String", "stringValue": dd_json_data} + }, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sns_event_without_datadog_message_attributes( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + event = { + "Records": [ + { + "Sns": { + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": {}, + } + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sns_event_with_malformed_datadog_message_attributes( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + dd_json_data = "not-json" + + event = { + "Records": [ + { + "Sns": { + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": {"Type": "String", "Value": dd_json_data} + }, + } + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sns_to_sqs_event_with_malformed_datadog_message_attributes( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test SNS->SQS case where SQS body contains SNS notification with malformed datadog ctx""" + dd_json_data = "not-json" + + sns_notification = { + "Type": "Notification", + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": { + "_datadog": {"Type": "String", "Value": dd_json_data} + }, + "Message": "test message", + } + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", + "body": json.dumps(sns_notification), + "messageAttributes": {}, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + -class TestExtractContextFromKinesisEvent(unittest.TestCase): +class TestExtractContextFromKinesisEventWithDSMLogic(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() @@ -2754,3 +2932,67 @@ def test_kinesis_event_with_datadog_data( "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", ) self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_kinesis_event_without_datadog_data( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + kinesis_data = {"message": "test"} + kinesis_data_str = json.dumps(kinesis_data) + encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + "kinesis": {"data": encoded_data}, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, + "kinesis", + "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_kinesis_event_with_malformed_data( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + encoded_data = "not-base64-or-json" + + event = { + "Records": [ + { + "eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + "kinesis": {"data": encoded_data}, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with( + {}, + "kinesis", + "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + ) + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) From 3ffd16fdd0cfb624e00a001d52bbcdcbebc6d3a1 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 18:52:26 -0400 Subject: [PATCH 21/42] remove not needed comment --- datadog_lambda/tracing.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4acbe0a6..383a03d8 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -70,10 +70,6 @@ def _dsm_set_checkpoint(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. - """ if not isinstance(context_json, dict): return From c84ea1442c4500a3bc8e3631cde3f721c114954c Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 19:10:15 -0400 Subject: [PATCH 22/42] fixes --- datadog_lambda/tracing.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 383a03d8..3a0e69be 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -68,8 +68,6 @@ def _dsm_set_checkpoint(context_json, event_type, arn): - from ddtrace.data_streams import set_consume_checkpoint - if not isinstance(context_json, dict): return @@ -77,7 +75,9 @@ def _dsm_set_checkpoint(context_json, event_type, arn): return try: - carrier_get = _create_carrier_get(context_json) + from ddtrace.data_streams import set_consume_checkpoint + + carrier_get = lambda k: context_json.get(k) set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) except Exception as e: logger.debug( @@ -85,13 +85,6 @@ def _dsm_set_checkpoint(context_json, event_type, arn): ) -def _create_carrier_get(context_json): - def carrier_get(key): - return context_json.get(key) - - return carrier_get - - def _convert_xray_trace_id(xray_trace_id): """ Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int). From 4ad2bab8dd468f7d2aab69ee08d38fcc69a386ce Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 19:12:24 -0400 Subject: [PATCH 23/42] lambda functions not allowed by lint --- datadog_lambda/tracing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 3a0e69be..1c71b66d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -77,7 +77,7 @@ def _dsm_set_checkpoint(context_json, event_type, arn): try: from ddtrace.data_streams import set_consume_checkpoint - carrier_get = lambda k: context_json.get(k) + carrier_get = _create_carrier_get(context_json) set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) except Exception as e: logger.debug( @@ -85,6 +85,13 @@ def _dsm_set_checkpoint(context_json, event_type, arn): ) +def _create_carrier_get(context_json): + def carrier_get(key): + return context_json.get(key) + + return carrier_get + + def _convert_xray_trace_id(xray_trace_id): """ Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int). From 9996fbbf813cedb96d06d95b12a24eeebf00603a Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 21:27:10 -0400 Subject: [PATCH 24/42] use lambda function, add checks before checkpoint --- datadog_lambda/tracing.py | 16 ++++---- tests/test_tracing.py | 77 +++++++++------------------------------ 2 files changed, 24 insertions(+), 69 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 1c71b66d..e6d27950 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -68,16 +68,21 @@ def _dsm_set_checkpoint(context_json, event_type, arn): + if not config.data_streams_enabled: + return + if not isinstance(context_json, dict): return - if not config.data_streams_enabled: + from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 + + if context_json and PROPAGATION_KEY_BASE_64 not in context_json: return try: from ddtrace.data_streams import set_consume_checkpoint - carrier_get = _create_carrier_get(context_json) + carrier_get = lambda k: context_json.get(k) # noqa: E731 set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) except Exception as e: logger.debug( @@ -85,13 +90,6 @@ def _dsm_set_checkpoint(context_json, event_type, arn): ) -def _create_carrier_get(context_json): - def carrier_get(key): - return context_json.get(key) - - return carrier_get - - def _convert_xray_trace_id(xray_trace_id): """ Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int). diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 9eafa091..f66b3606 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -43,7 +43,6 @@ propagator, emit_telemetry_on_exception_outside_of_handler, _dsm_set_checkpoint, - _create_carrier_get, extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled, extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled, ) @@ -60,6 +59,7 @@ fake_xray_header_value_root_decimal = "3995693151288333088" event_samples = "tests/event_samples/" +PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64" def with_trace_propagation_style(style): @@ -2457,7 +2457,7 @@ def setUp(self): @patch("datadog_lambda.config.Config.data_streams_enabled", False) def test_dsm_set_checkpoint_data_streams_disabled(self): - context_json = {"dd-pathway-ctx-base64": "12345"} + context_json = {PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2467,7 +2467,7 @@ def test_dsm_set_checkpoint_data_streams_disabled(self): @patch("datadog_lambda.config.Config.data_streams_enabled", True) def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self): - context_json = {"dd-pathway-ctx-base64": "12345"} + context_json = {PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2482,7 +2482,7 @@ def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self): @patch("datadog_lambda.config.Config.data_streams_enabled", True) def test_dsm_set_checkpoint_exception_path(self): - context_json = {"dd-pathway-ctx-base64": "12345"} + context_json = {PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2495,67 +2495,24 @@ def test_dsm_set_checkpoint_exception_path(self): self.mock_logger.debug.assert_called_once() @patch("ddtrace.data_streams.set_consume_checkpoint") - def test_dsm_set_checkpoint_non_dict_context_sqs(self, mock_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_non_dict_context_sns_to_sqs(self, mock_checkpoint): - _dsm_set_checkpoint( - ["not", "a", "dict"], "sqs", "arn:aws:sqs:us-east-1:123456789012:test" - ) - mock_checkpoint.assert_not_called() - - @patch("ddtrace.data_streams.set_consume_checkpoint") - def test_dsm_set_checkpoint_non_dict_context_kinesis(self, mock_checkpoint): - _dsm_set_checkpoint( - 12345, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test" - ) - mock_checkpoint.assert_not_called() - - @patch("ddtrace.data_streams.set_consume_checkpoint") - def test_dsm_set_checkpoint_non_dict_context_sns(self, mock_checkpoint): + def test_dsm_set_checkpoint_PROPAGATION_KEY_BASE_64_not_present( + self, mock_checkpoint + ): _dsm_set_checkpoint( - None, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + {"not_a_dict": "not_a_dict"}, + "sqs", + "arn:aws:sqs:us-east-1:123456789012:test-queue", ) mock_checkpoint.assert_not_called() -class TestCreateCarrierGet(unittest.TestCase): - def test_create_carrier_get_with_valid_data(self): - context_json = { - "x-datadog-trace-id": "12345", - "x-datadog-parent-id": "67890", - "x-datadog-sampling-priority": "1", - } - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") - self.assertEqual(carrier_get("x-datadog-parent-id"), "67890") - self.assertEqual(carrier_get("x-datadog-sampling-priority"), "1") - - def test_create_carrier_get_with_missing_key(self): - context_json = {"x-datadog-trace-id": "12345"} - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertEqual(carrier_get("x-datadog-trace-id"), "12345") - self.assertIsNone(carrier_get("x-datadog-parent-id")) - - def test_create_carrier_get_with_empty_context(self): - context_json = {} - - carrier_get = _create_carrier_get(context_json) - - self.assertTrue(callable(carrier_get)) - self.assertIsNone(carrier_get("any-key")) - - class TestExtractContextFromSqsOrSnsEventWithDSMLogic(unittest.TestCase): def setUp(self): self.lambda_context = get_mock_context() @@ -2565,7 +2522,7 @@ def setUp(self): def test_sqs_event_with_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2597,7 +2554,7 @@ def test_sqs_event_with_datadog_message_attributes( def test_sqs_event_with_binary_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) encoded_data = base64.b64encode(dd_json_data.encode()).decode() @@ -2630,7 +2587,7 @@ def test_sqs_event_with_binary_datadog_message_attributes( def test_sns_event_with_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2666,7 +2623,7 @@ def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present( self, mock_extract, mock_dsm_set_checkpoint ): """Test that is_sqs = True when eventSourceARN is present in first record""" - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2699,7 +2656,7 @@ def test_sns_to_sqs_event_detection_and_processing( self, mock_extract, mock_dsm_set_checkpoint ): """Test SNS->SQS case where SQS body contains SNS notification""" - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) sns_notification = { @@ -2904,7 +2861,7 @@ def setUp(self): def test_kinesis_event_with_datadog_data( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {"dd-pathway-ctx-base64": "12345"} + dd_data = {PROPAGATION_KEY_BASE_64: "12345"} kinesis_data = {"_datadog": dd_data, "message": "test"} kinesis_data_str = json.dumps(kinesis_data) encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() From 8b0313b6e2100450fecfc6257b06a286f94fc1fb Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 21:45:07 -0400 Subject: [PATCH 25/42] remove unneccesary import --- datadog_lambda/tracing.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index e6d27950..9f95f732 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -65,20 +65,19 @@ DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000" HIGHER_64_BITS = "HIGHER_64_BITS" LOWER_64_BITS = "LOWER_64_BITS" +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 not isinstance(context_json, dict): return - from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 - if context_json and PROPAGATION_KEY_BASE_64 not in context_json: return + if not config.data_streams_enabled: + return + try: from ddtrace.data_streams import set_consume_checkpoint From cfbb1f8a0fc071c0e4da590212b1d32c771ae70d Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Mon, 30 Jun 2025 21:48:06 -0400 Subject: [PATCH 26/42] move if statement with least work first --- datadog_lambda/tracing.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 9f95f732..214cc77f 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -69,13 +69,13 @@ def _dsm_set_checkpoint(context_json, event_type, arn): - if not isinstance(context_json, dict): + if not config.data_streams_enabled: return - if context_json and PROPAGATION_KEY_BASE_64 not in context_json: + if not isinstance(context_json, dict): return - if not config.data_streams_enabled: + if context_json and PROPAGATION_KEY_BASE_64 not in context_json: return try: From 0b01978e59720d81556b2e0b85421bf5ebf4bb8f Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Tue, 1 Jul 2025 19:57:31 -0400 Subject: [PATCH 27/42] changed function name to original, arn exception handle w test, return None instead of {} --- datadog_lambda/tracing.py | 37 ++++++------- tests/test_tracing.py | 110 ++++++++++++++++++++++++-------------- 2 files changed, 88 insertions(+), 59 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 214cc77f..6b655544 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -65,23 +65,24 @@ DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000" HIGHER_64_BITS = "HIGHER_64_BITS" LOWER_64_BITS = "LOWER_64_BITS" -PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64" +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 not isinstance(context_json, dict): - return - - if context_json and PROPAGATION_KEY_BASE_64 not in context_json: - 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 - carrier_get = lambda k: context_json.get(k) # noqa: E731 + carrier_get = lambda k: context_json and context_json.get(k) # noqa: E731 set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False) except Exception as e: logger.debug( @@ -224,9 +225,7 @@ def create_sns_event(message): } -def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( - event, lambda_context -): +def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): """ Extract Datadog trace context from an SQS event. @@ -241,6 +240,7 @@ def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported. """ is_sqs = False + arn = "" # EventBridge => SQS try: @@ -329,12 +329,12 @@ def extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled sampling_priority=float(x_ray_context["sampled"]), ) # Still want to set a DSM checkpoint even if DSM context not propagated - _dsm_set_checkpoint({}, "sqs" if is_sqs else "sns", arn) + _dsm_set_checkpoint(None, "sqs" if is_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({}, "sqs" if is_sqs else "sns", arn) + _dsm_set_checkpoint(None, "sqs" if is_sqs else "sns", arn) return extract_context_from_lambda_context(lambda_context) @@ -390,13 +390,12 @@ def extract_context_from_eventbridge_event(event, lambda_context): return extract_context_from_lambda_context(lambda_context) -def extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( - event, lambda_context -): +def extract_context_from_kinesis_event(event, lambda_context): """ Extract datadog trace context from a Kinesis Stream's base64 encoded data string Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported. """ + arn = "" try: record = get_first_record(event) arn = record.get("eventSourceARN", "") @@ -419,7 +418,7 @@ def extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( 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({}, "kinesis", arn) + _dsm_set_checkpoint(None, "kinesis", arn) return extract_context_from_lambda_context(lambda_context) @@ -636,15 +635,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 = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + context = 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 = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( - event, lambda_context - ) + context = extract_context_from_kinesis_event(event, lambda_context) elif event_source.equals(EventTypes.STEPFUNCTIONS): context = extract_context_from_step_functions(event, lambda_context) else: diff --git a/tests/test_tracing.py b/tests/test_tracing.py index f66b3606..f4e1f0d7 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -43,8 +43,8 @@ propagator, emit_telemetry_on_exception_outside_of_handler, _dsm_set_checkpoint, - extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled, - extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled, + extract_context_from_kinesis_event, + extract_context_from_sqs_or_sns_event_or_context, ) from tests.utils import get_mock_context @@ -59,7 +59,7 @@ fake_xray_header_value_root_decimal = "3995693151288333088" event_samples = "tests/event_samples/" -PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64" +DSM_PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64" def with_trace_propagation_style(style): @@ -2457,7 +2457,7 @@ def setUp(self): @patch("datadog_lambda.config.Config.data_streams_enabled", False) def test_dsm_set_checkpoint_data_streams_disabled(self): - context_json = {PROPAGATION_KEY_BASE_64: "12345"} + context_json = {DSM_PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2467,7 +2467,7 @@ def test_dsm_set_checkpoint_data_streams_disabled(self): @patch("datadog_lambda.config.Config.data_streams_enabled", True) def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self): - context_json = {PROPAGATION_KEY_BASE_64: "12345"} + context_json = {DSM_PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2482,7 +2482,7 @@ def test_dsm_set_checkpoint_data_streams_enabled_complete_context(self): @patch("datadog_lambda.config.Config.data_streams_enabled", True) def test_dsm_set_checkpoint_exception_path(self): - context_json = {PROPAGATION_KEY_BASE_64: "12345"} + context_json = {DSM_PROPAGATION_KEY_BASE_64: "12345"} event_type = "sqs" arn = "arn:aws:sqs:us-east-1:123456789012:test-queue" @@ -2502,7 +2502,7 @@ def test_dsm_set_checkpoint_non_dict_context(self, mock_checkpoint): mock_checkpoint.assert_not_called() @patch("ddtrace.data_streams.set_consume_checkpoint") - def test_dsm_set_checkpoint_PROPAGATION_KEY_BASE_64_not_present( + def test_dsm_set_checkpoint_DSM_PROPAGATION_KEY_BASE_64_not_present( self, mock_checkpoint ): _dsm_set_checkpoint( @@ -2522,7 +2522,7 @@ def setUp(self): def test_sqs_event_with_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2539,7 +2539,7 @@ def test_sqs_event_with_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) @@ -2554,7 +2554,7 @@ def test_sqs_event_with_datadog_message_attributes( def test_sqs_event_with_binary_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) encoded_data = base64.b64encode(dd_json_data.encode()).decode() @@ -2572,7 +2572,7 @@ def test_sqs_event_with_binary_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) @@ -2587,7 +2587,7 @@ def test_sqs_event_with_binary_datadog_message_attributes( def test_sns_event_with_datadog_message_attributes( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2607,7 +2607,7 @@ def test_sns_event_with_datadog_message_attributes( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) @@ -2623,7 +2623,7 @@ def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present( self, mock_extract, mock_dsm_set_checkpoint ): """Test that is_sqs = True when eventSourceARN is present in first record""" - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) event = { @@ -2640,7 +2640,7 @@ 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.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) @@ -2656,7 +2656,7 @@ def test_sns_to_sqs_event_detection_and_processing( self, mock_extract, mock_dsm_set_checkpoint ): """Test SNS->SQS case where SQS body contains SNS notification""" - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) sns_notification = { @@ -2681,7 +2681,7 @@ def test_sns_to_sqs_event_detection_and_processing( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) @@ -2708,12 +2708,12 @@ def test_sqs_event_without_datadog_message_attributes( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) mock_dsm_set_checkpoint.assert_called_once_with( - {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + None, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @@ -2739,12 +2739,12 @@ def test_sqs_event_with_malformed_datadog_message_attributes( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) mock_dsm_set_checkpoint.assert_called_once_with( - {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + None, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @@ -2768,12 +2768,12 @@ def test_sns_event_without_datadog_message_attributes( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) mock_dsm_set_checkpoint.assert_called_once_with( - {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + None, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @@ -2801,12 +2801,12 @@ def test_sns_event_with_malformed_datadog_message_attributes( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) mock_dsm_set_checkpoint.assert_called_once_with( - {}, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" + None, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic" ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @@ -2841,16 +2841,35 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_sqs_or_sns_event_or_context_and_set_dsm_ckpt_if_enabled( + result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context ) mock_dsm_set_checkpoint.assert_called_once_with( - {}, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" + None, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue" ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sqs_sns_event_with_exception_accessing_first_record( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test that arn is properly initialized to empty string when exception occur""" + event = {"Records": None} + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context + ) + + mock_dsm_set_checkpoint.assert_called_once_with(None, "sns", "") + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + class TestExtractContextFromKinesisEventWithDSMLogic(unittest.TestCase): def setUp(self): @@ -2861,7 +2880,7 @@ def setUp(self): def test_kinesis_event_with_datadog_data( self, mock_extract, mock_dsm_set_checkpoint ): - dd_data = {PROPAGATION_KEY_BASE_64: "12345"} + dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} kinesis_data = {"_datadog": dd_data, "message": "test"} kinesis_data_str = json.dumps(kinesis_data) encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() @@ -2878,9 +2897,7 @@ def test_kinesis_event_with_datadog_data( mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) mock_extract.return_value = mock_context - result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( - event, self.lambda_context - ) + result = extract_context_from_kinesis_event(event, self.lambda_context) mock_extract.assert_called_once_with(dd_data) mock_dsm_set_checkpoint.assert_called_once_with( @@ -2911,12 +2928,10 @@ def test_kinesis_event_without_datadog_data( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( - event, self.lambda_context - ) + result = extract_context_from_kinesis_event(event, self.lambda_context) mock_dsm_set_checkpoint.assert_called_once_with( - {}, + None, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", ) @@ -2942,14 +2957,31 @@ def test_kinesis_event_with_malformed_data( mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context - result = extract_context_from_kinesis_event_and_set_dsm_checkpoint_if_enabled( - event, self.lambda_context - ) + result = extract_context_from_kinesis_event(event, self.lambda_context) mock_dsm_set_checkpoint.assert_called_once_with( - {}, + None, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", ) mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_kinesis_event_with_exception_accessing_first_record( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test that arn is properly initialized to empty string when exception occur""" + event = {"Records": None} + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_kinesis_event(event, self.lambda_context) + + # 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_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) From 9dca396387dde70ed71ff20a80c0bb4923d47d81 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 11:56:07 -0400 Subject: [PATCH 28/42] some fixes --- datadog_lambda/tracing.py | 43 ++++++++++++++++++++++++--------------- tests/test_tracing.py | 25 +++++++++++++---------- 2 files changed, 41 insertions(+), 27 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 6b655544..8ecb2312 100644 --- a/datadog_lambda/tracing.py +++ b/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) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index f4e1f0d7..4ebb4ea8 100644 --- a/tests/test_tracing.py +++ b/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) From d849b75ac3b0907bd54b0a272179ae9d75a7556d Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 12:49:18 -0400 Subject: [PATCH 29/42] remove comments that are not needed --- tests/test_tracing.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 4ebb4ea8..28c62fbe 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2859,7 +2859,6 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes( def test_sqs_sns_event_with_exception_accessing_first_record( self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context ): - """Test that arn is properly initialized to empty string when exception occur""" event = {"Records": None} mock_context = Context(trace_id=123, span_id=456) @@ -2975,7 +2974,6 @@ def test_kinesis_event_with_malformed_data( def test_kinesis_event_with_exception_accessing_first_record( self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context ): - """Test that arn is properly initialized to empty string when exception occur""" event = {"Records": None} mock_context = Context(trace_id=123, span_id=456) @@ -2983,8 +2981,6 @@ def test_kinesis_event_with_exception_accessing_first_record( result = extract_context_from_kinesis_event(event, self.lambda_context) - # Verify that _dsm_set_checkpoint is called with empty string for arn - # even when exception occurs mock_dsm_set_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) From a375d81452f1d439324a16c47e08b62efa765090 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 12:53:30 -0400 Subject: [PATCH 30/42] fix --- tests/test_tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 28c62fbe..5c64ec7c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2495,11 +2495,11 @@ def test_dsm_set_checkpoint_exception_path(self): self.mock_logger.debug.assert_called_once() @patch("ddtrace.data_streams.set_consume_checkpoint") - def test_dsm_set_checkpoint_DSM_PROPAGATION_KEY_BASE_64_not_present( + def test_dsm_set_checkpoint_with_non_dict_context_does_not_set_checkpoint( self, mock_checkpoint ): _dsm_set_checkpoint( - {"not_a_dict": "not_a_dict"}, + [], "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue", ) From f465407bba3a0bcb80c837417bc1003512ada5bd Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 12:55:48 -0400 Subject: [PATCH 31/42] fix --- tests/test_tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 5c64ec7c..a134686f 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2618,7 +2618,7 @@ def test_sns_event_with_datadog_message_attributes( def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present( self, mock_extract, mock_dsm_set_checkpoint ): - """Test that is_sqs = True when eventSourceARN is present in first record""" + """Test that is_sqs = True when eventSource is SQS""" dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} dd_json_data = json.dumps(dd_data) From ec5dfe5316450d6584dade184860938d1ed7d6b8 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 13:09:37 -0400 Subject: [PATCH 32/42] extra check --- datadog_lambda/tracing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 8ecb2312..8279da54 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -422,7 +422,9 @@ def extract_context_from_kinesis_event(event, lambda_context): dd_ctx = data_obj.get("_datadog") if dd_ctx: context = propagator.extract(dd_ctx) - _dsm_set_checkpoint(dd_ctx, "kinesis", arn) + # Do not want to set checkpoint with "" arn + if arn: + _dsm_set_checkpoint(dd_ctx, "kinesis", arn) return context except Exception as e: logger.debug("The trace extractor returned with error %s", e) From 123e5a3843775ac18a1d974b0b72e3553bd863d8 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 20:06:58 -0400 Subject: [PATCH 33/42] remove unneccesary work associated with event_source --- datadog_lambda/tracing.py | 5 ++--- tests/test_tracing.py | 23 ++++++++++++----------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 8279da54..5a4fb0e4 100644 --- a/datadog_lambda/tracing.py +++ b/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) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index a134686f..87e1580d 100644 --- a/tests/test_tracing.py +++ b/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() From e0497d074195e348ef7b52262400ab63e8f56354 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Wed, 2 Jul 2025 20:09:23 -0400 Subject: [PATCH 34/42] fix lint --- datadog_lambda/tracing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 5a4fb0e4..cd40e307 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -217,7 +217,9 @@ def create_sns_event(message): } -def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context, event_source): +def extract_context_from_sqs_or_sns_event_or_context( + event, lambda_context, event_source +): """ Extract Datadog trace context from an SQS event. From 81237b3739a594d6fe9d067b3aa513aa34cf7386 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 07:27:50 -0400 Subject: [PATCH 35/42] add tests for empty arn logic --- tests/test_tracing.py | 119 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 87e1580d..dd08f185 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2873,6 +2873,97 @@ def test_sqs_sns_event_with_exception_accessing_first_record( mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sqs_event_with_empty_arn( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test SQS event with empty eventSourceARN""" + event = { + "Records": [ + { + "eventSourceARN": "", + "messageAttributes": {}, + "eventSource": "aws:sqs", + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + mock_dsm_set_checkpoint.assert_not_called() + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sns_event_with_empty_arn( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test SNS event with empty TopicArn""" + event = { + "Records": [ + { + "Sns": { + "TopicArn": "", + "MessageAttributes": {}, + }, + "eventSource": "aws:sns", + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + mock_dsm_set_checkpoint.assert_not_called() + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_sns_to_sqs_event_with_empty_arn( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test SNS->SQS event with empty eventSourceARN""" + sns_notification = { + "Type": "Notification", + "TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic", + "MessageAttributes": {}, + "Message": "test message", + } + + event = { + "Records": [ + { + "eventSourceARN": "", + "body": json.dumps(sns_notification), + "messageAttributes": {}, + "eventSource": "aws:sqs", + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_sqs_or_sns_event_or_context( + event, self.lambda_context, parse_event_source(event) + ) + + mock_dsm_set_checkpoint.assert_not_called() + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + class TestExtractContextFromKinesisEventWithDSMLogic(unittest.TestCase): def setUp(self): @@ -2970,6 +3061,34 @@ def test_kinesis_event_with_malformed_data( mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") + @patch("datadog_lambda.tracing._dsm_set_checkpoint") + def test_kinesis_event_with_empty_arn( + self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + ): + """Test Kinesis event with empty eventSourceARN""" + kinesis_data = {"message": "test"} + kinesis_data_str = json.dumps(kinesis_data) + encoded_data = base64.b64encode(kinesis_data_str.encode()).decode() + + event = { + "Records": [ + { + "eventSourceARN": "", + "kinesis": {"data": encoded_data}, + } + ] + } + + mock_context = Context(trace_id=123, span_id=456) + mock_extract_from_lambda_context.return_value = mock_context + + result = extract_context_from_kinesis_event(event, self.lambda_context) + + mock_dsm_set_checkpoint.assert_not_called() + mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + self.assertEqual(result, mock_context) + @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") def test_kinesis_event_with_exception_accessing_first_record( From 046386290955bb138173b63567d6e7e640430985 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 08:29:26 -0400 Subject: [PATCH 36/42] more descriptive name --- datadog_lambda/tracing.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index cd40e307..44305f42 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -233,7 +233,7 @@ def extract_context_from_sqs_or_sns_event_or_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. """ - arn = "" + source_arn = "" # EventBridge => SQS try: @@ -245,7 +245,7 @@ def extract_context_from_sqs_or_sns_event_or_context( try: first_record = event.get("Records")[0] - arn = first_record.get("eventSourceARN", "") + source_arn = first_record.get("eventSourceARN", "") # logic to deal with SNS => SQS event if "body" in first_record: @@ -297,12 +297,12 @@ def extract_context_from_sqs_or_sns_event_or_context( ) context = propagator.extract(dd_data) # Do not want to set checkpoint with "" arn - if arn: + if source_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, + source_arn, ) return context else: @@ -330,20 +330,20 @@ def extract_context_from_sqs_or_sns_event_or_context( # Still want to set a DSM checkpoint even if DSM context not propagated # 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: + if source_arn: _dsm_set_checkpoint( - None, "sqs" if event_source.equals(EventTypes.SQS) else "sns", arn + None, "sqs" if event_source.equals(EventTypes.SQS) else "sns", source_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 # Do not want to set checkpoint with "" arn - if arn: + if source_arn: _dsm_set_checkpoint( None, "sqs" if event_source.equals(EventTypes.SQS) else "sns", - arn, + source_arn, ) return extract_context_from_lambda_context(lambda_context) @@ -405,10 +405,10 @@ def extract_context_from_kinesis_event(event, lambda_context): Extract datadog trace context from a Kinesis Stream's base64 encoded data string Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported. """ - arn = "" + source_arn = "" try: record = get_first_record(event) - arn = record.get("eventSourceARN", "") + source_arn = record.get("eventSourceARN", "") kinesis = record.get("kinesis") if not kinesis: return extract_context_from_lambda_context(lambda_context) @@ -424,15 +424,15 @@ def extract_context_from_kinesis_event(event, lambda_context): if dd_ctx: context = propagator.extract(dd_ctx) # Do not want to set checkpoint with "" arn - if arn: - _dsm_set_checkpoint(dd_ctx, "kinesis", arn) + if source_arn: + _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn) return 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 # Do not want to set checkpoint with "" arn - if arn: - _dsm_set_checkpoint(None, "kinesis", arn) + if source_arn: + _dsm_set_checkpoint(None, "kinesis", source_arn) return extract_context_from_lambda_context(lambda_context) From 1a9df0a49ee13fa208ad406ff03513a137150c0b Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 08:30:02 -0400 Subject: [PATCH 37/42] fix lint --- datadog_lambda/tracing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 44305f42..aff43577 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -332,7 +332,9 @@ def extract_context_from_sqs_or_sns_event_or_context( # Do not want to set checkpoint with "" arn if source_arn: _dsm_set_checkpoint( - None, "sqs" if event_source.equals(EventTypes.SQS) else "sns", source_arn + None, + "sqs" if event_source.equals(EventTypes.SQS) else "sns", + source_arn, ) return extract_context_from_lambda_context(lambda_context) except Exception as e: From b8ff3fb197a7e94deab34d8a054030a80179d701 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 08:32:17 -0400 Subject: [PATCH 38/42] fix --- datadog_lambda/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index aff43577..0cf6f90f 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -263,7 +263,7 @@ def extract_context_from_sqs_or_sns_event_or_context( sns_record = first_record.get("Sns") or {} # SNS->SQS event would extract SNS arn without this check if event_source.equals(EventTypes.SNS): - arn = sns_record.get("TopicArn", "") + source_arn = sns_record.get("TopicArn", "") msg_attributes = sns_record.get("MessageAttributes") or {} dd_payload = msg_attributes.get("_datadog") if dd_payload: From 021c8f8d6d9585d4bcacd8628a26e52f98c9a986 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 12:53:52 -0400 Subject: [PATCH 39/42] moved arn check to checkpoint, remove comments, add variable dec --- datadog_lambda/tracing.py | 46 +++++++++++++++++---------------------- 1 file changed, 20 insertions(+), 26 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0cf6f90f..6426660c 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -71,6 +71,9 @@ def _dsm_set_checkpoint(context_json, event_type, arn): if not config.data_streams_enabled: return + if not arn: + return + try: from ddtrace.data_streams import set_consume_checkpoint @@ -234,6 +237,7 @@ def extract_context_from_sqs_or_sns_event_or_context( Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported. """ source_arn = "" + event_type = "sqs" if event_source.equals(EventTypes.SQS) else "sns" # EventBridge => SQS try: @@ -296,14 +300,11 @@ def extract_context_from_sqs_or_sns_event_or_context( "Failed to extract Step Functions context from SQS/SNS event." ) context = propagator.extract(dd_data) - # Do not want to set checkpoint with "" arn - if source_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", - source_arn, - ) + _dsm_set_checkpoint( + dd_data, + event_type, + source_arn, + ) return context else: # Handle case where trace context is injected into attributes.AWSTraceHeader @@ -328,25 +329,20 @@ def extract_context_from_sqs_or_sns_event_or_context( sampling_priority=float(x_ray_context["sampled"]), ) # Still want to set a DSM checkpoint even if DSM context not propagated - # In this function only recieves SQS and SNS events, if not SQS must be SNS - # Do not want to set checkpoint with "" arn - if source_arn: - _dsm_set_checkpoint( - None, - "sqs" if event_source.equals(EventTypes.SQS) else "sns", - source_arn, - ) + _dsm_set_checkpoint( + None, + event_type, + source_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 - # Do not want to set checkpoint with "" arn - if source_arn: - _dsm_set_checkpoint( - None, - "sqs" if event_source.equals(EventTypes.SQS) else "sns", - source_arn, - ) + _dsm_set_checkpoint( + None, + event_type, + source_arn, + ) return extract_context_from_lambda_context(lambda_context) @@ -425,9 +421,7 @@ def extract_context_from_kinesis_event(event, lambda_context): dd_ctx = data_obj.get("_datadog") if dd_ctx: context = propagator.extract(dd_ctx) - # Do not want to set checkpoint with "" arn - if source_arn: - _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn) + _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn) return context except Exception as e: logger.debug("The trace extractor returned with error %s", e) From ef98427ef30049c7e379e8729ba5a566ad473761 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 13:14:30 -0400 Subject: [PATCH 40/42] kinesis fix, tests fix --- datadog_lambda/tracing.py | 4 +-- tests/test_tracing.py | 55 +++++++++++++++++++++++++++------------ 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 6426660c..1c448b4b 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -426,9 +426,7 @@ 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 - # Do not want to set checkpoint with "" arn - if source_arn: - _dsm_set_checkpoint(None, "kinesis", source_arn) + _dsm_set_checkpoint(None, "kinesis", source_arn) return extract_context_from_lambda_context(lambda_context) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index dd08f185..9c87e75d 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2857,10 +2857,14 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes( @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_sqs_sns_event_with_exception_accessing_first_record( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, + mock_set_consume_checkpoint, + mock_dsm_set_checkpoint, + mock_extract_from_lambda_context, ): - event = {"Records": None} + event = {"Records": None, "eventSource": "aws:sqs"} mock_context = Context(trace_id=123, span_id=456) mock_extract_from_lambda_context.return_value = mock_context @@ -2868,15 +2872,18 @@ def test_sqs_sns_event_with_exception_accessing_first_record( result = extract_context_from_sqs_or_sns_event_or_context( event, self.lambda_context, parse_event_source(event) ) - - mock_dsm_set_checkpoint.assert_not_called() + mock_set_consume_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_sqs_event_with_empty_arn( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, + mock_set_consume_checkpoint, + mock_dsm_set_checkpoint, + mock_extract_from_lambda_context, ): """Test SQS event with empty eventSourceARN""" event = { @@ -2896,14 +2903,19 @@ def test_sqs_event_with_empty_arn( event, self.lambda_context, parse_event_source(event) ) - mock_dsm_set_checkpoint.assert_not_called() + mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "") + mock_set_consume_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_sns_event_with_empty_arn( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, + mock_set_consume_checkpoint, + mock_dsm_set_checkpoint, + mock_extract_from_lambda_context, ): """Test SNS event with empty TopicArn""" event = { @@ -2925,14 +2937,19 @@ def test_sns_event_with_empty_arn( event, self.lambda_context, parse_event_source(event) ) - mock_dsm_set_checkpoint.assert_not_called() + mock_dsm_set_checkpoint.assert_called_once_with(None, "sns", "") + mock_set_consume_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_sns_to_sqs_event_with_empty_arn( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, + mock_set_consume_checkpoint, + mock_dsm_set_checkpoint, + mock_extract_from_lambda_context, ): """Test SNS->SQS event with empty eventSourceARN""" sns_notification = { @@ -2960,7 +2977,8 @@ def test_sns_to_sqs_event_with_empty_arn( event, self.lambda_context, parse_event_source(event) ) - mock_dsm_set_checkpoint.assert_not_called() + mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "") + mock_set_consume_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @@ -3063,8 +3081,12 @@ def test_kinesis_event_with_malformed_data( @patch("datadog_lambda.tracing.extract_context_from_lambda_context") @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_kinesis_event_with_empty_arn( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, + mock_set_consume_checkpoint, + mock_dsm_set_checkpoint, + mock_extract_from_lambda_context, ): """Test Kinesis event with empty eventSourceARN""" kinesis_data = {"message": "test"} @@ -3085,14 +3107,15 @@ def test_kinesis_event_with_empty_arn( result = extract_context_from_kinesis_event(event, self.lambda_context) - mock_dsm_set_checkpoint.assert_not_called() + mock_dsm_set_checkpoint.assert_called_once_with(None, "kinesis", "") + mock_set_consume_checkpoint.assert_not_called() mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) self.assertEqual(result, mock_context) @patch("datadog_lambda.tracing.extract_context_from_lambda_context") - @patch("datadog_lambda.tracing._dsm_set_checkpoint") + @patch("ddtrace.data_streams.set_consume_checkpoint") def test_kinesis_event_with_exception_accessing_first_record( - self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context + self, mock_set_consume_checkpoint, mock_extract_from_lambda_context ): event = {"Records": None} @@ -3100,7 +3123,5 @@ def test_kinesis_event_with_exception_accessing_first_record( mock_extract_from_lambda_context.return_value = mock_context result = extract_context_from_kinesis_event(event, self.lambda_context) - - mock_dsm_set_checkpoint.assert_not_called() - mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context) + mock_set_consume_checkpoint.assert_not_called() self.assertEqual(result, mock_context) From 0f71f842990c000671674ffbaae3771fd16ea47d Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 13:53:36 -0400 Subject: [PATCH 41/42] remove not needed test --- tests/test_tracing.py | 34 ---------------------------------- 1 file changed, 34 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 9c87e75d..9f4547f0 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -2614,40 +2614,6 @@ def test_sns_event_with_datadog_message_attributes( ) self.assertEqual(result, mock_context) - @patch("datadog_lambda.tracing._dsm_set_checkpoint") - @patch("datadog_lambda.tracing.propagator.extract") - def test_sqs_event_determines_is_sqs_true_when_event_source_arn_present( - self, mock_extract, mock_dsm_set_checkpoint - ): - """Test that is_sqs = True when eventSource is SQS""" - dd_data = {DSM_PROPAGATION_KEY_BASE_64: "12345"} - dd_json_data = json.dumps(dd_data) - - event = { - "Records": [ - { - "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue", - "messageAttributes": { - "_datadog": {"dataType": "String", "stringValue": dd_json_data} - }, - "eventSource": "aws:sqs", - } - ] - } - - mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1) - mock_extract.return_value = mock_context - - result = extract_context_from_sqs_or_sns_event_or_context( - event, self.lambda_context, parse_event_source(event) - ) - - mock_extract.assert_called_once_with(dd_data) - mock_dsm_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._dsm_set_checkpoint") @patch("datadog_lambda.tracing.propagator.extract") def test_sns_to_sqs_event_detection_and_processing( From c4fa49bcd34cd0b54c4c5054fd61f63a946d1c59 Mon Sep 17 00:00:00 2001 From: Michael Zhao Date: Thu, 3 Jul 2025 14:03:16 -0400 Subject: [PATCH 42/42] formatting fix --- datadog_lambda/tracing.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 1c448b4b..f4057480 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -300,11 +300,7 @@ def extract_context_from_sqs_or_sns_event_or_context( "Failed to extract Step Functions context from SQS/SNS event." ) context = propagator.extract(dd_data) - _dsm_set_checkpoint( - dd_data, - event_type, - source_arn, - ) + _dsm_set_checkpoint(dd_data, event_type, source_arn) return context else: # Handle case where trace context is injected into attributes.AWSTraceHeader @@ -329,20 +325,12 @@ def extract_context_from_sqs_or_sns_event_or_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, - event_type, - source_arn, - ) + _dsm_set_checkpoint(None, event_type, source_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, - event_type, - source_arn, - ) + _dsm_set_checkpoint(None, event_type, source_arn) return extract_context_from_lambda_context(lambda_context)