From 368c9b81401831e466be1cd924325a153ddad518 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 8 Aug 2024 15:31:40 -0400 Subject: [PATCH 1/8] pull payload out of legacy event --- datadog_lambda/tracing.py | 11 +++++++++++ datadog_lambda/wrapper.py | 3 +++ 2 files changed, 14 insertions(+) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 64a19ccd..ec005b5c 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -407,6 +407,17 @@ def extract_context_from_step_functions(event, lambda_context): return extract_context_from_lambda_context(lambda_context) +def is_legacy_lambda_step_function(event): + """ + Check if the event is a step function that called a legacy lambda + """ + if "Payload" in event: + event = event["Payload"] + return "Execution" in event and "StateMachine" in event and "State" in event + + 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/wrapper.py b/datadog_lambda/wrapper.py index ed3d92b4..e97c0ebe 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -44,6 +44,7 @@ is_authorizer_response, tracer, propagator, + is_legacy_lambda_step_function, ) from datadog_lambda.trigger import ( extract_trigger_tags, @@ -277,6 +278,8 @@ 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( From ebe0b59ec7afda3d9dba100eaa34a1ba8f9ec73e Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 8 Aug 2024 16:48:28 -0400 Subject: [PATCH 2/8] add unit test for legacy lambda --- tests/test_tracing.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index b94e968f..89ea77ae 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -647,6 +647,45 @@ def test_step_function_trace_data(self): expected_context, ) + @with_trace_propagation_style("datadog") + def test_step_function_legacy_lambda_trace_data(self): + lambda_ctx = get_mock_context() + sqs_event = { + "Payload": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + }, + } + } + ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) + self.assertEqual(source, "event") + expected_context = Context( + trace_id=3675572987363469717, + span_id=6880978411788117524, + sampling_priority=1, + meta={"_dd.p.tid": "e987c84b36b11ab"}, + ) + self.assertEqual(ctx, expected_context) + self.assertEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "3675572987363469717", + TraceHeader.PARENT_ID: "10713633173203262661", + TraceHeader.SAMPLING_PRIORITY: "1", + "x-datadog-tags": "_dd.p.tid=e987c84b36b11ab", + }, + ) + 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 1fb832ad12e27e3cba9400e06b9840be05a14d94 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 8 Aug 2024 16:54:20 -0400 Subject: [PATCH 3/8] fix unit test --- tests/test_tracing.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 89ea77ae..1ff59dc3 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -40,6 +40,7 @@ service_mapping as global_service_mapping, propagator, emit_telemetry_on_exception_outside_of_handler, + is_legacy_lambda_step_function, ) from datadog_lambda.trigger import EventTypes @@ -662,6 +663,10 @@ def test_step_function_legacy_lambda_trace_data(self): }, } } + + if is_legacy_lambda_step_function(sqs_event): + sqs_event = sqs_event["Payload"] + ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) self.assertEqual(source, "event") expected_context = Context( From a31bead191c4c7e00c4f29426d211e3aefd8e6de Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 9 Aug 2024 14:22:00 -0400 Subject: [PATCH 4/8] integration test experiment --- datadog_lambda/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index fd8afb3e..f3790d3b 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -7,7 +7,7 @@ class SamplingPriority(object): - USER_REJECT = -1 + USER_REJECT = -2 AUTO_REJECT = 0 AUTO_KEEP = 1 USER_KEEP = 2 From 4bddfb0faa44a4061bc85a558cf83b39a1c3dd5a Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 9 Aug 2024 14:52:56 -0400 Subject: [PATCH 5/8] undo integration test experiment --- datadog_lambda/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index f3790d3b..fd8afb3e 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -7,7 +7,7 @@ class SamplingPriority(object): - USER_REJECT = -2 + USER_REJECT = -1 AUTO_REJECT = 0 AUTO_KEEP = 1 USER_KEEP = 2 From ebdfe1344dbac89e536d113e9aa945dc2def7b82 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Mon, 12 Aug 2024 09:03:19 -0400 Subject: [PATCH 6/8] refactor check --- datadog_lambda/tracing.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index ec005b5c..2a696046 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -411,11 +411,8 @@ def is_legacy_lambda_step_function(event): """ Check if the event is a step function that called a legacy lambda """ - if "Payload" in event: - event = event["Payload"] - return "Execution" in event and "StateMachine" in event and "State" in event - - return False + event = event.get("Payload", {}) + return "Execution" in event and "StateMachine" in event and "State" in event def extract_context_custom_extractor(extractor, event, lambda_context): From dae03a44a176c62c0751600bd4b120f734a12939 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Mon, 12 Aug 2024 10:14:34 -0400 Subject: [PATCH 7/8] added new unit test --- tests/test_tracing.py | 40 ++++++++++++---------------------------- 1 file changed, 12 insertions(+), 28 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 1ff59dc3..36bf2d21 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -648,9 +648,7 @@ def test_step_function_trace_data(self): expected_context, ) - @with_trace_propagation_style("datadog") - def test_step_function_legacy_lambda_trace_data(self): - lambda_ctx = get_mock_context() + def test_is_legacy_lambda_step_function(self): sqs_event = { "Payload": { "Execution": { @@ -663,33 +661,19 @@ def test_step_function_legacy_lambda_trace_data(self): }, } } + self.assertTrue(is_legacy_lambda_step_function(sqs_event)) - if is_legacy_lambda_step_function(sqs_event): - sqs_event = sqs_event["Payload"] - - ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) - self.assertEqual(source, "event") - expected_context = Context( - trace_id=3675572987363469717, - span_id=6880978411788117524, - sampling_priority=1, - meta={"_dd.p.tid": "e987c84b36b11ab"}, - ) - self.assertEqual(ctx, expected_context) - self.assertEqual( - get_dd_trace_context(), - { - TraceHeader.TRACE_ID: "3675572987363469717", - TraceHeader.PARENT_ID: "10713633173203262661", - TraceHeader.SAMPLING_PRIORITY: "1", - "x-datadog-tags": "_dd.p.tid=e987c84b36b11ab", + sqs_event = { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", }, - ) - create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) - self.mock_send_segment.assert_called_with( - XraySubsegment.TRACE_KEY, - expected_context, - ) + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + }, + } + self.assertFalse(is_legacy_lambda_step_function(sqs_event)) class TestXRayContextConversion(unittest.TestCase): From 11937c3066b690b09247ac28694177a02d64241d Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Mon, 12 Aug 2024 11:28:13 -0400 Subject: [PATCH 8/8] update var naming in test --- tests/test_tracing.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 36bf2d21..b21b9337 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -649,7 +649,7 @@ def test_step_function_trace_data(self): ) def test_is_legacy_lambda_step_function(self): - sqs_event = { + sf_event = { "Payload": { "Execution": { "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", @@ -661,9 +661,9 @@ def test_is_legacy_lambda_step_function(self): }, } } - self.assertTrue(is_legacy_lambda_step_function(sqs_event)) + self.assertTrue(is_legacy_lambda_step_function(sf_event)) - sqs_event = { + sf_event = { "Execution": { "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", }, @@ -673,7 +673,7 @@ def test_is_legacy_lambda_step_function(self): "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", }, } - self.assertFalse(is_legacy_lambda_step_function(sqs_event)) + self.assertFalse(is_legacy_lambda_step_function(sf_event)) class TestXRayContextConversion(unittest.TestCase):