From 3b3d981a4fc67e81800842dd1031527fce2d135f Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 5 Nov 2024 12:20:16 -0500 Subject: [PATCH 1/2] check if event is a dict before getting --- datadog_lambda/tracing.py | 3 +++ tests/test_tracing.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0aded4de..cd0dd1f3 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -411,6 +411,9 @@ def is_legacy_lambda_step_function(event): """ Check if the event is a step function that called a legacy lambda """ + if not isinstance(event, dict): + return False + event = event.get("Payload", {}) return "Execution" in event and "StateMachine" in event and "State" in event diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0fb2ee31..22ac7049 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -678,6 +678,9 @@ def test_is_legacy_lambda_step_function(self): } self.assertFalse(is_legacy_lambda_step_function(sf_event)) + other_event = ["foo", "bar"] + self.assertFalse(is_legacy_lambda_step_function(other_event)) + class TestXRayContextConversion(unittest.TestCase): def test_convert_xray_trace_id(self): From e2d265ab13a59cbfa437f89dbe58a58994bbee83 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 5 Nov 2024 12:50:18 -0500 Subject: [PATCH 2/2] early exit in is_legacy_lambda_step_function --- datadog_lambda/tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index cd0dd1f3..dfb08dd2 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -411,10 +411,10 @@ def is_legacy_lambda_step_function(event): """ Check if the event is a step function that called a legacy lambda """ - if not isinstance(event, dict): + if not isinstance(event, dict) or "Payload" not in event: return False - event = event.get("Payload", {}) + event = event.get("Payload") return "Execution" in event and "StateMachine" in event and "State" in event