diff --git a/CHANGELOG.md b/CHANGELOG.md index e17542b104..eead4dd886 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +### Added + +- `opentelemetry-sdk-extension-aws` Add AwsXrayLambdaPropagator + ([#2573](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2573)) + ### Breaking changes - `opentelemetry-instrumentation-asgi`, `opentelemetry-instrumentation-fastapi`, `opentelemetry-instrumentation-starlette` Use `tracer` and `meter` of originating components instead of one from `asgi` middleware diff --git a/propagator/opentelemetry-propagator-aws-xray/pyproject.toml b/propagator/opentelemetry-propagator-aws-xray/pyproject.toml index 69fd0bbbfa..4a3e22269a 100644 --- a/propagator/opentelemetry-propagator-aws-xray/pyproject.toml +++ b/propagator/opentelemetry-propagator-aws-xray/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ [project.entry-points.opentelemetry_propagator] xray = "opentelemetry.propagators.aws:AwsXRayPropagator" +xray_lambda = "opentelemetry.propagators.aws:AwsXRayLambdaPropagator" [project.urls] Homepage = "https://github.com/open-telemetry/opentelemetry-python-contrib/tree/main/propagator/opentelemetry-propagator-aws-xray" diff --git a/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py b/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py index 4e4a6872ea..4966218211 100644 --- a/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py +++ b/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py @@ -58,6 +58,7 @@ import logging import typing +from os import environ from opentelemetry import trace from opentelemetry.context import Context @@ -71,6 +72,7 @@ ) TRACE_HEADER_KEY = "X-Amzn-Trace-Id" +AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID" KV_PAIR_DELIMITER = ";" KEY_AND_VALUE_DELIMITER = "=" @@ -324,3 +326,33 @@ def fields(self): """Returns a set with the fields set in `inject`.""" return {TRACE_HEADER_KEY} + + +class AwsXrayLambdaPropagator(AwsXRayPropagator): + """Implementation of the AWS X-Ray Trace Header propagation protocol but + with special handling for Lambda's ``_X_AMZN_TRACE_ID` environment + variable. + """ + + def extract( + self, + carrier: CarrierT, + context: typing.Optional[Context] = None, + getter: Getter[CarrierT] = default_getter, + ) -> Context: + + xray_context = super().extract(carrier, context=context, getter=getter) + + if trace.get_current_span(context=context).get_span_context().is_valid: + return xray_context + + trace_header = environ.get(AWS_TRACE_HEADER_ENV_KEY) + + if trace_header is None: + return xray_context + + return super().extract( + {TRACE_HEADER_KEY: trace_header}, + context=xray_context, + getter=getter, + ) diff --git a/propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_lambda_propagator.py b/propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_lambda_propagator.py new file mode 100644 index 0000000000..a0432d1457 --- /dev/null +++ b/propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_lambda_propagator.py @@ -0,0 +1,164 @@ +# Copyright The OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from os import environ +from unittest import TestCase +from unittest.mock import patch + +from requests.structures import CaseInsensitiveDict + +from opentelemetry.context import get_current +from opentelemetry.propagators.aws.aws_xray_propagator import ( + TRACE_HEADER_KEY, + AwsXrayLambdaPropagator, +) +from opentelemetry.propagators.textmap import DefaultGetter +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace import ( + Link, + NonRecordingSpan, + SpanContext, + TraceState, + get_current_span, + use_span, +) + + +class AwsXRayLambdaPropagatorTest(TestCase): + + def test_extract_no_environment_variable(self): + + actual_context = get_current_span( + AwsXrayLambdaPropagator().extract( + {}, context=get_current(), getter=DefaultGetter() + ) + ).get_span_context() + + self.assertEqual(hex(actual_context.trace_id), "0x0") + self.assertEqual(hex(actual_context.span_id), "0x0") + self.assertFalse( + actual_context.trace_flags.sampled, + ) + self.assertEqual(actual_context.trace_state, TraceState.get_default()) + + def test_extract_no_environment_variable_valid_context(self): + + with use_span(NonRecordingSpan(SpanContext(1, 2, False))): + + actual_context = get_current_span( + AwsXrayLambdaPropagator().extract( + {}, context=get_current(), getter=DefaultGetter() + ) + ).get_span_context() + + self.assertEqual(hex(actual_context.trace_id), "0x1") + self.assertEqual(hex(actual_context.span_id), "0x2") + self.assertFalse( + actual_context.trace_flags.sampled, + ) + self.assertEqual( + actual_context.trace_state, TraceState.get_default() + ) + + @patch.dict( + environ, + { + "_X_AMZN_TRACE_ID": ( + "Root=1-00000001-d188f8fa79d48a391a778fa6;" + "Parent=53995c3f42cd8ad8;Sampled=1;Foo=Bar" + ) + }, + ) + def test_extract_from_environment_variable(self): + + actual_context = get_current_span( + AwsXrayLambdaPropagator().extract( + {}, context=get_current(), getter=DefaultGetter() + ) + ).get_span_context() + + self.assertEqual( + hex(actual_context.trace_id), "0x1d188f8fa79d48a391a778fa6" + ) + self.assertEqual(hex(actual_context.span_id), "0x53995c3f42cd8ad8") + self.assertTrue( + actual_context.trace_flags.sampled, + ) + self.assertEqual(actual_context.trace_state, TraceState.get_default()) + + @patch.dict( + environ, + { + "_X_AMZN_TRACE_ID": ( + "Root=1-00000002-240000000000000000000002;" + "Parent=1600000000000002;Sampled=1;Foo=Bar" + ) + }, + ) + def test_add_link_from_environment_variable(self): + + propagator = AwsXrayLambdaPropagator() + + default_getter = DefaultGetter() + + carrier = CaseInsensitiveDict( + { + TRACE_HEADER_KEY: ( + "Root=1-00000001-240000000000000000000001;" + "Parent=1600000000000001;Sampled=1" + ) + } + ) + + extracted_context = propagator.extract( + carrier, context=get_current(), getter=default_getter + ) + + link_context = propagator.extract( + carrier, context=extracted_context, getter=default_getter + ) + + span = ReadableSpan( + "test", parent=extracted_context, links=[Link(link_context)] + ) + + span_parent_context = get_current_span(span.parent).get_span_context() + + self.assertEqual( + hex(span_parent_context.trace_id), "0x2240000000000000000000002" + ) + self.assertEqual( + hex(span_parent_context.span_id), "0x1600000000000002" + ) + self.assertTrue( + span_parent_context.trace_flags.sampled, + ) + self.assertEqual( + span_parent_context.trace_state, TraceState.get_default() + ) + + span_link_context = get_current_span( + span.links[0].context + ).get_span_context() + + self.assertEqual( + hex(span_link_context.trace_id), "0x1240000000000000000000001" + ) + self.assertEqual(hex(span_link_context.span_id), "0x1600000000000001") + self.assertTrue( + span_link_context.trace_flags.sampled, + ) + self.assertEqual( + span_link_context.trace_state, TraceState.get_default() + )