diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 64a19ccd..2a696046 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -407,6 +407,14 @@ 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 + """ + 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): """ 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( diff --git a/tests/test_tracing.py b/tests/test_tracing.py index b94e968f..b21b9337 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 @@ -647,6 +648,33 @@ def test_step_function_trace_data(self): expected_context, ) + def test_is_legacy_lambda_step_function(self): + sf_event = { + "Payload": { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + }, + } + } + self.assertTrue(is_legacy_lambda_step_function(sf_event)) + + sf_event = { + "Execution": { + "Id": "665c417c-1237-4742-aaca-8b3becbb9e75", + }, + "StateMachine": {}, + "State": { + "Name": "my-awesome-state", + "EnteredTime": "Mon Nov 13 12:43:33 PST 2023", + }, + } + self.assertFalse(is_legacy_lambda_step_function(sf_event)) + class TestXRayContextConversion(unittest.TestCase): def test_convert_xray_trace_id(self):