From 49b61aeeee6451c8fce08fac824799476b4bb27c Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 2 Dec 2022 16:48:03 -0500 Subject: [PATCH] use context.aws_request_id instead of http payload's requestId in the spans' tags --- datadog_lambda/tracing.py | 12 ++++++------ tests/test_tracing.py | 22 +++++++--------------- 2 files changed, 13 insertions(+), 21 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 62c64c61..ead72767 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -584,7 +584,7 @@ def create_inferred_span( ): logger.debug("API Gateway event detected. Inferring a span") return create_inferred_span_from_api_gateway_event( - event, decode_authorizer_context + event, context, decode_authorizer_context ) elif event_source.equals(EventTypes.LAMBDA_FUNCTION_URL): logger.debug("Function URL event detected. Inferring a span") @@ -601,7 +601,7 @@ def create_inferred_span( ): logger.debug("API Gateway Websocket event detected. Inferring a span") return create_inferred_span_from_api_gateway_websocket_event( - event, decode_authorizer_context + event, context, decode_authorizer_context ) elif event_source.equals(EventTypes.SQS): logger.debug("SQS event detected. Inferring a span") @@ -732,7 +732,7 @@ def process_injected_data(event, request_time_epoch_ms, args, tags): def create_inferred_span_from_api_gateway_websocket_event( - event, decode_authorizer_context: bool = True + event, context, decode_authorizer_context: bool = True ): request_context = event.get("requestContext") domain = request_context.get("domainName") @@ -745,7 +745,7 @@ def create_inferred_span_from_api_gateway_websocket_event( "apiid": request_context.get("apiId"), "apiname": request_context.get("apiId"), "stage": request_context.get("stage"), - "request_id": request_context.get("requestId"), + "request_id": context.aws_request_id, "connection_id": request_context.get("connectionId"), "event_type": request_context.get("eventType"), "message_direction": request_context.get("messageDirection"), @@ -781,7 +781,7 @@ def create_inferred_span_from_api_gateway_websocket_event( def create_inferred_span_from_api_gateway_event( - event, decode_authorizer_context: bool = True + event, context, decode_authorizer_context: bool = True ): request_context = event.get("requestContext") domain = request_context.get("domainName", "") @@ -797,7 +797,7 @@ def create_inferred_span_from_api_gateway_event( "apiid": request_context.get("apiId"), "apiname": request_context.get("apiId"), "stage": request_context.get("stage"), - "request_id": request_context.get("requestId"), + "request_id": context.aws_request_id, } request_time_epoch_ms = int(request_context.get("requestTimeEpoch")) if is_api_gateway_invocation_async(event): diff --git a/tests/test_tracing.py b/tests/test_tracing.py index fed90588..25865d5e 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -755,9 +755,7 @@ def test_create_inferred_span_from_api_gateway_event(self): span.get_tag("resource_names"), "POST /path/to/resource", ) - self.assertEqual( - span.get_tag("request_id"), "c6af9ac6-7b61-11e6-9a41-93e8deadbeef" - ) + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "1234567890") self.assertEqual(span.get_tag("apiname"), "1234567890") self.assertEqual(span.get_tag("stage"), "prod") @@ -789,9 +787,7 @@ def test_create_inferred_span_from_api_gateway_non_proxy_event_async(self): span.get_tag("resource_names"), "GET /http/get", ) - self.assertEqual( - span.get_tag("request_id"), "7bf3b161-f698-432c-a639-6fef8b445137" - ) + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "lgxbo6a518") self.assertEqual(span.get_tag("apiname"), "lgxbo6a518") self.assertEqual(span.get_tag("stage"), "dev") @@ -823,9 +819,7 @@ def test_create_inferred_span_from_api_gateway_non_proxy_event_sync(self): span.get_tag("resource_names"), "GET /http/get", ) - self.assertEqual( - span.get_tag("request_id"), "7bf3b161-f698-432c-a639-6fef8b445137" - ) + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "lgxbo6a518") self.assertEqual(span.get_tag("apiname"), "lgxbo6a518") self.assertEqual(span.get_tag("stage"), "dev") @@ -892,7 +886,7 @@ def test_create_inferred_span_from_api_gateway_websocket_default_event(self): span.get_tag("resource_names"), "$default", ) - self.assertEqual(span.get_tag("request_id"), "Fc5S3EvdGjQFtsQ=") + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "p62c47itsb") self.assertEqual(span.get_tag("apiname"), "p62c47itsb") self.assertEqual(span.get_tag("stage"), "dev") @@ -927,7 +921,7 @@ def test_create_inferred_span_from_api_gateway_websocket_connect_event(self): span.get_tag("resource_names"), "$connect", ) - self.assertEqual(span.get_tag("request_id"), "Fc2tgH1RmjQFnOg=") + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "p62c47itsb") self.assertEqual(span.get_tag("apiname"), "p62c47itsb") self.assertEqual(span.get_tag("stage"), "dev") @@ -962,7 +956,7 @@ def test_create_inferred_span_from_api_gateway_websocket_disconnect_event(self): span.get_tag("resource_names"), "$disconnect", ) - self.assertEqual(span.get_tag("request_id"), "Fc2ydE4LmjQFhdg=") + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "p62c47itsb") self.assertEqual(span.get_tag("apiname"), "p62c47itsb") self.assertEqual(span.get_tag("stage"), "dev") @@ -1358,9 +1352,7 @@ def test_create_inferred_span_from_api_gateway_event_no_apiid(self): span.get_tag("resource_names"), "POST /path/to/resource", ) - self.assertEqual( - span.get_tag("request_id"), "c6af9ac6-7b61-11e6-9a41-93e8deadbeef" - ) + self.assertEqual(span.get_tag("request_id"), "123") self.assertEqual(span.get_tag("apiid"), "None") self.assertEqual(span.get_tag("apiname"), "None") self.assertEqual(span.get_tag("stage"), "prod")