use context.aws_request_id instead of http payload's requestId in the spans' tags

Author: joeyzhao2018

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):

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")