fix lambda authorizer trace context injection issue (#453)

Author: joeyzhao2018

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

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