diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 73d17887..057f048c 100644 --- 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 index ee1ef99c..141e245a 100644 --- 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()