Skip to content

[SVLS-4647] Fix lambda authorizer trace context injection issue #451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[SVLS-4647] Fix lambda authorizer trace context injection issue #451

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions datadog_lambda/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
TraceContextSource,
XraySubsegment,
Headers,
TraceHeader,
)
from datadog_lambda.metric import (
flush_stats,
Expand All @@ -44,6 +43,7 @@
InferredSpanInfo,
is_authorizer_response,
tracer,
propagator,
)
from datadog_lambda.trigger import (
extract_trigger_tags,
Expand Down Expand Up @@ -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
Expand Down
19 changes: 9 additions & 10 deletions tests/test_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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()
Expand All @@ -680,7 +680,6 @@ def flush():
({"DD_LOCAL_TEST": ""}, False),
({}, False),
):

os.environ = environ
flushes.clear()

Expand Down