From 41d49efa2034da5c7f1f4863336ced1538289546 Mon Sep 17 00:00:00 2001 From: Sergio Prada Date: Fri, 15 Jan 2021 14:26:30 -0500 Subject: [PATCH 1/5] add extractor and tests --- datadog_lambda/handler.py | 2 +- datadog_lambda/tracing.py | 27 ++++++++- datadog_lambda/wrapper.py | 14 ++++- tests/integration/serverless-plugin.yml | 1 + tests/integration/serverless.yml | 1 + tests/test_tracing.py | 81 +++++++++++++++++++++++++ tests/test_wrapper.py | 24 +++++++- 7 files changed, 145 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/handler.py b/datadog_lambda/handler.py index 09cc5e7d..5c90355b 100644 --- a/datadog_lambda/handler.py +++ b/datadog_lambda/handler.py @@ -16,6 +16,7 @@ class HandlerError(Exception): path = os.environ.get("DD_LAMBDA_HANDLER", None) + if path is None: raise HandlerError( "DD_LAMBDA_HANDLER is not defined. Can't use prebuilt datadog handler" @@ -24,7 +25,6 @@ class HandlerError(Exception): if len(parts) != 2: raise HandlerError("Value %s for DD_LAMBDA_HANDLER has invalid format." % path) - (mod_name, handler_name) = parts modified_mod_name = modify_module_name(mod_name) handler_module = import_module(modified_mod_name) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index d8ecca1d..5c91e5e7 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -150,7 +150,24 @@ def extract_context_from_sqs_event_or_context(event, lambda_context): return extract_context_from_lambda_context(lambda_context) -def extract_dd_trace_context(event, lambda_context): +def extract_context_custom_extractor(extractor, event, lambda_context): + """ + Extract Datadog trace context using a custom trace extractor function + """ + try: + ( + trace_id, + parent_id, + sampling_priority, + ) = extractor(event, lambda_context) + return trace_id, parent_id, sampling_priority + except Exception as e: + logger.debug("The trace extractor returned with error %s", e) + + return None, None, None + + +def extract_dd_trace_context(event, lambda_context, extractor=None): """ Extract Datadog trace context from the Lambda `event` object. @@ -163,7 +180,13 @@ def extract_dd_trace_context(event, lambda_context): """ global dd_trace_context - if "headers" in event: + if extractor is not None: + ( + trace_id, + parent_id, + sampling_priority, + ) = extract_context_custom_extractor(extractor, event, lambda_context) + elif "headers" in event: ( trace_id, parent_id, diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 34ce6168..430fd7bb 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -6,6 +6,7 @@ import os import logging import traceback +from importlib import import_module from datadog_lambda.extension import should_use_extension, flush_extension from datadog_lambda.cold_start import set_cold_start, is_cold_start @@ -14,6 +15,7 @@ submit_invocations_metric, submit_errors_metric, ) +from datadog_lambda.module_name import modify_module_name from datadog_lambda.patch import patch_all from datadog_lambda.tracing import ( extract_dd_trace_context, @@ -88,6 +90,16 @@ def __init__(self, func): os.environ.get("DD_MERGE_XRAY_TRACES", "false").lower() == "true" ) self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function") + self.extractor_env = os.environ.get("DD_TRACE_EXTRACTOR", None) + self.trace_extractor = None + + if self.extractor_env: + extractor_parts = self.extractor_env.rsplit(".", 1) + if len(extractor_parts) == 2: + (mod_name, extractor_name) = extractor_parts + modified_extractor_name = modify_module_name(mod_name) + extractor_module = import_module(modified_extractor_name) + self.trace_extractor = getattr(extractor_module, extractor_name) # Inject trace correlation ids to logs if self.logs_injection: @@ -119,7 +131,7 @@ def _before(self, event, context): set_cold_start() submit_invocations_metric(context) # Extract Datadog trace context from incoming requests - dd_context = extract_dd_trace_context(event, context) + dd_context = extract_dd_trace_context(event, context, extractor=self.trace_extractor) self.span = None if dd_tracing_enabled: diff --git a/tests/integration/serverless-plugin.yml b/tests/integration/serverless-plugin.yml index 2a69fa05..386429e1 100644 --- a/tests/integration/serverless-plugin.yml +++ b/tests/integration/serverless-plugin.yml @@ -6,6 +6,7 @@ provider: DD_INTEGRATION_TEST: true DD_API_KEY: ${env:DD_API_KEY} WITH_PLUGIN: true + lambdaHashingVersion: 20201221 layers: python27: diff --git a/tests/integration/serverless.yml b/tests/integration/serverless.yml index e7badc55..cc605397 100644 --- a/tests/integration/serverless.yml +++ b/tests/integration/serverless.yml @@ -8,6 +8,7 @@ provider: environment: DD_INTEGRATION_TEST: true DD_API_KEY: ${env:DD_API_KEY} + lambdaHashingVersion: 20201221 layers: python27: diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0c8f8154..efae0b2b 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -150,6 +150,87 @@ def test_with_complete_datadog_trace_headers(self): XraySubsegment.NAMESPACE, ) + def test_with_extractor_function(self): + def extractor_foo(event, context): + foo = event.get("foo", {}) + lowercase_foo = {k.lower(): v for k, v in foo.items()} + + trace_id = lowercase_foo.get(TraceHeader.TRACE_ID) + parent_id = lowercase_foo.get(TraceHeader.PARENT_ID) + sampling_priority = lowercase_foo.get(TraceHeader.SAMPLING_PRIORITY) + return trace_id, parent_id, sampling_priority + + lambda_ctx = get_mock_context() + ctx = extract_dd_trace_context( + { + "foo": { + TraceHeader.TRACE_ID: "123", + TraceHeader.PARENT_ID: "321", + TraceHeader.SAMPLING_PRIORITY: "1", + } + }, + lambda_ctx, + extractor=extractor_foo, + ) + self.assertDictEqual( + ctx, + { + "trace-id": "123", + "parent-id": "321", + "sampling-priority": "1", + "source": "event", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "123", + TraceHeader.PARENT_ID: "65535", + TraceHeader.SAMPLING_PRIORITY: "1", + }, + ) + self.mock_xray_recorder.begin_subsegment.assert_called() + self.mock_xray_recorder.end_subsegment.assert_called() + self.mock_current_subsegment.put_metadata.assert_called_with( + XraySubsegment.KEY, + {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, + XraySubsegment.NAMESPACE, + ) + + def test_graceful_fail_of_extractor_function(self): + def extractor_raiser(event, context): + raise Exception("kreator") + + lambda_ctx = get_mock_context() + ctx = extract_dd_trace_context( + { + "foo": { + TraceHeader.TRACE_ID: "123", + TraceHeader.PARENT_ID: "321", + TraceHeader.SAMPLING_PRIORITY: "1", + } + }, + lambda_ctx, + extractor=extractor_raiser, + ) + self.assertDictEqual( + ctx, + { + "trace-id": "4369", + "parent-id": "65535", + "sampling-priority": "2", + "source": "xray", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "4369", + TraceHeader.PARENT_ID: "65535", + TraceHeader.SAMPLING_PRIORITY: "2", + }, + ) + def test_with_sqs_distributed_datadog_trace_data(self): lambda_ctx = get_mock_context() sqs_event = { diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index f581122d..711c04f9 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -99,7 +99,29 @@ def lambda_handler(event, context): ) self.mock_wrapper_lambda_stats.flush.assert_called() self.mock_extract_dd_trace_context.assert_called_with( - lambda_event, lambda_context + lambda_event, lambda_context, extractor=None + ) + self.mock_set_correlation_ids.assert_called() + self.mock_inject_correlation_ids.assert_called() + self.mock_patch_all.assert_called() + + def test_datadog_lambda_wrapper_with_extractor(self): + @datadog_lambda_wrapper + def lambda_handler(event, context): + lambda_metric("test.metric", 100) + + lambda_event = {} + + lambda_context = get_mock_context() + + lambda_handler(lambda_event, lambda_context) + + self.mock_metric_lambda_stats.distribution.assert_has_calls( + [call("test.metric", 100, timestamp=None, tags=ANY)] + ) + self.mock_wrapper_lambda_stats.flush.assert_called() + self.mock_extract_dd_trace_context.assert_called_with( + lambda_event, lambda_context, extractor=None ) self.mock_set_correlation_ids.assert_called() self.mock_inject_correlation_ids.assert_called() From c7d34ea675e41288bbd54649b2fe1ba2f97bbec8 Mon Sep 17 00:00:00 2001 From: Sergio Prada Date: Fri, 15 Jan 2021 14:30:22 -0500 Subject: [PATCH 2/5] lint and other fixes --- datadog_lambda/handler.py | 2 +- datadog_lambda/tracing.py | 14 ++++---------- tests/test_wrapper.py | 22 ---------------------- 3 files changed, 5 insertions(+), 33 deletions(-) diff --git a/datadog_lambda/handler.py b/datadog_lambda/handler.py index 5c90355b..09cc5e7d 100644 --- a/datadog_lambda/handler.py +++ b/datadog_lambda/handler.py @@ -16,7 +16,6 @@ class HandlerError(Exception): path = os.environ.get("DD_LAMBDA_HANDLER", None) - if path is None: raise HandlerError( "DD_LAMBDA_HANDLER is not defined. Can't use prebuilt datadog handler" @@ -25,6 +24,7 @@ class HandlerError(Exception): if len(parts) != 2: raise HandlerError("Value %s for DD_LAMBDA_HANDLER has invalid format." % path) + (mod_name, handler_name) = parts modified_mod_name = modify_module_name(mod_name) handler_module = import_module(modified_mod_name) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 5c91e5e7..09f543c3 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -155,11 +155,7 @@ def extract_context_custom_extractor(extractor, event, lambda_context): Extract Datadog trace context using a custom trace extractor function """ try: - ( - trace_id, - parent_id, - sampling_priority, - ) = extractor(event, lambda_context) + (trace_id, parent_id, sampling_priority,) = extractor(event, lambda_context) return trace_id, parent_id, sampling_priority except Exception as e: logger.debug("The trace extractor returned with error %s", e) @@ -181,11 +177,9 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): global dd_trace_context if extractor is not None: - ( - trace_id, - parent_id, - sampling_priority, - ) = extract_context_custom_extractor(extractor, event, lambda_context) + (trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor( + extractor, event, lambda_context + ) elif "headers" in event: ( trace_id, diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 711c04f9..9788f292 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -105,28 +105,6 @@ def lambda_handler(event, context): self.mock_inject_correlation_ids.assert_called() self.mock_patch_all.assert_called() - def test_datadog_lambda_wrapper_with_extractor(self): - @datadog_lambda_wrapper - def lambda_handler(event, context): - lambda_metric("test.metric", 100) - - lambda_event = {} - - lambda_context = get_mock_context() - - lambda_handler(lambda_event, lambda_context) - - self.mock_metric_lambda_stats.distribution.assert_has_calls( - [call("test.metric", 100, timestamp=None, tags=ANY)] - ) - self.mock_wrapper_lambda_stats.flush.assert_called() - self.mock_extract_dd_trace_context.assert_called_with( - lambda_event, lambda_context, extractor=None - ) - self.mock_set_correlation_ids.assert_called() - self.mock_inject_correlation_ids.assert_called() - self.mock_patch_all.assert_called() - def test_datadog_lambda_wrapper_flush_to_log(self): os.environ["DD_FLUSH_TO_LOG"] = "True" From 342cc21a5a1239b338c9fea237cefab1474ba564 Mon Sep 17 00:00:00 2001 From: Sergio Prada Date: Fri, 15 Jan 2021 15:01:58 -0500 Subject: [PATCH 3/5] fix lint --- datadog_lambda/wrapper.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 430fd7bb..8dc58277 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -131,7 +131,9 @@ def _before(self, event, context): set_cold_start() submit_invocations_metric(context) # Extract Datadog trace context from incoming requests - dd_context = extract_dd_trace_context(event, context, extractor=self.trace_extractor) + dd_context = extract_dd_trace_context( + event, context, extractor=self.trace_extractor + ) self.span = None if dd_tracing_enabled: From 15a113225d8c2d51cc7af9f168a1612562e4222e Mon Sep 17 00:00:00 2001 From: Tian Chu Date: Wed, 17 Feb 2021 15:56:17 -0500 Subject: [PATCH 4/5] Update tests --- tests/test_tracing.py | 38 ++++++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 1ab7b0ff..dd257744 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -73,7 +73,8 @@ def test_without_datadog_trace_headers(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, + {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -93,7 +94,8 @@ def test_with_incomplete_datadog_trace_headers(self): ) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, + {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -118,7 +120,8 @@ def test_with_complete_datadog_trace_headers(self): ) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, + ctx, + {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -147,7 +150,7 @@ def extractor_foo(event, context): return trace_id, parent_id, sampling_priority lambda_ctx = get_mock_context() - ctx = extract_dd_trace_context( + ctx, ctx_source = extract_dd_trace_context( { "foo": { TraceHeader.TRACE_ID: "123", @@ -158,13 +161,13 @@ def extractor_foo(event, context): lambda_ctx, extractor=extractor_foo, ) + self.assertEquals(ctx_source, "event") self.assertDictEqual( ctx, { "trace-id": "123", "parent-id": "321", "sampling-priority": "1", - "source": "event", }, ) self.assertDictEqual( @@ -175,20 +178,13 @@ def extractor_foo(event, context): TraceHeader.SAMPLING_PRIORITY: "1", }, ) - self.mock_xray_recorder.begin_subsegment.assert_called() - self.mock_xray_recorder.end_subsegment.assert_called() - self.mock_current_subsegment.put_metadata.assert_called_with( - XraySubsegment.KEY, - {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, - XraySubsegment.NAMESPACE, - ) def test_graceful_fail_of_extractor_function(self): def extractor_raiser(event, context): raise Exception("kreator") lambda_ctx = get_mock_context() - ctx = extract_dd_trace_context( + ctx, ctx_source = extract_dd_trace_context( { "foo": { TraceHeader.TRACE_ID: "123", @@ -199,13 +195,13 @@ def extractor_raiser(event, context): lambda_ctx, extractor=extractor_raiser, ) + self.assertEquals(ctx_source, "xray") self.assertDictEqual( ctx, { "trace-id": "4369", "parent-id": "65535", "sampling-priority": "2", - "source": "xray", }, ) self.assertDictEqual( @@ -252,7 +248,12 @@ def test_with_sqs_distributed_datadog_trace_data(self): ctx, source = extract_dd_trace_context(sqs_event, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",}, + ctx, + { + "trace-id": "123", + "parent-id": "321", + "sampling-priority": "1", + }, ) self.assertDictEqual( get_dd_trace_context(), @@ -284,7 +285,12 @@ def test_with_client_context_datadog_trace_data(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "666", "parent-id": "777", "sampling-priority": "1",}, + ctx, + { + "trace-id": "666", + "parent-id": "777", + "sampling-priority": "1", + }, ) self.assertDictEqual( get_dd_trace_context(), From b0fd4a05f74bb52f622bf3c758b37c726c078f72 Mon Sep 17 00:00:00 2001 From: Tian Chu Date: Wed, 17 Feb 2021 17:26:53 -0500 Subject: [PATCH 5/5] Fix format --- tests/test_tracing.py | 37 +++++++------------------------------ 1 file changed, 7 insertions(+), 30 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index dd257744..bf8ebba1 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -73,8 +73,7 @@ def test_without_datadog_trace_headers(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, - {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -94,8 +93,7 @@ def test_with_incomplete_datadog_trace_headers(self): ) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, - {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -120,8 +118,7 @@ def test_with_complete_datadog_trace_headers(self): ) self.assertEqual(source, "event") self.assertDictEqual( - ctx, - {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, + ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -163,12 +160,7 @@ def extractor_foo(event, context): ) self.assertEquals(ctx_source, "event") self.assertDictEqual( - ctx, - { - "trace-id": "123", - "parent-id": "321", - "sampling-priority": "1", - }, + ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",}, ) self.assertDictEqual( get_dd_trace_context(), @@ -197,12 +189,7 @@ def extractor_raiser(event, context): ) self.assertEquals(ctx_source, "xray") self.assertDictEqual( - ctx, - { - "trace-id": "4369", - "parent-id": "65535", - "sampling-priority": "2", - }, + ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2",}, ) self.assertDictEqual( get_dd_trace_context(), @@ -248,12 +235,7 @@ def test_with_sqs_distributed_datadog_trace_data(self): ctx, source = extract_dd_trace_context(sqs_event, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, - { - "trace-id": "123", - "parent-id": "321", - "sampling-priority": "1", - }, + ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",}, ) self.assertDictEqual( get_dd_trace_context(), @@ -285,12 +267,7 @@ def test_with_client_context_datadog_trace_data(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, - { - "trace-id": "666", - "parent-id": "777", - "sampling-priority": "1", - }, + ctx, {"trace-id": "666", "parent-id": "777", "sampling-priority": "1",}, ) self.assertDictEqual( get_dd_trace_context(),