aws-xray-propagator: ensure trace state is not overwritten

Author: majanjua-amzn

propagator/opentelemetry-propagator-aws-xray/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Update `opentelemetry-api` version to 1.16
   ([#2961](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2961))
+- aws-xray-propagator: ensure trace state is not overwritten
+  ([#3774](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3774))
 
 ## Version 1.0.2 (2024-08-05)
 

propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py

@@ -70,6 +70,7 @@
     default_getter,
     default_setter,
 )
+from opentelemetry.trace.span import TraceState
 
 TRACE_HEADER_KEY = "X-Amzn-Trace-Id"
 AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID"
@@ -92,6 +93,8 @@
 IS_SAMPLED = "1"
 NOT_SAMPLED = "0"
 
+TRACESTATE_HEADER_KEY = "tracestate"
+
 
 _logger = logging.getLogger(__name__)
 
@@ -144,12 +147,18 @@ def extract(
         if sampled:
             options |= trace.TraceFlags.SAMPLED
 
+        tracestate_headers = getter.get(carrier, TRACESTATE_HEADER_KEY)
+        if tracestate_headers is None:
+            tracestate = trace.TraceState()
+        else:
+            tracestate = TraceState.from_header(tracestate_headers)
+
         span_context = trace.SpanContext(
             trace_id=trace_id,
             span_id=span_id,
             is_remote=True,
             trace_flags=trace.TraceFlags(options),
-            trace_state=trace.TraceState(),
+            trace_state=tracestate,
         )
 
         if not span_context.is_valid:

propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_propagator.py

@@ -21,6 +21,7 @@
 from opentelemetry.context import Context
 from opentelemetry.propagators.aws.aws_xray_propagator import (
     TRACE_HEADER_KEY,
+    TRACESTATE_HEADER_KEY,
     AwsXRayPropagator,
 )
 from opentelemetry.trace import (
@@ -335,3 +336,19 @@ def test_fields(self, mock_trace):
         self.assertEqual(
             AwsXRayPropagatorTest.XRAY_PROPAGATOR.fields, inject_fields
         )
+
+    def test_extract_trace_state_from_headers(self):
+        """Test that extract properly propagates the trace state without modification."""
+        context_with_extracted = AwsXRayPropagatorTest.XRAY_PROPAGATOR.extract(
+            CaseInsensitiveDict(
+                {
+                    TRACE_HEADER_KEY: "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=0",
+                    TRACESTATE_HEADER_KEY: "foo=bar,baz=qux"
+                }
+            ),
+        )
+
+        extracted_span_context = get_nested_span_context(context_with_extracted)
+        expected_trace_state = TraceState([("foo", "bar"), ("baz", "qux")])
+        
+        self.assertEqual(extracted_span_context.trace_state, expected_trace_state)