From 64eb98f6f625828c055e322f837bb39efd494c6c Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 10:14:03 -0400 Subject: [PATCH 1/7] Check for the new direct invocation trace format and use that --- datadog_lambda/tracing.py | 42 ++++++++++++++++++++++----------------- tests/test_tracing.py | 37 +++++++++++++++++++++++++++++++++- 2 files changed, 60 insertions(+), 19 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index ae48a161..909e79c7 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -6,6 +6,7 @@ import logging import os import json +from typing import Optional from aws_xray_sdk.core import xray_recorder from aws_xray_sdk.core.lambda_launcher import LambdaContext @@ -129,16 +130,27 @@ def extract_context_from_lambda_context(lambda_context): dd_trace libraries inject this trace context on synchronous invocations """ client_context = lambda_context.client_context - + trace_id = None # type: Optional[str] + parent_id = None # type: Optional[str] + sampling_priority = None # type: Optional[str] if client_context and client_context.custom: - dd_data = client_context.custom.get("_datadog", {}) - trace_id = dd_data.get(TraceHeader.TRACE_ID) - parent_id = dd_data.get(TraceHeader.PARENT_ID) - sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) + if "_datadog" in client_context.custom: + # Legacy trace propagation dict + dd_data = client_context.custom.get("_datadog", {}) + trace_id = dd_data.get(TraceHeader.TRACE_ID) + parent_id = dd_data.get(TraceHeader.PARENT_ID) + sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) + elif ( + TraceHeader.TRACE_ID in client_context.custom + and TraceHeader.PARENT_ID in client_context.custom + and TraceHeader.SAMPLING_PRIORITY in client_context.custom + ): + # New trace propagation keys + trace_id = client_context.custom.get(TraceHeader.TRACE_ID) + parent_id = client_context.custom.get(TraceHeader.PARENT_ID) + sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY) - return trace_id, parent_id, sampling_priority - - return None, None, None + return trace_id, parent_id, sampling_priority def extract_context_from_http_event_or_context(event, lambda_context): @@ -186,11 +198,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) @@ -209,11 +217,9 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): trace_context_source = None 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_tracing.py b/tests/test_tracing.py index 431f339f..64b42e18 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self): XraySubsegment.NAMESPACE, ) - def test_with_client_context_datadog_trace_data(self): + def test_with_legacy_client_context_datadog_trace_data(self): lambda_ctx = get_mock_context( custom={ "_datadog": { @@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self): XraySubsegment.NAMESPACE, ) + def test_with_new_client_context_datadog_trace_data(self): + lambda_ctx = get_mock_context( + custom={ + TraceHeader.TRACE_ID: "666", + TraceHeader.PARENT_ID: "777", + TraceHeader.SAMPLING_PRIORITY: "1", + } + ) + ctx, source = extract_dd_trace_context({}, lambda_ctx) + self.assertEqual(source, "event") + self.assertDictEqual( + ctx, + { + "trace-id": "666", + "parent-id": "777", + "sampling-priority": "1", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: "666", + TraceHeader.PARENT_ID: "65535", + TraceHeader.SAMPLING_PRIORITY: "1", + }, + ) + create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY) + 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.TRACE_KEY, + {"trace-id": "666", "parent-id": "777", "sampling-priority": "1"}, + XraySubsegment.NAMESPACE, + ) + def test_with_complete_datadog_trace_headers_with_mixed_casing(self): lambda_ctx = get_mock_context() extract_dd_trace_context( From 6719312f857ac20e7f6ef9e1f7073a1d1fa949d9 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 12:04:45 -0400 Subject: [PATCH 2/7] run Black --- datadog_lambda/metric.py | 4 +--- tests/test_tracing.py | 44 ++++++++-------------------------------- 2 files changed, 9 insertions(+), 39 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 62100464..49e83958 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -237,9 +237,7 @@ def decrypt_kms_api_key(kms_client, ciphertext): try: plaintext = kms_client.decrypt( CiphertextBlob=decoded_bytes, - EncryptionContext={ - KMS_ENCRYPTION_CONTEXT_KEY: function_name, - }, + EncryptionContext={KMS_ENCRYPTION_CONTEXT_KEY: function_name,}, )["Plaintext"].decode("utf-8") except ClientError: logger.debug( diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 64b42e18..0e05ade4 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_legacy_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(), @@ -320,12 +297,7 @@ def test_with_new_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 535988a5452e9599b62108242d7ce55d49f1d5f0 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 12:19:42 -0400 Subject: [PATCH 3/7] noqa Optional --- datadog_lambda/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 909e79c7..881dec2a 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -6,7 +6,7 @@ import logging import os import json -from typing import Optional +from typing import Optional # noqa: F401 from aws_xray_sdk.core import xray_recorder from aws_xray_sdk.core.lambda_launcher import LambdaContext From 275903dae0080c138e8c11ed29debeed9e4a61fb Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 13:55:45 -0400 Subject: [PATCH 4/7] Do what the lint build step is telling me to do --- datadog_lambda/metric.py | 4 +++- datadog_lambda/tracing.py | 12 ++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 49e83958..62100464 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -237,7 +237,9 @@ def decrypt_kms_api_key(kms_client, ciphertext): try: plaintext = kms_client.decrypt( CiphertextBlob=decoded_bytes, - EncryptionContext={KMS_ENCRYPTION_CONTEXT_KEY: function_name,}, + EncryptionContext={ + KMS_ENCRYPTION_CONTEXT_KEY: function_name, + }, )["Plaintext"].decode("utf-8") except ClientError: logger.debug( diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 881dec2a..0a03586d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -198,7 +198,11 @@ 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) @@ -217,7 +221,11 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): trace_context_source = None if extractor is not None: - (trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor( + ( + trace_id, + parent_id, + sampling_priority, + ) = extract_context_custom_extractor( extractor, event, lambda_context ) elif "headers" in event: From 72aaa2da47244b862a77cd63db0a913545bc0cb3 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 14:01:54 -0400 Subject: [PATCH 5/7] installed the correct version of Black --- datadog_lambda/tracing.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0a03586d..30aeb5c2 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -225,9 +225,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): trace_id, parent_id, sampling_priority, - ) = extract_context_custom_extractor( - extractor, event, lambda_context - ) + ) = extract_context_custom_extractor(extractor, event, lambda_context) elif "headers" in event: ( trace_id, From 25e3efd38f95c8e7119edd86c5765bc11bc53ddb Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 16 Jun 2021 14:08:34 -0400 Subject: [PATCH 6/7] Black against the tests now --- tests/test_tracing.py | 44 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0e05ade4..64b42e18 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(), @@ -160,7 +163,12 @@ 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(), @@ -189,7 +197,12 @@ 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(), @@ -235,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(), @@ -267,7 +285,12 @@ def test_with_legacy_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(), @@ -297,7 +320,12 @@ def test_with_new_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 a10a1fe812de295233b0abff945cecfc6229cd7a Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Thu, 17 Jun 2021 12:33:03 -0400 Subject: [PATCH 7/7] Remove type hints --- datadog_lambda/tracing.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 30aeb5c2..49e1d05f 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -6,7 +6,6 @@ import logging import os import json -from typing import Optional # noqa: F401 from aws_xray_sdk.core import xray_recorder from aws_xray_sdk.core.lambda_launcher import LambdaContext @@ -130,9 +129,9 @@ def extract_context_from_lambda_context(lambda_context): dd_trace libraries inject this trace context on synchronous invocations """ client_context = lambda_context.client_context - trace_id = None # type: Optional[str] - parent_id = None # type: Optional[str] - sampling_priority = None # type: Optional[str] + trace_id = None + parent_id = None + sampling_priority = None if client_context and client_context.custom: if "_datadog" in client_context.custom: # Legacy trace propagation dict