From a38c77baa4f5ffa622efca28a83f32605f4f51f2 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 11 Mar 2025 12:57:15 -0400 Subject: [PATCH 01/15] support step function context in eventbridge extraction --- datadog_lambda/tracing.py | 43 ++++++++++---- datadog_lambda/trigger.py | 6 +- datadog_lambda/wrapper.py | 3 - tests/test_tracing.py | 114 +++++++++++++++++++++++++++++++++++++- 4 files changed, 149 insertions(+), 17 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index a73423e1..03f10d19 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -320,12 +320,17 @@ def extract_context_from_eventbridge_event(event, lambda_context): """ Extract datadog trace context from an EventBridge message's Details. This is only possible if Details is a JSON string. + + If we find a Step Function context, try to extract the trace context from + that header. """ try: detail = event.get("detail") dd_context = detail.get("_datadog") if not dd_context: return extract_context_from_lambda_context(lambda_context) + if is_step_function_event(dd_context): + return extract_context_from_step_functions(detail, lambda_context) return propagator.extract(dd_context) except Exception as e: logger.debug("The trace extractor returned with error %s", e) @@ -424,7 +429,7 @@ def _generate_sfn_trace_id(execution_id: str, part: str): def extract_context_from_step_functions(event, lambda_context): """ Only extract datadog trace context when Step Functions Context Object is injected - into lambda's event dict. + into lambda's event dict. Unwrap "Payload" if it exists to handle Legacy Lambda cases. If '_datadog' header is present, we have two cases: 1. Root is a Lambda and we use its traceID @@ -435,6 +440,8 @@ def extract_context_from_step_functions(event, lambda_context): object. """ try: + event = event.get("Payload", event) + meta = {} dd_data = event.get("_datadog") @@ -472,18 +479,30 @@ def extract_context_from_step_functions(event, lambda_context): return extract_context_from_lambda_context(lambda_context) -def is_legacy_lambda_step_function(event): +def is_step_function_event(event): """ - Check if the event is a step function that called a legacy lambda + Check if the event is a step function that invoked the current lambda. + + The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a + "_datadog" for JSONata style context propagation. + + The actual event must contain "Execution", "StateMachine", and "State" fields. """ - if not isinstance(event, dict) or "Payload" not in event: - return False + event = event.get("Payload", event) + + # JSONPath style + if all(field in event for field in ("Execution", "StateMachine", "State")): + return True + + # JSONata style + if "_datadog" in event: + event = event["_datadog"] + return all( + field in event + for field in ("Execution", "StateMachine", "State", "serverless-version") + ) - event = event.get("Payload") - return isinstance(event, dict) and ( - "_datadog" in event - or ("Execution" in event and "StateMachine" in event and "State" in event) - ) + return False def extract_context_custom_extractor(extractor, event, lambda_context): @@ -1320,6 +1339,10 @@ def create_inferred_span_from_eventbridge_event(event, context): if span: span.set_tags(tags) span.start = dt.replace(tzinfo=timezone.utc).timestamp() + + # Since inferred span will later parent Lambda, preserve Lambda's current parent + span.parent_id = dd_trace_context.span_id + return span diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index 11759a0a..c583763a 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -10,6 +10,8 @@ from enum import Enum from typing import Any +from datadog_lambda.tracing import is_step_function_event + class _stringTypedEnum(Enum): """ @@ -146,9 +148,7 @@ def parse_event_source(event: dict) -> _EventSource: if event.get("source") == "aws.events" or has_event_categories: event_source = _EventSource(EventTypes.CLOUDWATCH_EVENTS) - if ( - "_datadog" in event and event.get("_datadog").get("serverless-version") == "v1" - ) or ("Execution" in event and "StateMachine" in event and "State" in event): + if is_step_function_event(event): event_source = _EventSource(EventTypes.STEPFUNCTIONS) event_record = get_first_record(event) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 2632d22e..b52a145a 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -45,7 +45,6 @@ is_authorizer_response, tracer, propagator, - is_legacy_lambda_step_function, ) from datadog_lambda.trigger import ( extract_trigger_tags, @@ -279,8 +278,6 @@ def _before(self, event, context): self.response = None set_cold_start(init_timestamp_ns) submit_invocations_metric(context) - if is_legacy_lambda_step_function(event): - event = event["Payload"] self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 5480a92c..1b4ad0d7 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -42,7 +42,7 @@ service_mapping as global_service_mapping, propagator, emit_telemetry_on_exception_outside_of_handler, - is_legacy_lambda_step_function, + is_step_function_event, ) from datadog_lambda.trigger import EventTypes @@ -836,6 +836,55 @@ def test_step_function_trace_data_sfn_root(self): expected_context, ) + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_event_bridge(self): + lambda_ctx = get_mock_context() + sfn_event = { + "_datadog": { + "Execution": { + "StartTime": "2025-03-11T01:16:31.408Z", + "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", + "RedriveCount": 0, + "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", + "Name": "eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", + }, + "StateMachine": { + "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine", + "Name": "abhinav-inner-state-machine", + }, + "State": { + "EnteredTime": "2025-03-11T01:16:31.448Z", + "RetryCount": 0, + "Name": "EventBridge PutEvents", + }, + "serverless-version": "v1", + "RootExecutionId": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", + } + } + ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=4521899030418994483, + span_id=6880978411788117524, + sampling_priority=1, + meta={"_dd.p.tid": "12d1270d99cc5e03"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "4521899030418994483", + TraceHeader.PARENT_ID: "2685222157636933868", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: "_dd.p.tid=12d1270d99cc5e03", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + class TestXRayContextConversion(unittest.TestCase): def test_convert_xray_trace_id(self): @@ -2282,6 +2331,69 @@ def test_deterministic_m5_hash__always_leading_with_zero(self): if len(result_in_binary) == 66: # "0b" + 64 bits. self.assertTrue(result_in_binary.startswith("0b0")) + def test_is_step_function_event_jsonata(self): + event = { + "_datadog": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + "x-datadog-trace-id": "5821803790426892636", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", + "serverless-version": "v1", + } + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_jsonpath(self): + event = { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_legacy_lambda(self): + event = { + "Payload": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + } + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_dd_header(self): + event = { + "_datadog": { + "x-datadog-trace-id": "5821803790426892636", + "x-datadog-parent-id": "5821803790426892636", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", + "x-datadog-sampling-priority": "1", + } + } + self.assertFalse(is_step_function_event(event)) + class TestExceptionOutsideHandler(unittest.TestCase): @patch("datadog_lambda.tracing.dd_tracing_enabled", True) From 0a51475a1fb04fde3aeda3443f8f949aa3430bf0 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 11 Mar 2025 13:38:54 -0400 Subject: [PATCH 02/15] fix circular import --- datadog_lambda/tracing.py | 27 +--------------- datadog_lambda/trigger.py | 27 ++++++++++++++-- tests/test_tracing.py | 64 -------------------------------------- tests/test_trigger.py | 65 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 92 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 03f10d19..95f5af30 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -39,6 +39,7 @@ _EventSource, parse_event_source, get_first_record, + is_step_function_event, EventTypes, EventSubtypes, ) @@ -479,32 +480,6 @@ def extract_context_from_step_functions(event, lambda_context): return extract_context_from_lambda_context(lambda_context) -def is_step_function_event(event): - """ - Check if the event is a step function that invoked the current lambda. - - The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a - "_datadog" for JSONata style context propagation. - - The actual event must contain "Execution", "StateMachine", and "State" fields. - """ - event = event.get("Payload", event) - - # JSONPath style - if all(field in event for field in ("Execution", "StateMachine", "State")): - return True - - # JSONata style - if "_datadog" in event: - event = event["_datadog"] - return all( - field in event - for field in ("Execution", "StateMachine", "State", "serverless-version") - ) - - return False - - def extract_context_custom_extractor(extractor, event, lambda_context): """ Extract Datadog trace context using a custom trace extractor function diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index c583763a..1dd73dd7 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -10,8 +10,6 @@ from enum import Enum from typing import Any -from datadog_lambda.tracing import is_step_function_event - class _stringTypedEnum(Enum): """ @@ -369,3 +367,28 @@ def extract_http_status_code_tag(trigger_tags, response): status_code = response.status_code return str(status_code) + +def is_step_function_event(event): + """ + Check if the event is a step function that invoked the current lambda. + + The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a + "_datadog" for JSONata style context propagation. + + The actual event must contain "Execution", "StateMachine", and "State" fields. + """ + event = event.get("Payload", event) + + # JSONPath style + if all(field in event for field in ("Execution", "StateMachine", "State")): + return True + + # JSONata style + if "_datadog" in event: + event = event["_datadog"] + return all( + field in event + for field in ("Execution", "StateMachine", "State", "serverless-version") + ) + + return False \ No newline at end of file diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 1b4ad0d7..aeddb2c7 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -42,7 +42,6 @@ service_mapping as global_service_mapping, propagator, emit_telemetry_on_exception_outside_of_handler, - is_step_function_event, ) from datadog_lambda.trigger import EventTypes @@ -2331,69 +2330,6 @@ def test_deterministic_m5_hash__always_leading_with_zero(self): if len(result_in_binary) == 66: # "0b" + 64 bits. self.assertTrue(result_in_binary.startswith("0b0")) - def test_is_step_function_event_jsonata(self): - event = { - "_datadog": { - "Execution": { - "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", - "RedriveCount": 0, - }, - "StateMachine": {}, - "State": { - "Name": "my-awesome-state", - "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", - "RetryCount": 0, - }, - "x-datadog-trace-id": "5821803790426892636", - "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", - "serverless-version": "v1", - } - } - self.assertTrue(is_step_function_event(event)) - - def test_is_step_function_event_jsonpath(self): - event = { - "Execution": { - "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", - "RedriveCount": 0, - }, - "StateMachine": {}, - "State": { - "Name": "my-awesome-state", - "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", - "RetryCount": 0, - }, - } - self.assertTrue(is_step_function_event(event)) - - def test_is_step_function_event_legacy_lambda(self): - event = { - "Payload": { - "Execution": { - "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", - "RedriveCount": 0, - }, - "StateMachine": {}, - "State": { - "Name": "my-awesome-state", - "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", - "RetryCount": 0, - }, - } - } - self.assertTrue(is_step_function_event(event)) - - def test_is_step_function_event_dd_header(self): - event = { - "_datadog": { - "x-datadog-trace-id": "5821803790426892636", - "x-datadog-parent-id": "5821803790426892636", - "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", - "x-datadog-sampling-priority": "1", - } - } - self.assertFalse(is_step_function_event(event)) - class TestExceptionOutsideHandler(unittest.TestCase): @patch("datadog_lambda.tracing.dd_tracing_enabled", True) diff --git a/tests/test_trigger.py b/tests/test_trigger.py index be028a23..a005100f 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -9,6 +9,7 @@ get_event_source_arn, extract_trigger_tags, extract_http_status_code_tag, + is_step_function_event, ) from tests.utils import get_mock_context @@ -543,3 +544,67 @@ def test_extract_http_status_code_tag_from_response_object(self): response.status_code = 403 status_code = extract_http_status_code_tag(trigger_tags, response) self.assertEqual(status_code, "403") + +class IsStepFunctionEvent(unittest.TestCase): + def test_is_step_function_event_jsonata(self): + event = { + "_datadog": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + "x-datadog-trace-id": "5821803790426892636", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", + "serverless-version": "v1", + } + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_jsonpath(self): + event = { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_legacy_lambda(self): + event = { + "Payload": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + "RedriveCount": 0, + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + "RetryCount": 0, + }, + } + } + self.assertTrue(is_step_function_event(event)) + + def test_is_step_function_event_dd_header(self): + event = { + "_datadog": { + "x-datadog-trace-id": "5821803790426892636", + "x-datadog-parent-id": "5821803790426892636", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000", + "x-datadog-sampling-priority": "1", + } + } + self.assertFalse(is_step_function_event(event)) \ No newline at end of file From 5700d22611d408aa438a29d685a4c2e1960583f2 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 11 Mar 2025 13:46:08 -0400 Subject: [PATCH 03/15] fix lint --- datadog_lambda/trigger.py | 3 ++- tests/test_tracing.py | 10 +++++----- tests/test_trigger.py | 3 ++- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index 1dd73dd7..cc8fffed 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -368,6 +368,7 @@ def extract_http_status_code_tag(trigger_tags, response): return str(status_code) + def is_step_function_event(event): """ Check if the event is a step function that invoked the current lambda. @@ -391,4 +392,4 @@ def is_step_function_event(event): for field in ("Execution", "StateMachine", "State", "serverless-version") ) - return False \ No newline at end of file + return False diff --git a/tests/test_tracing.py b/tests/test_tracing.py index aeddb2c7..28e2746e 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -863,19 +863,19 @@ def test_step_function_trace_data_event_bridge(self): ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) self.assertEqual(source, "event") expected_context = Context( - trace_id=4521899030418994483, - span_id=6880978411788117524, + trace_id=4728686021345621131, + span_id=2685222157636933868, sampling_priority=1, - meta={"_dd.p.tid": "12d1270d99cc5e03"}, + meta={"_dd.p.tid": "7683d2257c051fce"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( get_dd_trace_context(), { - TraceHeader.TRACE_ID: "4521899030418994483", + TraceHeader.TRACE_ID: "4728686021345621131", TraceHeader.PARENT_ID: "2685222157636933868", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=12d1270d99cc5e03", + TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) diff --git a/tests/test_trigger.py b/tests/test_trigger.py index a005100f..9cb088f1 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -545,6 +545,7 @@ def test_extract_http_status_code_tag_from_response_object(self): status_code = extract_http_status_code_tag(trigger_tags, response) self.assertEqual(status_code, "403") + class IsStepFunctionEvent(unittest.TestCase): def test_is_step_function_event_jsonata(self): event = { @@ -607,4 +608,4 @@ def test_is_step_function_event_dd_header(self): "x-datadog-sampling-priority": "1", } } - self.assertFalse(is_step_function_event(event)) \ No newline at end of file + self.assertFalse(is_step_function_event(event)) From f524509a3ae7ea7799ac98402d5df5e89e4ecaab Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 11 Mar 2025 13:51:20 -0400 Subject: [PATCH 04/15] fix parent_id value in unit test --- 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 28e2746e..bada79fa 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -873,7 +873,7 @@ def test_step_function_trace_data_event_bridge(self): get_dd_trace_context(), { TraceHeader.TRACE_ID: "4728686021345621131", - TraceHeader.PARENT_ID: "2685222157636933868", + TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", }, From 55ff07505c1b5e05d8479e2921497dab9d49c641 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 11 Mar 2025 17:10:41 -0400 Subject: [PATCH 05/15] use provided timestamp from step function --- datadog_lambda/tracing.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 95f5af30..12d3aa2d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -1303,8 +1303,15 @@ def create_inferred_span_from_eventbridge_event(event, context): synchronicity="async", tag_source="self", ) - dt_format = "%Y-%m-%dT%H:%M:%SZ" - timestamp = event.get("time") + + # Use more granular timestamp from upstream Step Function if possible + if is_step_function_event(event.get("detail")): + timestamp = event.get("detail").get("_datadog").get("State").get("EnteredTime") + dt_format = "%Y-%m-%dT%H:%M:%S.%fZ" + else: + timestamp = event.get("time") + dt_format = "%Y-%m-%dT%H:%M:%SZ" + dt = datetime.strptime(timestamp, dt_format) tracer.set_tags(_dd_origin) From d7daf0d6817f9020374f705e057e01f1a86de9c3 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Wed, 12 Mar 2025 14:34:13 -0400 Subject: [PATCH 06/15] extract context in eventbridge sqs case --- datadog_lambda/tracing.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 12d3aa2d..6983180e 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -314,6 +314,14 @@ def _extract_context_from_eventbridge_sqs_event(event): body = json.loads(body_str) detail = body.get("detail") dd_context = detail.get("_datadog") + if is_step_function_event(dd_context): + try: + return extract_context_from_step_functions(detail, None) + except Exception: + logger.debug( + "Failed to extract Step Functions context from EventBridge to SQS event." + ) + return propagator.extract(dd_context) From e490610979326502fc8b431291072362a8d333c2 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 10:16:45 -0400 Subject: [PATCH 07/15] Update datadog_lambda/trigger.py Co-authored-by: kimi <47579703+kimi-p@users.noreply.github.com> --- datadog_lambda/trigger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index cc8fffed..5f291a86 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -386,7 +386,7 @@ def is_step_function_event(event): # JSONata style if "_datadog" in event: - event = event["_datadog"] + datadog_context = event["_datadog"] return all( field in event for field in ("Execution", "StateMachine", "State", "serverless-version") From e31c8f87e426792daf4dab9853d79033ea28748c Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 11:53:39 -0400 Subject: [PATCH 08/15] add sqs/sns support --- datadog_lambda/tracing.py | 35 +++++++++++++++++++++++------------ datadog_lambda/trigger.py | 3 +-- 2 files changed, 24 insertions(+), 14 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 6983180e..d2941a85 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -272,6 +272,15 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): if dd_json_data: dd_data = json.loads(dd_json_data) + + if is_step_function_event(dd_data): + try: + return extract_context_from_step_functions(dd_data, None) + except Exception: + logger.debug( + "Failed to extract Step Functions context from SQS/SNS event." + ) + return propagator.extract(dd_data) else: # Handle case where trace context is injected into attributes.AWSTraceHeader @@ -314,9 +323,10 @@ def _extract_context_from_eventbridge_sqs_event(event): body = json.loads(body_str) detail = body.get("detail") dd_context = detail.get("_datadog") + if is_step_function_event(dd_context): try: - return extract_context_from_step_functions(detail, None) + return extract_context_from_step_functions(dd_context, None) except Exception: logger.debug( "Failed to extract Step Functions context from EventBridge to SQS event." @@ -338,8 +348,10 @@ def extract_context_from_eventbridge_event(event, lambda_context): dd_context = detail.get("_datadog") if not dd_context: return extract_context_from_lambda_context(lambda_context) + if is_step_function_event(dd_context): - return extract_context_from_step_functions(detail, lambda_context) + return extract_context_from_step_functions(dd_context, lambda_context) + return propagator.extract(dd_context) except Exception as e: logger.debug("The trace extractor returned with error %s", e) @@ -450,26 +462,24 @@ def extract_context_from_step_functions(event, lambda_context): """ try: event = event.get("Payload", event) + event = event.get("_datadog", event) meta = {} - dd_data = event.get("_datadog") - if dd_data and dd_data.get("serverless-version") == "v1": - if "x-datadog-trace-id" in dd_data: # lambda root - trace_id = int(dd_data.get("x-datadog-trace-id")) - high_64_bit_trace_id = _parse_high_64_bits( - dd_data.get("x-datadog-tags") - ) + if event.get("serverless-version") == "v1": + if "x-datadog-trace-id" in event: # lambda root + trace_id = int(event.get("x-datadog-trace-id")) + high_64_bit_trace_id = _parse_high_64_bits(event.get("x-datadog-tags")) if high_64_bit_trace_id: meta["_dd.p.tid"] = high_64_bit_trace_id else: # sfn root - root_execution_id = dd_data.get("RootExecutionId") + root_execution_id = event.get("RootExecutionId") trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS) meta["_dd.p.tid"] = _generate_sfn_trace_id( root_execution_id, HIGHER_64_BITS ) - parent_id = _generate_sfn_parent_id(dd_data) + parent_id = _generate_sfn_parent_id(event) else: execution_id = event.get("Execution").get("Id") trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS) @@ -1331,7 +1341,8 @@ def create_inferred_span_from_eventbridge_event(event, context): span.start = dt.replace(tzinfo=timezone.utc).timestamp() # Since inferred span will later parent Lambda, preserve Lambda's current parent - span.parent_id = dd_trace_context.span_id + if dd_trace_context.span_id: + span.parent_id = dd_trace_context.span_id return span diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index 5f291a86..d09d23b2 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -386,9 +386,8 @@ def is_step_function_event(event): # JSONata style if "_datadog" in event: - datadog_context = event["_datadog"] return all( - field in event + field in event["_datadog"] for field in ("Execution", "StateMachine", "State", "serverless-version") ) From dc4a283db4e64df39fdf71d87841b83572e2983c Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 11:54:03 -0400 Subject: [PATCH 09/15] add more test cases --- tests/test_tracing.py | 276 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 251 insertions(+), 25 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index bada79fa..7a4097bc 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -836,35 +836,55 @@ def test_step_function_trace_data_sfn_root(self): ) @with_trace_propagation_style("datadog") - def test_step_function_trace_data_event_bridge(self): + def test_step_function_trace_data_eventbridge(self): lambda_ctx = get_mock_context() - sfn_event = { - "_datadog": { - "Execution": { - "StartTime": "2025-03-11T01:16:31.408Z", - "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", - "RedriveCount": 0, - "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", - "Name": "eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", - }, - "StateMachine": { - "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine", - "Name": "abhinav-inner-state-machine", - }, - "State": { - "EnteredTime": "2025-03-11T01:16:31.448Z", - "RetryCount": 0, - "Name": "EventBridge PutEvents", + eventbridge_event = { + "version": "0", + "id": "eaacd8db-02de-ab13-ed5a-8ffb84048294", + "detail-type": "StepFunctionTask", + "source": "my.eventbridge", + "account": "425362996713", + "time": "2025-03-13T15:17:34Z", + "region": "sa-east-1", + "resources": [ + "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine", + "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:912eaa4c-291a-488a-bda3-d06bcc21203d", + ], + "detail": { + "Message": "Hello from Step Functions!", + "TaskToken": "AQCEAAAAKgAAAAMAAAAAAAAAAeMHr6sb8Ll5IKntjIiLGaBkaNeweo84kKYKDTvDaSAP1vjuYRJEGqFdHsKMyZL8ZcgAdanKpkbhPEN5hpoCe+BH9KblWeDsJxkDCk/meN5SaPlC1qS7Q/7/KqBq+tmAOCSy+MjdqFsnihy5Yo6g6C9uuPn7ccSB/609d8pznFm9nigEos/82emwi18lm67/+/bn4RTX4S7qV4RoGWUWUPeHfr34xWOipCt4SVDkoQPZdRVpq3wyRJP2zcK0zup24/opJqKKSCI5Q9orALNB2jEjDyQ9LE4mSrafoe0tcm/bOAGfrcpR3AwtArUiF6JPYd7Nw0XWWyPXFBjiQTJDhZFlGfllJ1N91eiN8wlzUX1+I0vw/t2PoEmuQ2VCJYCbl1ybjX/tQ97GZ9ogjY9N7VYy5uD5xfZ6VAyetUR06HUtbUIXTVxULm7wmsHb979W/fIQXsrxbFzc0+ypKaqGXJBoq7xX//irjpuNhWg1Wgfn0hxuXl5oN/LkqI83T8f9SdnJMxRDpaHDpttqbjVESB/Pf9o7gakjJj12+r2uiJNc81k50uhuHdFOGsImFHKV8hb1LGcq0ZzUKT5SbEDV2k+ezOP+O9Sk4c0unbpNLM3PKLKxVLhu2gtiIIVCHUHGmumW", + "_datadog": { + "Execution": { + "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:912eaa4c-291a-488a-bda3-d06bcc21203d", + "StartTime": "2025-03-13T15:17:33.972Z", + "Name": "912eaa4c-291a-488a-bda3-d06bcc21203d", + "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", + "RedriveCount": 0, + }, + "StateMachine": { + "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine", + "Name": "abhinav-inner-state-machine", + }, + "State": { + "Name": "EventBridge PutEvents", + "EnteredTime": "2025-03-13T15:17:34.008Z", + "RetryCount": 0, + }, + "Task": { + "Token": "AQCEAAAAKgAAAAMAAAAAAAAAAeMHr6sb8Ll5IKntjIiLGaBkaNeweo84kKYKDTvDaSAP1vjuYRJEGqFdHsKMyZL8ZcgAdanKpkbhPEN5hpoCe+BH9KblWeDsJxkDCk/meN5SaPlC1qS7Q/7/KqBq+tmAOCSy+MjdqFsnihy5Yo6g6C9uuPn7ccSB/609d8pznFm9nigEos/82emwi18lm67/+/bn4RTX4S7qV4RoGWUWUPeHfr34xWOipCt4SVDkoQPZdRVpq3wyRJP2zcK0zup24/opJqKKSCI5Q9orALNB2jEjDyQ9LE4mSrafoe0tcm/bOAGfrcpR3AwtArUiF6JPYd7Nw0XWWyPXFBjiQTJDhZFlGfllJ1N91eiN8wlzUX1+I0vw/t2PoEmuQ2VCJYCbl1ybjX/tQ97GZ9ogjY9N7VYy5uD5xfZ6VAyetUR06HUtbUIXTVxULm7wmsHb979W/fIQXsrxbFzc0+ypKaqGXJBoq7xX//irjpuNhWg1Wgfn0hxuXl5oN/LkqI83T8f9SdnJMxRDpaHDpttqbjVESB/Pf9o7gakjJj12+r2uiJNc81k50uhuHdFOGsImFHKV8hb1LGcq0ZzUKT5SbEDV2k+ezOP+O9Sk4c0unbpNLM3PKLKxVLhu2gtiIIVCHUHGmumW" + }, + "RootExecutionId": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:912eaa4c-291a-488a-bda3-d06bcc21203d", + "serverless-version": "v1", }, - "serverless-version": "v1", - "RootExecutionId": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed", - } + }, } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) + ctx, source, event_source = extract_dd_trace_context( + eventbridge_event, lambda_ctx + ) self.assertEqual(source, "event") expected_context = Context( - trace_id=4728686021345621131, - span_id=2685222157636933868, + trace_id=3401561763239692811, + span_id=10430178702434539423, sampling_priority=1, meta={"_dd.p.tid": "7683d2257c051fce"}, ) @@ -872,7 +892,7 @@ def test_step_function_trace_data_event_bridge(self): self.assertEqual( get_dd_trace_context(), { - TraceHeader.TRACE_ID: "4728686021345621131", + TraceHeader.TRACE_ID: "3401561763239692811", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", @@ -884,6 +904,212 @@ def test_step_function_trace_data_event_bridge(self): expected_context, ) + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_sqs(self): + lambda_ctx = get_mock_context() + sqs_event = { + "Records": [ + { + "EventSource": "aws:sns", + "EventVersion": "1.0", + "EventSubscriptionArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "Sns": { + "Type": "Notification", + "MessageId": "e39184ea-bfd8-5efa-96fe-e4a64a457ff7", + "TopicArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic", + "Subject": None, + "Message": "{}", + "Timestamp": "2025-03-13T15:01:49.942Z", + "SignatureVersion": "1", + "Signature": "WJHKq+pNOLgxa7+dB1dud02RM/30Jvz+KiMZzjRl38/Pphz90H24eGyIbnq3BJXYEyawFCHC6sq/5HcwXouGc5gbah6he+JpqXahMEs6cyMs2tg9SXxooRHEGv5iiZXKhnDcJYOrQ+iFExO9w+WFWfJjO2m/EDVVSYvuDjDV7mmTwAgEOD0zUvWpT7wOeKGG5Uk916Ppy3iMV7sCoHV/RwVikdhCWDDmxbdqteGduAXPdGESE/aj6kUx9ibEOKXyhC+7H1/j0tlhUchl6LZsTf1Gaiq2yEqKXKvsupcG3hRZ6FtIWP0jGlFhpW5EHc2oiHIVOsQceCYPqXYMCZvFuA==", + "SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem", + "UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "MessageAttributes": { + "_datadog": { + "Type": "String", + "Value": '{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:79478846-0cff-44de-91f5-02c96ff65762","StartTime":"2025-03-13T15:01:49.738Z","Name":"79478846-0cff-44de-91f5-02c96ff65762","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"SNS Publish","EnteredTime":"2025-03-13T15:01:49.768Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:79478846-0cff-44de-91f5-02c96ff65762","serverless-version":"v1"}', + } + }, + }, + } + ] + } + ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=3818106616964044169, + span_id=15912108710769293902, + sampling_priority=1, + meta={"_dd.p.tid": "7683d2257c051fce"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "3818106616964044169", + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_eventbridge_sqs(self): + lambda_ctx = get_mock_context() + eventbridge_sqs_event = { + "Records": [ + { + "messageId": "9ed082ad-2f4d-4309-ab99-9553d2be5613", + "receiptHandle": "AQEB6z7FatNIXbWOTC4Bx+udD0flrnT7XMehruTohl8O2KI2t9hvo5oxGIOhwcb+QtS5aRXsFE35TgGE8kZHlHK7Sa8jQUen6XmsPG7qB6BPdXjr0eunM2SDAtLj0mDSKx907VIKRYQG+qpI9ZyNK7Bi786oQIz2UkZGZru9zlXxJtAQiXBqfJ+OfTzhIwkPu04czU6lYfAbxdyNaBNdBEsTNJKPjquvcq1ZBVCHkn9L6wo8jha6XreoeS2WJ5N26ZLKtAl3wlSUByB92OKZU2mEuNboyY7bgK+nkx4N8fVVrafVXnY9YHuq60eQcZ/nusWFeJlVyN7NFypYP2IOn25xylltEACKbgUdEsFU2h5k7yI2DVk5eAt9vB6qmAJlgfkGsXG0SZrCADoIKXl9jpwajw==", + "body": '{"version":"0","id":"ff6d828b-b35e-abdf-64b6-6ea2cf698c0b","detail-type":"StepFunctionTask","source":"my.eventbridge","account":"425362996713","time":"2025-03-13T15:14:21Z","region":"sa-east-1","resources":["arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985"],"detail":{"Message":"Hello from Step Functions!","TaskToken":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p","_datadog":{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","StartTime":"2025-03-13T15:14:21.730Z","Name":"fe087266-fe48-4a31-a21b-691f4e7ea985","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"EventBridge PutEvents","EnteredTime":"2025-03-13T15:14:21.765Z","RetryCount":0},"Task":{"Token":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p"},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","serverless-version":"v1"}}}', + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1741878862068", + "SenderId": "AROAWGCM4HXUUNHLDXVER:6145b5ba998f311c8ac27f5cade2b915", + "ApproximateFirstReceiveTimestamp": "1741878862075", + }, + "messageAttributes": {}, + "md5OfBody": "e5cf8197b304a4dd4fd5db8e4842484b", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", + "awsRegion": "sa-east-1", + } + ] + } + ctx, source, event_source = extract_dd_trace_context( + eventbridge_sqs_event, lambda_ctx + ) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=6527209323865742984, + span_id=14276854885394865473, + sampling_priority=1, + meta={"_dd.p.tid": "7683d2257c051fce"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "6527209323865742984", + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_sns(self): + lambda_ctx = get_mock_context() + sns_event = { + "Records": [ + { + "EventSource": "aws:sns", + "EventVersion": "1.0", + "EventSubscriptionArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "Sns": { + "Type": "Notification", + "MessageId": "7bc0c17d-bf88-5ff4-af7f-a131463a0d90", + "TopicArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic", + "Subject": None, + "Message": "{}", + "Timestamp": "2025-03-13T15:19:14.245Z", + "SignatureVersion": "1", + "Signature": "r8RoYzq4uNcq0yj7sxcp8sTbFiDk8zqtocG7mJuE2MPVuR8O5eNg2ohofokUnC84xADlCq5k6ElP55lbbY36tQO+qDGdV6+TGN4bAL9FiQrDE6tQYYJdlv/sYE7iOOgnRBC9ljEdCIDNtQNGCfND/8JzatPg8KAy7xMRcLrGWu4xIMEysqNTz7rETfhdZjLQPssAht44KcoUJCH4/VuB+B9W1RhwA+M8Q3tqxzahIXzcgDM8OlmfkBlXo4FDVF3WUzjXLf9AMOg+66GupjQFtUpmRMkA8KXSV1HCso7e6nIIWtOnUoWeDDUfQPFFq4TNSlb6h2NuebaHdnW5nhxnJQ==", + "SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem", + "UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "MessageAttributes": { + "_datadog": { + "Type": "String", + "Value": '{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","StartTime":"2025-03-13T15:19:14.019Z","Name":"11623e4f-70ee-4330-8fbe-955152dea54c","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"SNS Publish","EnteredTime":"2025-03-13T15:19:14.061Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","serverless-version":"v1"}', + } + }, + }, + } + ] + } + ctx, source, event_source = extract_dd_trace_context(sns_event, lambda_ctx) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=1459500239678510857, + span_id=13193042003602978730, + sampling_priority=1, + meta={"_dd.p.tid": "7683d2257c051fce"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "1459500239678510857", + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_sns_sqs(self): + lambda_ctx = get_mock_context() + sns_sqs_event = { + "Records": [ + { + "messageId": "9ec3339f-cd1a-43ba-9681-3e9113b430d3", + "receiptHandle": "AQEBJ5gIvqEWQt39NHPMAoK57cGgKtrgTtckWeWdDRi2FeucYr6pBhNjzXuUrmoHZMozX1WaoABtfQ5+kX5ucDBpA2Ci3Q07Z4MYvA6X0Sw13HCkiBnLrHPmH/F3rUBjvdRkIIKqA2ACX58MdkaYGNpqsHTJHB613wa8z4zurK0u7eUIXrr+e+gtsuPD39hiWlJo7cpBVv7y178rzMX8gPQTnRJv1cjhCHENtjWTSmfFC5N+BIQNIcjFsTTDRSovZlNIfAEuS+uowgzk0DUyoTJD5nFTL8lQHeXGRCUQe58/UY9OwRXEFVPGZOQR4OI9Wa4Kf/keFypTk9YwC9DhSeKvzZ0wBvejyl1n0ztT45+XYoWfi0mxGWM5b7r9wT36RDmjnM6vszH/d3fhZSRPASxWBQ==", + "body": '{\n "Type" : "Notification",\n "MessageId" : "1f3078d0-c792-5cf3-a130-189c3b846a3f",\n "TopicArn" : "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic",\n "Message" : "{}",\n "Timestamp" : "2025-03-13T15:29:26.348Z",\n "SignatureVersion" : "1",\n "Signature" : "mxOqAQ5o/isJrMS0PezHKRaA3g8Z/8YDbkToqhJub6I66LGtl+NYhyfTyllbgxvRP2XD2meKPRSgPI3nLyq8UHsWgyYwe3Tsv8QpRunCVE9Pebh+V1LGPWfjOiL0e+bnaj956QJD99560LJ6bzWP9QO584/zfOdcw6E5XQZfAI+pvEsf28Dy0WJO/lWTATRZDf8wGhmc7uKI1ZMsrOaNoUD8PXVqsI4yrJHxhzMb3SrC7YjI/PnNIbcn6ezwprbUdbZvyNAfJiE0k5IlppA089tMXC/ItgC7AgQhG9huPdKi5KdWGACK7gEwqmFwL+5T33sUXDaH2g58WhCs76pKEw==",\n "SigningCertURL" : "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem",\n "UnsubscribeURL" : "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:5f64545d-ae9a-4a5f-a7ee-798a0bd8519e",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"Execution\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"StartTime\\":\\"2025-03-13T15:29:26.144Z\\",\\"Name\\":\\"37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"RoleArn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j\\",\\"RedriveCount\\":0},\\"StateMachine\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine\\",\\"Name\\":\\"abhinav-inner-state-machine\\"},\\"State\\":{\\"Name\\":\\"SNS Publish\\",\\"EnteredTime\\":\\"2025-03-13T15:29:26.182Z\\",\\"RetryCount\\":0},\\"RootExecutionId\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"serverless-version\\":\\"v1\\"}"}\n }\n}', + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1741879766424", + "SenderId": "AIDAIOA2GYWSHW4E2VXIO", + "ApproximateFirstReceiveTimestamp": "1741879766432", + }, + "messageAttributes": {}, + "md5OfBody": "52af59de28507d7e67324b46c95337d8", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", + "awsRegion": "sa-east-1", + } + ] + } + ctx, source, event_source = extract_dd_trace_context( + sns_sqs_event, lambda_ctx + ) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=5708348677301000120, + span_id=18223515719478572006, + sampling_priority=1, + meta={"_dd.p.tid": "7683d2257c051fce"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "5708348677301000120", + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + class TestXRayContextConversion(unittest.TestCase): def test_convert_xray_trace_id(self): From f59de194ebd5c6a7c08af5e8627a7e337930ebdf Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 12:05:58 -0400 Subject: [PATCH 10/15] fix expected _dd.p.tid vals --- tests/test_tracing.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 7a4097bc..0234358c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -886,7 +886,7 @@ def test_step_function_trace_data_eventbridge(self): trace_id=3401561763239692811, span_id=10430178702434539423, sampling_priority=1, - meta={"_dd.p.tid": "7683d2257c051fce"}, + meta={"_dd.p.tid": "a49ff3b7fb47b0b"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( @@ -895,7 +895,7 @@ def test_step_function_trace_data_eventbridge(self): TraceHeader.TRACE_ID: "3401561763239692811", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + TraceHeader.TAGS: "_dd.p.tid=a49ff3b7fb47b0b", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) @@ -940,7 +940,7 @@ def test_step_function_trace_data_sqs(self): trace_id=3818106616964044169, span_id=15912108710769293902, sampling_priority=1, - meta={"_dd.p.tid": "7683d2257c051fce"}, + meta={"_dd.p.tid": "3a4fd1a254eb514a"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( @@ -949,7 +949,7 @@ def test_step_function_trace_data_sqs(self): TraceHeader.TRACE_ID: "3818106616964044169", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + TraceHeader.TAGS: "_dd.p.tid=3a4fd1a254eb514a", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) @@ -989,7 +989,7 @@ def test_step_function_trace_data_eventbridge_sqs(self): trace_id=6527209323865742984, span_id=14276854885394865473, sampling_priority=1, - meta={"_dd.p.tid": "7683d2257c051fce"}, + meta={"_dd.p.tid": "2ee7d9862d048173"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( @@ -998,7 +998,7 @@ def test_step_function_trace_data_eventbridge_sqs(self): TraceHeader.TRACE_ID: "6527209323865742984", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + TraceHeader.TAGS: "_dd.p.tid=2ee7d9862d048173", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) @@ -1043,7 +1043,7 @@ def test_step_function_trace_data_sns(self): trace_id=1459500239678510857, span_id=13193042003602978730, sampling_priority=1, - meta={"_dd.p.tid": "7683d2257c051fce"}, + meta={"_dd.p.tid": "fafc98885fd4647"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( @@ -1052,7 +1052,7 @@ def test_step_function_trace_data_sns(self): TraceHeader.TRACE_ID: "1459500239678510857", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + TraceHeader.TAGS: "_dd.p.tid=fafc98885fd4647", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) @@ -1092,7 +1092,7 @@ def test_step_function_trace_data_sns_sqs(self): trace_id=5708348677301000120, span_id=18223515719478572006, sampling_priority=1, - meta={"_dd.p.tid": "7683d2257c051fce"}, + meta={"_dd.p.tid": "45457f5f3fde3fa1"}, ) self.assertEqual(ctx, expected_context) self.assertEqual( @@ -1101,7 +1101,7 @@ def test_step_function_trace_data_sns_sqs(self): TraceHeader.TRACE_ID: "5708348677301000120", TraceHeader.PARENT_ID: "10713633173203262661", TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce", + TraceHeader.TAGS: "_dd.p.tid=45457f5f3fde3fa1", }, ) create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) From 6222016bf460f704ea2bf0bae2bcfce6f41f8cf1 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 15:50:15 -0400 Subject: [PATCH 11/15] add lru cache and remove usage of all() --- datadog_lambda/trigger.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index d09d23b2..edc04d12 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -5,6 +5,8 @@ import base64 import gzip +from functools import lru_cache + import ujson as json from io import BytesIO, BufferedReader from enum import Enum @@ -369,6 +371,7 @@ def extract_http_status_code_tag(trigger_tags, response): return str(status_code) +@lru_cache(maxsize=8) def is_step_function_event(event): """ Check if the event is a step function that invoked the current lambda. @@ -381,14 +384,15 @@ def is_step_function_event(event): event = event.get("Payload", event) # JSONPath style - if all(field in event for field in ("Execution", "StateMachine", "State")): + if "Execution" in event and "StateMachine" in event and "State" in event: return True # JSONata style - if "_datadog" in event: - return all( - field in event["_datadog"] - for field in ("Execution", "StateMachine", "State", "serverless-version") - ) - - return False + dd_context = event.get("_datadog") + return ( + dd_context + and "Execution" in dd_context + and "StateMachine" in dd_context + and "State" in dd_context + and "serverless-version" in dd_context + ) From 45ba94591d4386f12e741cf63a5c3f251db6489f Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 15:50:36 -0400 Subject: [PATCH 12/15] exception handling for timestamp parsing --- datadog_lambda/tracing.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index d2941a85..0fae76dd 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -349,8 +349,12 @@ def extract_context_from_eventbridge_event(event, lambda_context): if not dd_context: return extract_context_from_lambda_context(lambda_context) - if is_step_function_event(dd_context): - return extract_context_from_step_functions(dd_context, lambda_context) + try: + return extract_context_from_step_functions(dd_context, None) + except Exception: + logger.debug( + "Failed to extract Step Functions context from EventBridge event." + ) return propagator.extract(dd_context) except Exception as e: @@ -1322,13 +1326,16 @@ def create_inferred_span_from_eventbridge_event(event, context): tag_source="self", ) + timestamp = event.get("time") + dt_format = "%Y-%m-%dT%H:%M:%SZ" + # Use more granular timestamp from upstream Step Function if possible - if is_step_function_event(event.get("detail")): - timestamp = event.get("detail").get("_datadog").get("State").get("EnteredTime") - dt_format = "%Y-%m-%dT%H:%M:%S.%fZ" - else: - timestamp = event.get("time") - dt_format = "%Y-%m-%dT%H:%M:%SZ" + try: + if is_step_function_event(event.get("detail")): + timestamp = event["detail"]["_datadog"]["State"]["EnteredTime"] + dt_format = "%Y-%m-%dT%H:%M:%S.%fZ" + except (TypeError, KeyError, AttributeError): + logger.debug("Error parsing timestamp from Step Functions event") dt = datetime.strptime(timestamp, dt_format) From dc43c9e84351d53fa826682942aaccdccd64b14c Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 13 Mar 2025 15:55:55 -0400 Subject: [PATCH 13/15] remove memo --- datadog_lambda/trigger.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index edc04d12..708138bf 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -5,8 +5,6 @@ import base64 import gzip -from functools import lru_cache - import ujson as json from io import BytesIO, BufferedReader from enum import Enum @@ -371,7 +369,6 @@ def extract_http_status_code_tag(trigger_tags, response): return str(status_code) -@lru_cache(maxsize=8) def is_step_function_event(event): """ Check if the event is a step function that invoked the current lambda. From eec4743203497e8d519bdd599dfef0aef7e614bb Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Mon, 17 Mar 2025 09:36:39 -0400 Subject: [PATCH 14/15] refactor tests to use common helper --- tests/test_tracing.py | 491 ++++++++++++++---------------------------- 1 file changed, 161 insertions(+), 330 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0234358c..4097113b 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -612,9 +612,40 @@ def test_with_complete_datadog_trace_headers_with_trigger_tags(self): ] ) + def _test_step_function_trace_data_common( + self, event, expected_trace_id, expected_span_id, expected_tid + ): + """Common test logic for step function trace data tests""" + lambda_ctx = get_mock_context() + expected_context = Context( + trace_id=expected_trace_id, + span_id=expected_span_id, + sampling_priority=1, + meta={"_dd.p.tid": expected_tid}, + ) + expected_headers = { + TraceHeader.TRACE_ID: str(expected_trace_id), + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + TraceHeader.TAGS: f"_dd.p.tid={expected_tid}", + } + + ctx, source, event_source = extract_dd_trace_context(event, lambda_ctx) + + self.assertEqual(source, "event") + self.assertEqual(event_source.event_type, EventTypes.STEPFUNCTIONS) + self.assertEqual(ctx, expected_context) + self.assertEqual(get_dd_trace_context(), expected_headers) + + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + self.mock_send_segment.assert_called_with( + XraySubsegment.TRACE_KEY, + expected_context, + ) + @with_trace_propagation_style("datadog") def test_step_function_trace_data(self): - lambda_ctx = get_mock_context() + """Test basic step function trace data extraction""" sfn_event = { "Execution": { "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", @@ -633,79 +664,39 @@ def test_step_function_trace_data(self): "Name": "abhinav-activity-state-machine", }, } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=435175499815315247, - span_id=3929055471293792800, - sampling_priority=1, - meta={"_dd.p.tid": "3e7a89d1b7310603"}, + self._test_step_function_trace_data_common( + sfn_event, 435175499815315247, 3929055471293792800, "3e7a89d1b7310603" ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "435175499815315247", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=3e7a89d1b7310603", + + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_retry(self): + """Test step function trace data extraction with non-zero retry count""" + sfn_event = { + "Execution": { + "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", + "Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316", + "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", + "StartTime": "2024-12-04T19:38:04.069Z", + "RedriveCount": 0, }, + "State": { + "Name": "Lambda Invoke", + "EnteredTime": "2024-12-04T19:38:04.118Z", + "RetryCount": 1, + }, + "StateMachine": { + "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine", + "Name": "abhinav-activity-state-machine", + }, + } + self._test_step_function_trace_data_common( + sfn_event, 435175499815315247, 5063839446130725204, "3e7a89d1b7310603" ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) - - @with_trace_propagation_style("datadog") - def test_step_function_trace_data_retry(self): - lambda_ctx = get_mock_context() - sfn_event = { - "Execution": { - "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", - "Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316", - "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", - "StartTime": "2024-12-04T19:38:04.069Z", - "RedriveCount": 0, - }, - "State": { - "Name": "Lambda Invoke", - "EnteredTime": "2024-12-04T19:38:04.118Z", - "RetryCount": 1, - }, - "StateMachine": { - "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine", - "Name": "abhinav-activity-state-machine", - }, - } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=435175499815315247, - span_id=5063839446130725204, - sampling_priority=1, - meta={"_dd.p.tid": "3e7a89d1b7310603"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "435175499815315247", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=3e7a89d1b7310603", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) - # https://github.com/DataDog/logs-backend/blob/c17618cb552fc369ca40282bae0a65803f82f694/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L46 + # https://github.com/DataDog/logs-backend/blob/65ea567150f24e5498008f3cf8cabef9ea995f5d/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L45-L46 @with_trace_propagation_style("datadog") def test_step_function_trace_data_redrive(self): - lambda_ctx = get_mock_context() + """Test step function trace data extraction with non-zero redrive count""" sfn_event = { "Execution": { "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", @@ -724,33 +715,13 @@ def test_step_function_trace_data_redrive(self): "Name": "abhinav-activity-state-machine", }, } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=435175499815315247, - span_id=8782364156266188026, - sampling_priority=1, - meta={"_dd.p.tid": "3e7a89d1b7310603"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "435175499815315247", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=3e7a89d1b7310603", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, + self._test_step_function_trace_data_common( + sfn_event, 435175499815315247, 8782364156266188026, "3e7a89d1b7310603" ) @with_trace_propagation_style("datadog") def test_step_function_trace_data_lambda_root(self): - lambda_ctx = get_mock_context() + """Test JSONata style step function trace data extraction where there's an upstream Lambda""" sfn_event = { "_datadog": { "Execution": { @@ -768,33 +739,13 @@ def test_step_function_trace_data_lambda_root(self): "serverless-version": "v1", } } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=5821803790426892636, - span_id=6880978411788117524, - sampling_priority=1, - meta={"_dd.p.tid": "672a7cb100000000"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "5821803790426892636", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=672a7cb100000000", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, + self._test_step_function_trace_data_common( + sfn_event, 5821803790426892636, 6880978411788117524, "672a7cb100000000" ) @with_trace_propagation_style("datadog") def test_step_function_trace_data_sfn_root(self): - lambda_ctx = get_mock_context() + """Test JSONata style step function trace data extraction where there's an upstream step function""" sfn_event = { "_datadog": { "Execution": { @@ -811,33 +762,13 @@ def test_step_function_trace_data_sfn_root(self): "serverless-version": "v1", } } - ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=4521899030418994483, - span_id=6880978411788117524, - sampling_priority=1, - meta={"_dd.p.tid": "12d1270d99cc5e03"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "4521899030418994483", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=12d1270d99cc5e03", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, + self._test_step_function_trace_data_common( + sfn_event, 4521899030418994483, 6880978411788117524, "12d1270d99cc5e03" ) @with_trace_propagation_style("datadog") def test_step_function_trace_data_eventbridge(self): - lambda_ctx = get_mock_context() + """Test step function trace data extraction through EventBridge""" eventbridge_event = { "version": "0", "id": "eaacd8db-02de-ab13-ed5a-8ffb84048294", @@ -878,35 +809,16 @@ def test_step_function_trace_data_eventbridge(self): }, }, } - ctx, source, event_source = extract_dd_trace_context( - eventbridge_event, lambda_ctx - ) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=3401561763239692811, - span_id=10430178702434539423, - sampling_priority=1, - meta={"_dd.p.tid": "a49ff3b7fb47b0b"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "3401561763239692811", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=a49ff3b7fb47b0b", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, + self._test_step_function_trace_data_common( + eventbridge_event, + 3401561763239692811, + 10430178702434539423, + "a49ff3b7fb47b0b", ) @with_trace_propagation_style("datadog") def test_step_function_trace_data_sqs(self): - lambda_ctx = get_mock_context() + """Test step function trace data extraction through SQS""" sqs_event = { "Records": [ { @@ -934,181 +846,100 @@ def test_step_function_trace_data_sqs(self): } ] } - ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=3818106616964044169, - span_id=15912108710769293902, - sampling_priority=1, - meta={"_dd.p.tid": "3a4fd1a254eb514a"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "3818106616964044169", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=3a4fd1a254eb514a", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, + self._test_step_function_trace_data_common( + sqs_event, 3818106616964044169, 15912108710769293902, "3a4fd1a254eb514a" ) - @with_trace_propagation_style("datadog") - def test_step_function_trace_data_eventbridge_sqs(self): - lambda_ctx = get_mock_context() - eventbridge_sqs_event = { - "Records": [ - { - "messageId": "9ed082ad-2f4d-4309-ab99-9553d2be5613", - "receiptHandle": "AQEB6z7FatNIXbWOTC4Bx+udD0flrnT7XMehruTohl8O2KI2t9hvo5oxGIOhwcb+QtS5aRXsFE35TgGE8kZHlHK7Sa8jQUen6XmsPG7qB6BPdXjr0eunM2SDAtLj0mDSKx907VIKRYQG+qpI9ZyNK7Bi786oQIz2UkZGZru9zlXxJtAQiXBqfJ+OfTzhIwkPu04czU6lYfAbxdyNaBNdBEsTNJKPjquvcq1ZBVCHkn9L6wo8jha6XreoeS2WJ5N26ZLKtAl3wlSUByB92OKZU2mEuNboyY7bgK+nkx4N8fVVrafVXnY9YHuq60eQcZ/nusWFeJlVyN7NFypYP2IOn25xylltEACKbgUdEsFU2h5k7yI2DVk5eAt9vB6qmAJlgfkGsXG0SZrCADoIKXl9jpwajw==", - "body": '{"version":"0","id":"ff6d828b-b35e-abdf-64b6-6ea2cf698c0b","detail-type":"StepFunctionTask","source":"my.eventbridge","account":"425362996713","time":"2025-03-13T15:14:21Z","region":"sa-east-1","resources":["arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985"],"detail":{"Message":"Hello from Step Functions!","TaskToken":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p","_datadog":{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","StartTime":"2025-03-13T15:14:21.730Z","Name":"fe087266-fe48-4a31-a21b-691f4e7ea985","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"EventBridge PutEvents","EnteredTime":"2025-03-13T15:14:21.765Z","RetryCount":0},"Task":{"Token":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p"},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","serverless-version":"v1"}}}', - "attributes": { - "ApproximateReceiveCount": "1", - "SentTimestamp": "1741878862068", - "SenderId": "AROAWGCM4HXUUNHLDXVER:6145b5ba998f311c8ac27f5cade2b915", - "ApproximateFirstReceiveTimestamp": "1741878862075", - }, - "messageAttributes": {}, - "md5OfBody": "e5cf8197b304a4dd4fd5db8e4842484b", - "eventSource": "aws:sqs", - "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", - "awsRegion": "sa-east-1", - } - ] - } - ctx, source, event_source = extract_dd_trace_context( - eventbridge_sqs_event, lambda_ctx - ) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=6527209323865742984, - span_id=14276854885394865473, - sampling_priority=1, - meta={"_dd.p.tid": "2ee7d9862d048173"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_eventbridge_sqs(self): + """Test step function trace data extraction through EventBridge and SQS""" + eventbridge_sqs_event = { + "Records": [ { - TraceHeader.TRACE_ID: "6527209323865742984", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=2ee7d9862d048173", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) + "messageId": "9ed082ad-2f4d-4309-ab99-9553d2be5613", + "receiptHandle": "AQEB6z7FatNIXbWOTC4Bx+udD0flrnT7XMehruTohl8O2KI2t9hvo5oxGIOhwcb+QtS5aRXsFE35TgGE8kZHlHK7Sa8jQUen6XmsPG7qB6BPdXjr0eunM2SDAtLj0mDSKx907VIKRYQG+qpI9ZyNK7Bi786oQIz2UkZGZru9zlXxJtAQiXBqfJ+OfTzhIwkPu04czU6lYfAbxdyNaBNdBEsTNJKPjquvcq1ZBVCHkn9L6wo8jha6XreoeS2WJ5N26ZLKtAl3wlSUByB92OKZU2mEuNboyY7bgK+nkx4N8fVVrafVXnY9YHuq60eQcZ/nusWFeJlVyN7NFypYP2IOn25xylltEACKbgUdEsFU2h5k7yI2DVk5eAt9vB6qmAJlgfkGsXG0SZrCADoIKXl9jpwajw==", + "body": '{"version":"0","id":"ff6d828b-b35e-abdf-64b6-6ea2cf698c0b","detail-type":"StepFunctionTask","source":"my.eventbridge","account":"425362996713","time":"2025-03-13T15:14:21Z","region":"sa-east-1","resources":["arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985"],"detail":{"Message":"Hello from Step Functions!","TaskToken":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p","_datadog":{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","StartTime":"2025-03-13T15:14:21.730Z","Name":"fe087266-fe48-4a31-a21b-691f4e7ea985","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"EventBridge PutEvents","EnteredTime":"2025-03-13T15:14:21.765Z","RetryCount":0},"Task":{"Token":"AQCEAAAAKgAAAAMAAAAAAAAAAfi3HMLTw3u9h0vSmkjyHlK1tv5bQUyA7i+6LIvrBWu+3S+DMuQ79JpMtAuCaMN/AGSuGPO7OPeTNA/9v7/kzAsLoPzwPhbrDPXP4SVF1YIO663PvtX/tEWxnAfwLqwDyx8G8VEsVLcmiiOafFCKJwn0OP/DoAWc0sjhWwRxIoQ0ipBGhOqU8rO8SFZVvxUbkosNejnhT7B6314pC89JZLpXU7SxFe+XrgN+uRAvFxsH/+RwDf94xk5hhtukH7HzhJKWN2WCtUISd84pM/1V7ppDuJ3FHgJT22xQIbEGA9Q4o+pLLehzE2SHCdo7eWYQqN+7BanxBNMI6kBMaf5nuh9izAp38lsrmHJyO8NvXgWg+F9hoTZX4RpV9CCwvRFrCRcCeDq4/uJzbvB4AwwA2q2Llm0X8yH0pKvPZ2v7pl4nCWdnEgj920I8AmBCuozbKP7gJRnAqfx3MnOSkpZTeGnHkp0ly8EevwCT2zX/1GQnCAx02kBaDJgUMputFeruMBzwVtlEVBFUUgaWbJwHzz2htuAw282pdATrKfv4VV1N962uLBJ32wd9a92rX7VXXToitvZGIvf/Z7cu4xfAzxQH1rIQ3M4ojkR9r48qoYtnYDlEf+BkIL8L4+xpbRFSBk3p"},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:fe087266-fe48-4a31-a21b-691f4e7ea985","serverless-version":"v1"}}}', + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1741878862068", + "SenderId": "AROAWGCM4HXUUNHLDXVER:6145b5ba998f311c8ac27f5cade2b915", + "ApproximateFirstReceiveTimestamp": "1741878862075", + }, + "messageAttributes": {}, + "md5OfBody": "e5cf8197b304a4dd4fd5db8e4842484b", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", + "awsRegion": "sa-east-1", + } + ] + } + self._test_step_function_trace_data_common( + eventbridge_sqs_event, + 6527209323865742984, + 14276854885394865473, + "2ee7d9862d048173", + ) - @with_trace_propagation_style("datadog") - def test_step_function_trace_data_sns(self): - lambda_ctx = get_mock_context() - sns_event = { - "Records": [ - { - "EventSource": "aws:sns", - "EventVersion": "1.0", - "EventSubscriptionArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", - "Sns": { - "Type": "Notification", - "MessageId": "7bc0c17d-bf88-5ff4-af7f-a131463a0d90", - "TopicArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic", - "Subject": None, - "Message": "{}", - "Timestamp": "2025-03-13T15:19:14.245Z", - "SignatureVersion": "1", - "Signature": "r8RoYzq4uNcq0yj7sxcp8sTbFiDk8zqtocG7mJuE2MPVuR8O5eNg2ohofokUnC84xADlCq5k6ElP55lbbY36tQO+qDGdV6+TGN4bAL9FiQrDE6tQYYJdlv/sYE7iOOgnRBC9ljEdCIDNtQNGCfND/8JzatPg8KAy7xMRcLrGWu4xIMEysqNTz7rETfhdZjLQPssAht44KcoUJCH4/VuB+B9W1RhwA+M8Q3tqxzahIXzcgDM8OlmfkBlXo4FDVF3WUzjXLf9AMOg+66GupjQFtUpmRMkA8KXSV1HCso7e6nIIWtOnUoWeDDUfQPFFq4TNSlb6h2NuebaHdnW5nhxnJQ==", - "SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem", - "UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", - "MessageAttributes": { - "_datadog": { - "Type": "String", - "Value": '{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","StartTime":"2025-03-13T15:19:14.019Z","Name":"11623e4f-70ee-4330-8fbe-955152dea54c","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"SNS Publish","EnteredTime":"2025-03-13T15:19:14.061Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","serverless-version":"v1"}', - } - }, - }, - } - ] - } - ctx, source, event_source = extract_dd_trace_context(sns_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=1459500239678510857, - span_id=13193042003602978730, - sampling_priority=1, - meta={"_dd.p.tid": "fafc98885fd4647"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_sns(self): + """Test step function trace data extraction through SNS""" + sns_event = { + "Records": [ { - TraceHeader.TRACE_ID: "1459500239678510857", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=fafc98885fd4647", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) - - @with_trace_propagation_style("datadog") - def test_step_function_trace_data_sns_sqs(self): - lambda_ctx = get_mock_context() - sns_sqs_event = { - "Records": [ - { - "messageId": "9ec3339f-cd1a-43ba-9681-3e9113b430d3", - "receiptHandle": "AQEBJ5gIvqEWQt39NHPMAoK57cGgKtrgTtckWeWdDRi2FeucYr6pBhNjzXuUrmoHZMozX1WaoABtfQ5+kX5ucDBpA2Ci3Q07Z4MYvA6X0Sw13HCkiBnLrHPmH/F3rUBjvdRkIIKqA2ACX58MdkaYGNpqsHTJHB613wa8z4zurK0u7eUIXrr+e+gtsuPD39hiWlJo7cpBVv7y178rzMX8gPQTnRJv1cjhCHENtjWTSmfFC5N+BIQNIcjFsTTDRSovZlNIfAEuS+uowgzk0DUyoTJD5nFTL8lQHeXGRCUQe58/UY9OwRXEFVPGZOQR4OI9Wa4Kf/keFypTk9YwC9DhSeKvzZ0wBvejyl1n0ztT45+XYoWfi0mxGWM5b7r9wT36RDmjnM6vszH/d3fhZSRPASxWBQ==", - "body": '{\n "Type" : "Notification",\n "MessageId" : "1f3078d0-c792-5cf3-a130-189c3b846a3f",\n "TopicArn" : "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic",\n "Message" : "{}",\n "Timestamp" : "2025-03-13T15:29:26.348Z",\n "SignatureVersion" : "1",\n "Signature" : "mxOqAQ5o/isJrMS0PezHKRaA3g8Z/8YDbkToqhJub6I66LGtl+NYhyfTyllbgxvRP2XD2meKPRSgPI3nLyq8UHsWgyYwe3Tsv8QpRunCVE9Pebh+V1LGPWfjOiL0e+bnaj956QJD99560LJ6bzWP9QO584/zfOdcw6E5XQZfAI+pvEsf28Dy0WJO/lWTATRZDf8wGhmc7uKI1ZMsrOaNoUD8PXVqsI4yrJHxhzMb3SrC7YjI/PnNIbcn6ezwprbUdbZvyNAfJiE0k5IlppA089tMXC/ItgC7AgQhG9huPdKi5KdWGACK7gEwqmFwL+5T33sUXDaH2g58WhCs76pKEw==",\n "SigningCertURL" : "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem",\n "UnsubscribeURL" : "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:5f64545d-ae9a-4a5f-a7ee-798a0bd8519e",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"Execution\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"StartTime\\":\\"2025-03-13T15:29:26.144Z\\",\\"Name\\":\\"37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"RoleArn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j\\",\\"RedriveCount\\":0},\\"StateMachine\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine\\",\\"Name\\":\\"abhinav-inner-state-machine\\"},\\"State\\":{\\"Name\\":\\"SNS Publish\\",\\"EnteredTime\\":\\"2025-03-13T15:29:26.182Z\\",\\"RetryCount\\":0},\\"RootExecutionId\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"serverless-version\\":\\"v1\\"}"}\n }\n}', - "attributes": { - "ApproximateReceiveCount": "1", - "SentTimestamp": "1741879766424", - "SenderId": "AIDAIOA2GYWSHW4E2VXIO", - "ApproximateFirstReceiveTimestamp": "1741879766432", + "EventSource": "aws:sns", + "EventVersion": "1.0", + "EventSubscriptionArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "Sns": { + "Type": "Notification", + "MessageId": "7bc0c17d-bf88-5ff4-af7f-a131463a0d90", + "TopicArn": "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic", + "Subject": None, + "Message": "{}", + "Timestamp": "2025-03-13T15:19:14.245Z", + "SignatureVersion": "1", + "Signature": "r8RoYzq4uNcq0yj7sxcp8sTbFiDk8zqtocG7mJuE2MPVuR8O5eNg2ohofokUnC84xADlCq5k6ElP55lbbY36tQO+qDGdV6+TGN4bAL9FiQrDE6tQYYJdlv/sYE7iOOgnRBC9ljEdCIDNtQNGCfND/8JzatPg8KAy7xMRcLrGWu4xIMEysqNTz7rETfhdZjLQPssAht44KcoUJCH4/VuB+B9W1RhwA+M8Q3tqxzahIXzcgDM8OlmfkBlXo4FDVF3WUzjXLf9AMOg+66GupjQFtUpmRMkA8KXSV1HCso7e6nIIWtOnUoWeDDUfQPFFq4TNSlb6h2NuebaHdnW5nhxnJQ==", + "SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem", + "UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:f1653ba3-2ff7-4c8e-9381-45a7a62f9708", + "MessageAttributes": { + "_datadog": { + "Type": "String", + "Value": '{"Execution":{"Id":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","StartTime":"2025-03-13T15:19:14.019Z","Name":"11623e4f-70ee-4330-8fbe-955152dea54c","RoleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine","Name":"abhinav-inner-state-machine"},"State":{"Name":"SNS Publish","EnteredTime":"2025-03-13T15:19:14.061Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:11623e4f-70ee-4330-8fbe-955152dea54c","serverless-version":"v1"}', + } }, - "messageAttributes": {}, - "md5OfBody": "52af59de28507d7e67324b46c95337d8", - "eventSource": "aws:sqs", - "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", - "awsRegion": "sa-east-1", - } - ] - } - ctx, source, event_source = extract_dd_trace_context( - sns_sqs_event, lambda_ctx - ) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=5708348677301000120, - span_id=18223515719478572006, - sampling_priority=1, - meta={"_dd.p.tid": "45457f5f3fde3fa1"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), + }, + } + ] + } + self._test_step_function_trace_data_common( + sns_event, 1459500239678510857, 13193042003602978730, "fafc98885fd4647" + ) + + @with_trace_propagation_style("datadog") + def test_step_function_trace_data_sns_sqs(self): + """Test step function trace data extraction through SNS and SQS""" + sns_sqs_event = { + "Records": [ { - TraceHeader.TRACE_ID: "5708348677301000120", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - TraceHeader.TAGS: "_dd.p.tid=45457f5f3fde3fa1", - }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) + "messageId": "9ec3339f-cd1a-43ba-9681-3e9113b430d3", + "receiptHandle": "AQEBJ5gIvqEWQt39NHPMAoK57cGgKtrgTtckWeWdDRi2FeucYr6pBhNjzXuUrmoHZMozX1WaoABtfQ5+kX5ucDBpA2Ci3Q07Z4MYvA6X0Sw13HCkiBnLrHPmH/F3rUBjvdRkIIKqA2ACX58MdkaYGNpqsHTJHB613wa8z4zurK0u7eUIXrr+e+gtsuPD39hiWlJo7cpBVv7y178rzMX8gPQTnRJv1cjhCHENtjWTSmfFC5N+BIQNIcjFsTTDRSovZlNIfAEuS+uowgzk0DUyoTJD5nFTL8lQHeXGRCUQe58/UY9OwRXEFVPGZOQR4OI9Wa4Kf/keFypTk9YwC9DhSeKvzZ0wBvejyl1n0ztT45+XYoWfi0mxGWM5b7r9wT36RDmjnM6vszH/d3fhZSRPASxWBQ==", + "body": '{\n "Type" : "Notification",\n "MessageId" : "1f3078d0-c792-5cf3-a130-189c3b846a3f",\n "TopicArn" : "arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic",\n "Message" : "{}",\n "Timestamp" : "2025-03-13T15:29:26.348Z",\n "SignatureVersion" : "1",\n "Signature" : "mxOqAQ5o/isJrMS0PezHKRaA3g8Z/8YDbkToqhJub6I66LGtl+NYhyfTyllbgxvRP2XD2meKPRSgPI3nLyq8UHsWgyYwe3Tsv8QpRunCVE9Pebh+V1LGPWfjOiL0e+bnaj956QJD99560LJ6bzWP9QO584/zfOdcw6E5XQZfAI+pvEsf28Dy0WJO/lWTATRZDf8wGhmc7uKI1ZMsrOaNoUD8PXVqsI4yrJHxhzMb3SrC7YjI/PnNIbcn6ezwprbUdbZvyNAfJiE0k5IlppA089tMXC/ItgC7AgQhG9huPdKi5KdWGACK7gEwqmFwL+5T33sUXDaH2g58WhCs76pKEw==",\n "SigningCertURL" : "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem",\n "UnsubscribeURL" : "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:425362996713:logs-to-traces-dev-topic:5f64545d-ae9a-4a5f-a7ee-798a0bd8519e",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"Execution\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"StartTime\\":\\"2025-03-13T15:29:26.144Z\\",\\"Name\\":\\"37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"RoleArn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j\\",\\"RedriveCount\\":0},\\"StateMachine\\":{\\"Id\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine\\",\\"Name\\":\\"abhinav-inner-state-machine\\"},\\"State\\":{\\"Name\\":\\"SNS Publish\\",\\"EnteredTime\\":\\"2025-03-13T15:29:26.182Z\\",\\"RetryCount\\":0},\\"RootExecutionId\\":\\"arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:37ff72b8-0ee0-49e2-93c0-8a1764206a03\\",\\"serverless-version\\":\\"v1\\"}"}\n }\n}', + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1741879766424", + "SenderId": "AIDAIOA2GYWSHW4E2VXIO", + "ApproximateFirstReceiveTimestamp": "1741879766432", + }, + "messageAttributes": {}, + "md5OfBody": "52af59de28507d7e67324b46c95337d8", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:sa-east-1:425362996713:abhinav-q", + "awsRegion": "sa-east-1", + } + ] + } + self._test_step_function_trace_data_common( + sns_sqs_event, 5708348677301000120, 18223515719478572006, "45457f5f3fde3fa1" + ) class TestXRayContextConversion(unittest.TestCase): From 1b350b92c9a67c6af2af33fd0fc0383e81004bda Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Mon, 17 Mar 2025 09:43:43 -0400 Subject: [PATCH 15/15] removed event_type assertion --- tests/test_tracing.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 4097113b..0a961a62 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -36,14 +36,12 @@ _convert_xray_trace_id, _convert_xray_entity_id, _convert_xray_sampling, - InferredSpanInfo, create_service_mapping, determine_service_name, service_mapping as global_service_mapping, propagator, emit_telemetry_on_exception_outside_of_handler, ) -from datadog_lambda.trigger import EventTypes from tests.utils import get_mock_context @@ -630,10 +628,9 @@ def _test_step_function_trace_data_common( TraceHeader.TAGS: f"_dd.p.tid={expected_tid}", } - ctx, source, event_source = extract_dd_trace_context(event, lambda_ctx) + ctx, source, _ = extract_dd_trace_context(event, lambda_ctx) self.assertEqual(source, "event") - self.assertEqual(event_source.event_type, EventTypes.STEPFUNCTIONS) self.assertEqual(ctx, expected_context) self.assertEqual(get_dd_trace_context(), expected_headers)