Skip to content

Commit 314aa88

Browse files
authored
fix(lambda): align Python context propagation with JS implementation (#727)
Closes #663 ## Summary - **Problem:** `custom_event_context_extractor` in `otel_wrapper.py` always used the X-Ray env var (`_X_AMZN_TRACE_ID`) as parent context because it is always set and valid when active tracing is enabled, causing upstream W3C trace context (`traceparent`) from event headers to be ignored and breaking trace continuity. - **Fix:** Align with the [JS implementation](https://github.com/aws-observability/aws-otel-js-instrumentation/blob/main/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts#L100-L123) — inject the X-Ray env var into the event headers (replacing any existing X-Ray header, case-insensitive), then delegate to the global composite propagator via `get_global_textmap().extract()`. This respects the propagator priority ordering configured in `otel-instrument` (`baggage,xray,tracecontext`), where `tracecontext` (W3C) takes precedence over `xray` when both are present. - **Tests updated:** Existing tests adapted to explicitly set the global textmap propagator. Two new tests added: - `test_xray_ignored_when_propagator_does_not_include_xray` — X-Ray active tracing on, but only `tracecontext` propagator configured; W3C headers should be used. - `test_w3c_takes_precedence_over_xray_when_both_present` — both X-Ray and W3C headers present with composite propagator; W3C should win. ## Test plan - [x] All 4 unit tests pass locally (`pytest -v`) - [x] Coverage: `otel_wrapper.py` at 90% (uncovered lines are error-handling edge cases) - [x] Lint passes: `black`, `isort`, `flake8` all clean
1 parent 5cdb6d7 commit 314aa88

3 files changed

Lines changed: 160 additions & 60 deletions

File tree

CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ If your change does not need a CHANGELOG entry, add the "skip changelog" label t
1212

1313
## Unreleased
1414

15+
- fix(lambda-layer): align context propagation with JS — delegate to global propagator so W3C traceparent is no longer ignored when X-Ray active tracing is enabled
16+
([#727](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/727))
1517
## v0.17.0 - 2026-04-08
1618

1719
- feat: support environment-configured endpoint visibility for HTTP operation names

lambda-layer/src/otel_wrapper.py

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,6 @@
4141
from opentelemetry.context import Context
4242
from opentelemetry.instrumentation.aws_lambda import _X_AMZN_TRACE_ID, AwsLambdaInstrumentor
4343
from opentelemetry.propagate import get_global_textmap
44-
from opentelemetry.propagators.aws import AwsXRayPropagator
4544
from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY
4645
from opentelemetry.trace import get_current_span
4746

@@ -57,21 +56,24 @@ class HandlerError(Exception):
5756

5857
def custom_event_context_extractor(lambda_event: Any) -> Context:
5958
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
60-
lambda_trace_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var})
61-
parent_span_context = get_current_span(lambda_trace_context).get_span_context()
6259

63-
if parent_span_context is None or not parent_span_context.is_valid:
60+
try:
61+
headers = lambda_event["headers"]
62+
except (TypeError, KeyError):
6463
headers = None
65-
try:
66-
headers = lambda_event["headers"]
67-
except (TypeError, KeyError):
68-
pass
69-
if not isinstance(headers, dict):
70-
headers = {}
64+
if not isinstance(headers, dict):
65+
headers = {}
66+
else:
67+
headers = headers.copy()
7168

72-
return get_global_textmap().extract(headers)
69+
if xray_env_var:
70+
headers = {k: v for k, v in headers.items() if k.lower() != TRACE_HEADER_KEY.lower()}
71+
headers[TRACE_HEADER_KEY] = xray_env_var
7372

74-
return lambda_trace_context
73+
extracted_context = get_global_textmap().extract(headers)
74+
if get_current_span(extracted_context).get_span_context():
75+
return extracted_context
76+
return Context()
7577

7678

7779
AwsLambdaInstrumentor().instrument(event_context_extractor=custom_event_context_extractor)
@@ -82,7 +84,7 @@ def custom_event_context_extractor(lambda_event: Any) -> Context:
8284
raise HandlerError("ORIG_HANDLER is not defined.")
8385

8486
try:
85-
(mod_name, handler_name) = path.rsplit(".", 1)
87+
mod_name, handler_name = path.rsplit(".", 1)
8688
except ValueError as e:
8789
raise HandlerError("Bad path '{}' for ORIG_HANDLER: {}".format(path, str(e)))
8890

lambda-layer/src/tests/test_lambda_instrumentation.py

Lines changed: 143 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,11 @@
2525
from shutil import which
2626
from unittest import mock
2727

28-
from opentelemetry.environment_variables import OTEL_PROPAGATORS
2928
from opentelemetry.instrumentation.aws_lambda import _HANDLER, _X_AMZN_TRACE_ID, ORIG_HANDLER, AwsLambdaInstrumentor
29+
from opentelemetry.propagate import get_global_textmap, set_global_textmap
30+
from opentelemetry.propagators.aws import AwsXRayPropagator
3031
from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_ID_FIRST_PART_LENGTH, TRACE_ID_VERSION
32+
from opentelemetry.propagators.composite import CompositePropagator
3133
from opentelemetry.semconv.resource import ResourceAttributes
3234
from opentelemetry.semconv.trace import SpanAttributes
3335
from opentelemetry.test.test_base import TestBase
@@ -196,6 +198,9 @@ def tearDownClass(cls):
196198
)
197199

