From 3006cddda660d3f494b4f3e7328d326853ee53c6 Mon Sep 17 00:00:00 2001 From: AJ Stuyvenberg Date: Mon, 27 Jun 2022 15:59:21 -0400 Subject: [PATCH 1/2] fix: type check event before treating it as a dict --- datadog_lambda/tracing.py | 2 +- tests/test_tracing.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b5324b1e..3cadc538 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -332,7 +332,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): parent_id, sampling_priority, ) = extract_context_custom_extractor(extractor, event, lambda_context) - elif "headers" in event: + elif isinstance(event, (list, dict)) and "headers" in event: ( trace_id, parent_id, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index be52697e..f69bc6c7 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -101,6 +101,28 @@ def test_without_datadog_trace_headers(self): {}, ) + def test_with_non_object_event(self): + lambda_ctx = get_mock_context() + ctx, source = extract_dd_trace_context(b'', lambda_ctx) + self.assertEqual(source, "xray") + self.assertDictEqual( + ctx, + { + "trace-id": fake_xray_header_value_root_decimal, + "parent-id": fake_xray_header_value_parent_decimal, + "sampling-priority": "2", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: fake_xray_header_value_root_decimal, + TraceHeader.PARENT_ID: fake_xray_header_value_parent_decimal, + TraceHeader.SAMPLING_PRIORITY: "2", + }, + {}, + ) + def test_with_incomplete_datadog_trace_headers(self): lambda_ctx = get_mock_context() ctx, source = extract_dd_trace_context( From 74b106abd645a640ef92f9c858f2b194ea8f93f8 Mon Sep 17 00:00:00 2001 From: AJ Stuyvenberg Date: Mon, 27 Jun 2022 16:11:44 -0400 Subject: [PATCH 2/2] fix: Use doublequotes for binary string --- 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 f69bc6c7..e44e9b0c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -103,7 +103,7 @@ def test_without_datadog_trace_headers(self): def test_with_non_object_event(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context(b'', lambda_ctx) + ctx, source = extract_dd_trace_context(b"", lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( ctx,