fix lambda authorizer trace context injection issue (#453)

DataDog · Mar 27, 2024 · 8e79a70 · 8e79a70
1 parent 6b9eda0
commit 8e79a70
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 18 deletions.
diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
@@ -22,7 +22,6 @@
     TraceContextSource,
     XraySubsegment,
     Headers,
-    TraceHeader,
 )
 from datadog_lambda.metric import (
     flush_stats,
@@ -44,6 +43,7 @@
     InferredSpanInfo,
     is_authorizer_response,
     tracer,
+    propagator,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -254,13 +254,7 @@ def _inject_authorizer_span_headers(self, request_id):
         injected_headers = {}
         source_span = self.inferred_span if self.inferred_span else self.span
         span_context = source_span.context
-        injected_headers[TraceHeader.TRACE_ID] = str(span_context.trace_id)
-        injected_headers[TraceHeader.PARENT_ID] = str(span_context.span_id)
-        sampling_priority = span_context.sampling_priority
-        if sampling_priority is not None:
-            injected_headers[TraceHeader.SAMPLING_PRIORITY] = str(
-                span_context.sampling_priority
-            )
+        propagator.inject(span_context, injected_headers)
         injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns
         if request_id is not None:
             injected_headers[Headers.Authorizing_Request_Id] = request_id

diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
@@ -10,6 +10,7 @@
 from datadog_lambda.metric import lambda_metric
 from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 from ddtrace import Span, tracer
+from ddtrace.internal.constants import MAX_UINT_64BITS
 
 
 def get_mock_context(
@@ -543,19 +544,19 @@ def lambda_handler(event, context):
         lambda_event = {}
 
         lambda_context = get_mock_context()
-        mock_span = Span(name="my_inferred_span", span_id=123, trace_id=456)
-        mock_span.context.sampling_priority = "1"
-        mock_span.context.dd_origin = None
-        mock_span.start_ns = 1668127541671386817
-        mock_span.duration_ns = 1e8
-        lambda_handler.inferred_span = mock_span
+        test_span = tracer.trace("test_span")
+        trace_ctx = tracer.current_trace_context()
+        test_span.finish()
+        lambda_handler.inferred_span = test_span
         lambda_handler.make_inferred_span = False
         result = lambda_handler(lambda_event, lambda_context)
         raw_inject_data = result["context"]["_datadog"]
         self.assertIsInstance(raw_inject_data, str)
         inject_data = json.loads(base64.b64decode(raw_inject_data))
-        self.assertEqual(inject_data[TraceHeader.PARENT_ID], "123")
-        self.assertEqual(inject_data[TraceHeader.TRACE_ID], "456")
+        self.assertEqual(inject_data[TraceHeader.PARENT_ID], str(trace_ctx.span_id))
+        self.assertEqual(
+            inject_data[TraceHeader.TRACE_ID], str(MAX_UINT_64BITS & trace_ctx.trace_id)
+        )
         self.assertEqual(inject_data[TraceHeader.SAMPLING_PRIORITY], "1")
         self.assertEqual(result["context"]["scope"], "still here")
 
@@ -662,7 +663,6 @@ def tearDown(self):
 
     @patch("datadog_lambda.wrapper.should_use_extension", True)
     def test_local_test_envvar_flushing(self):
-
         flushes = []
         lambda_event = {}
         lambda_context = get_mock_context()
@@ -680,7 +680,6 @@ def flush():
             ({"DD_LOCAL_TEST": ""}, False),
             ({}, False),
         ):
-
             os.environ = environ
             flushes.clear()