198200
def test_active_tracing(self):
201+
original_propagator = get_global_textmap()
202+
set_global_textmap(AwsXRayPropagator())
203+
199204
test_env_patch = mock.patch.dict(
200205
"os.environ",
201206
{
@@ -206,71 +211,162 @@ def test_active_tracing(self):
206211
)
207212
test_env_patch.start()
208213

209-
mock_execute_lambda()
210-
211-
spans = self.memory_exporter.get_finished_spans()
214+
try:
215+
mock_execute_lambda()
216+
217+
spans = self.memory_exporter.get_finished_spans()
218+
219+
assert spans
220+
221+
self.assertEqual(len(spans), 1)
222+
span = spans[0]
223+
self.assertEqual(span.name, os.environ[ORIG_HANDLER])
224+
self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
225+
self.assertEqual(span.kind, SpanKind.SERVER)
226+
self.assertSpanHasAttributes(
227+
span,
228+
{
229+
ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
230+
SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id,
231+
},
232+
)
233+
234+
parent_context = span.parent
235+
self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
236+
self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
237+
self.assertTrue(parent_context.is_remote)
238+
finally:
239+
test_env_patch.stop()
240+
set_global_textmap(original_propagator)
212241

213-
assert spans
242+
def test_parent_context_from_lambda_event(self):
243+
original_propagator = get_global_textmap()
244+
set_global_textmap(TraceContextTextMapPropagator())
214245

215-
self.assertEqual(len(spans), 1)
216-
span = spans[0]
217-
self.assertEqual(span.name, os.environ[ORIG_HANDLER])
218-
self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
219-
self.assertEqual(span.kind, SpanKind.SERVER)
220-
self.assertSpanHasAttributes(
221-
span,
246+
test_env_patch = mock.patch.dict(
247+
"os.environ",
222248
{
223-
ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
224-
SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id,
249+
**os.environ,
250+
# NOT Active Tracing
251+
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH,
225252
},
226253
)
254+
test_env_patch.start()
227255

228-
parent_context = span.parent
229-
self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
230-
self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
231-
self.assertTrue(parent_context.is_remote)
232-
233-
test_env_patch.stop()
256+
try:
257+
trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
258+
mock_execute_lambda(
259+
{
260+
"headers": {
261+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
262+
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
263+
}
264+
}
265+
)
266+
267+
spans = self.memory_exporter.get_finished_spans()
268+
269+
assert spans
270+
271+
self.assertEqual(len(spans), 1)
272+
span = spans[0]
273+
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
274+
275+
parent_context = span.parent
276+
self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
277+
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
278+
self.assertEqual(len(parent_context.trace_state), 3)
279+
self.assertEqual(
280+
parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
281+
MOCK_W3C_TRACE_STATE_VALUE,
282+
)
283+
self.assertTrue(parent_context.is_remote)
284+
finally:
285+
test_env_patch.stop()
286+
set_global_textmap(original_propagator)
287+
288+
def test_xray_ignored_when_propagator_does_not_include_xray(self):
289+
"""When X-Ray active tracing is on but the user only configures
290+
tracecontext propagator (no xray), the X-Ray env var is injected
291+
into headers but the propagator cannot parse it. The W3C context
292+
from event headers should be used instead."""
293+
original_propagator = get_global_textmap()
294+
set_global_textmap(TraceContextTextMapPropagator())
234295

235-
def test_parent_context_from_lambda_event(self):
236296
test_env_patch = mock.patch.dict(
237297
"os.environ",
238298
{
239299
**os.environ,
240-
# NOT Active Tracing
241-
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH,
242-
# NOT using the X-Ray Propagator
243-
OTEL_PROPAGATORS: "tracecontext",
300+
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
244301
},
245302
)
246303
test_env_patch.start()
247304

248-
trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
249-
mock_execute_lambda(
250-
{
251-
"headers": {
252-
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
253-
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
305+
try:
306+
mock_execute_lambda(
307+
{
308+
"headers": {
309+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
310+
}
254311
}
255-
}
256-
)
312+
)
313+
314+
spans = self.memory_exporter.get_finished_spans()
257315

258-
spans = self.memory_exporter.get_finished_spans()
316+
assert spans
317+
self.assertEqual(len(spans), 1)
318+
span = spans[0]
259319

260-
assert spans
320+
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
261321

262-
self.assertEqual(len(spans), 1)
263-
span = spans[0]
264-
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
322+
parent_context = span.parent
323+
self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID)
324+
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
325+
self.assertTrue(parent_context.is_remote)
326+
finally:
327+
test_env_patch.stop()
328+
set_global_textmap(original_propagator)
265329

266-
parent_context = span.parent
267-
self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
268-
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
269-
self.assertEqual(len(parent_context.trace_state), 3)
270-
self.assertEqual(
271-
parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
272-
MOCK_W3C_TRACE_STATE_VALUE,
330+
def test_w3c_takes_precedence_over_xray_when_both_present(self):
331+
"""When both X-Ray active tracing and W3C traceparent headers are
332+
present, the W3C context should be used as the parent because
333+
tracecontext comes after xray in the composite propagator — the
334+
last propagator to extract a valid context wins."""
335+
original_propagator = get_global_textmap()
336+
set_global_textmap(CompositePropagator([AwsXRayPropagator(), TraceContextTextMapPropagator()]))
337+
338+
test_env_patch = mock.patch.dict(
339+
"os.environ",
340+
{
341+
**os.environ,
342+
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
343+
},
273344
)
274-
self.assertTrue(parent_context.is_remote)
345+
test_env_patch.start()
346+
347+
try:
348+
trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
349+
mock_execute_lambda(
350+
{
351+
"headers": {
352+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
353+
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
354+
}
355+
}
356+
)
357+
358+
spans = self.memory_exporter.get_finished_spans()
359+
360+
assert spans
361+
self.assertEqual(len(spans), 1)
362+
span = spans[0]
363+
364+
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
275365

276-
test_env_patch.stop()
366+
parent_context = span.parent
367+
self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID)
368+
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
369+
self.assertTrue(parent_context.is_remote)
370+
finally:
371+
test_env_patch.stop()
372+
set_global_textmap(original_propagator)

0 commit comments

Comments
 (0)