Fix: Support new fields in X-Ray API responses

Author: majanjua-amzn

CHANGELOG.md

@@ -11,6 +11,9 @@ For any change that affects end users of this package, please add an entry under
 If your change does not need a CHANGELOG entry, add the "skip changelog" label to your PR.
 
 ## Unreleased
+
+- Fix: Support new fields in X-Ray API responses
+  ([#577](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/577))
 - Sign Lambda layer by AWS Signer
   ([#573](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/573))
 - Support PyPI Signature

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/_sampling_rule.py

@@ -1,6 +1,10 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from logging import getLogger
+
+_logger = getLogger(__name__)
+
 
 # Disable snake_case naming style so this class can match the sampling rules response from X-Ray
 # pylint: disable=invalid-name
@@ -20,6 +24,7 @@ def __init__(
         ServiceType: str = None,
         URLPath: str = None,
         Version: int = None,
+        **kwargs,
     ):
         self.Attributes = Attributes if Attributes is not None else {}
         self.FixedRate = FixedRate if FixedRate is not None else 0.0
@@ -36,6 +41,10 @@ def __init__(
         self.URLPath = URLPath if URLPath is not None else ""
         self.Version = Version if Version is not None else 0
 
+        # Log unknown fields for debugging/monitoring
+        if kwargs:
+            _logger.debug("Ignoring unknown fields in _SamplingRule: %s", list(kwargs.keys()))
+
     def __lt__(self, other: "_SamplingRule") -> bool:
         if self.Priority == other.Priority:
             # String order priority example:

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/_sampling_target.py

@@ -1,6 +1,7 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 from logging import getLogger
+from typing import List
 
 _logger = getLogger(__name__)
 
@@ -15,47 +16,62 @@ def __init__(
         ReservoirQuota: int = None,
         ReservoirQuotaTTL: float = None,
         RuleName: str = None,
+        **kwargs,
     ):
         self.FixedRate = FixedRate if FixedRate is not None else 0.0
         self.Interval = Interval  # can be None
         self.ReservoirQuota = ReservoirQuota  # can be None
         self.ReservoirQuotaTTL = ReservoirQuotaTTL  # can be None
         self.RuleName = RuleName if RuleName is not None else ""
 
+        # Log unknown fields for debugging/monitoring
+        if kwargs:
+            _logger.debug("Ignoring unknown fields in _SamplingTarget: %s", list(kwargs.keys()))
+
 
 class _UnprocessedStatistics:
     def __init__(
         self,
         ErrorCode: str = None,
         Message: str = None,
         RuleName: str = None,
+        **kwargs,
     ):
         self.ErrorCode = ErrorCode if ErrorCode is not None else ""
         self.Message = Message if ErrorCode is not None else ""
         self.RuleName = RuleName if ErrorCode is not None else ""
 
+        # Log unknown fields for debugging/monitoring
+        if kwargs:
+            _logger.debug("Ignoring unknown fields in _UnprocessedStatistics: %s", list(kwargs.keys()))
+
 
 class _SamplingTargetResponse:
     def __init__(
         self,
         LastRuleModification: float,
-        SamplingTargetDocuments: [dict] = None,
-        UnprocessedStatistics: [dict] = None,
+        SamplingTargetDocuments: List[dict] = None,
+        UnprocessedStatistics: List[dict] = None,
+        **kwargs,
     ):
         self.LastRuleModification: float = LastRuleModification if LastRuleModification is not None else 0.0
 
-        self.SamplingTargetDocuments: [_SamplingTarget] = []
+        self.SamplingTargetDocuments: List[_SamplingTarget] = []
         if SamplingTargetDocuments is not None:
             for document in SamplingTargetDocuments:
                 try:
                     self.SamplingTargetDocuments.append(_SamplingTarget(**document))
                 except TypeError as e:
                     _logger.debug("TypeError occurred: %s", e)
 
-        self.UnprocessedStatistics: [_UnprocessedStatistics] = []
+        self.UnprocessedStatistics: List[_UnprocessedStatistics] = []
         if UnprocessedStatistics is not None:
             for unprocessed in UnprocessedStatistics:
                 try:
                     self.UnprocessedStatistics.append(_UnprocessedStatistics(**unprocessed))
                 except TypeError as e:
                     _logger.debug("TypeError occurred: %s", e)
+
+        # Log unknown fields for debugging/monitoring
+        if kwargs:
+            _logger.debug("Ignoring unknown fields in _SamplingTargetResponse: %s", list(kwargs.keys()))

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_sampling_rule.py

@@ -1,6 +1,7 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 from unittest import TestCase
+from unittest.mock import patch
 
 from amazon.opentelemetry.distro.sampler._sampling_rule import _SamplingRule
 
@@ -16,6 +17,36 @@ def test_sampling_rule_ordering(self):
 
         self.assertTrue(rule1 < rule2 < rule3 < rule4 < rule5 < rule6)
 
+    def test_sampling_rule_with_extra_fields(self):
+        inputs = {
+            "Attributes": {},
+            "FixedRate": 0.1,
+            "HTTPMethod": "GET",
+            "Host": "localhost",
+            "Priority": 20,
+            "ReservoirSize": 1,
+            "ResourceARN": "*",
+            "RuleARN": "arn:aws:xray:us-east-1:999999999999:sampling-rule/test",
+            "RuleName": "test",
+            "ServiceName": "myServiceName",
+            "ServiceType": "AWS::EKS::Container",
+            "URLPath": "/helloworld",
+            "Version": 1,
+            "ExtraField1": "cat",
+            "ExtraField2": 123,
+        }
+
+        # Does not throw an error and logs debug message about unknown fields
+        with patch("amazon.opentelemetry.distro.sampler._sampling_rule._logger") as mock_logger:
+            rule = _SamplingRule(**inputs)
+            mock_logger.debug.assert_called_once_with(
+                "Ignoring unknown fields in _SamplingRule: %s", ["ExtraField1", "ExtraField2"]
+            )
+
+            self.assertEqual(rule.FixedRate, 0.1)
+            self.assertEqual(rule.RuleName, "test")
+            self.assertEqual(rule.ServiceName, "myServiceName")
+
     def test_sampling_rule_equality(self):
         sampling_rule = _SamplingRule(
             Attributes={"abc": "123", "def": "4?6", "ghi": "*89"},

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_sampling_target.py

@@ -1,11 +1,44 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 from unittest import TestCase
+from unittest.mock import patch
 
-from amazon.opentelemetry.distro.sampler._sampling_target import _SamplingTargetResponse
+from amazon.opentelemetry.distro.sampler._sampling_target import _SamplingTarget, _SamplingTargetResponse
 
 
 class TestSamplingTarget(TestCase):
+    def test_sampling_target_with_none_inputs(self):
+        target = _SamplingTarget()
+        self.assertEqual(target.FixedRate, 0.0)
+        self.assertEqual(target.RuleName, "")
+        self.assertIsNone(target.Interval)
+        self.assertIsNone(target.ReservoirQuota)
+        self.assertIsNone(target.ReservoirQuotaTTL)
+
+    def test_sampling_target_with_extra_inputs(self):
+        inputs = {
+            "FixedRate": 1.0,
+            "RuleName": "cat",
+            "Interval": 123,
+            "ReservoirQuota": 456,
+            "ReservoirQuotaTTL": 789,
+            "ExtraField1": "cat",
+            "ExtraField2": 123,
+        }
+
+        with patch("amazon.opentelemetry.distro.sampler._sampling_target._logger") as mock_logger:
+            target = _SamplingTarget(**inputs)
+            mock_logger.debug.assert_called_once_with(
+                "Ignoring unknown fields in _SamplingTarget: %s", ["ExtraField1", "ExtraField2"]
+            )
+
+            self.assertEqual(target.FixedRate, 1.0)
+            self.assertEqual(target.RuleName, "cat")
+            self.assertEqual(target.Interval, 123)
+            self.assertEqual(target.ReservoirQuota, 456)
+            self.assertEqual(target.ReservoirQuotaTTL, 789)
+            self.assertFalse(hasattr(target, "ExtraField2"))
+
     def test_sampling_target_response_with_none_inputs(self):
         target_response = _SamplingTargetResponse(None, None, None)
         self.assertEqual(target_response.LastRuleModification, 0.0)
@@ -28,5 +61,29 @@ def test_sampling_target_response_with_invalid_inputs(self):
         self.assertEqual(target_response.UnprocessedStatistics[0].RuleName, "")
 
         target_response = _SamplingTargetResponse(1.0, [{"foo": "bar"}], [{"dog": "cat"}])
-        self.assertEqual(len(target_response.SamplingTargetDocuments), 0)
-        self.assertEqual(len(target_response.UnprocessedStatistics), 0)
+        self.assertEqual(len(target_response.SamplingTargetDocuments), 1)
+        self.assertEqual(len(target_response.UnprocessedStatistics), 1)
+
+    def test_sampling_target_response_with_extra_inputs(self):
+        inputs = {
+            "LastRuleModification": 1.0,
+            "SamplingTargetDocuments": [{}],
+            "UnprocessedStatistics": [{}],
+            "ExtraField1": "cat",
+            "ExtraField2": 123,
+        }
+
+        # Does not throw an error and logs debug message about unknown fields
+        with patch("amazon.opentelemetry.distro.sampler._sampling_target._logger") as mock_logger:
+            target_response = _SamplingTargetResponse(**inputs)
+            mock_logger.debug.assert_called_once_with(
+                "Ignoring unknown fields in _SamplingTargetResponse: %s", ["ExtraField1", "ExtraField2"]
+            )
+
+            self.assertEqual(target_response.LastRuleModification, 1.0)
+            self.assertEqual(len(target_response.SamplingTargetDocuments), 1)
+            self.assertEqual(target_response.SamplingTargetDocuments[0].FixedRate, 0)
+            self.assertEqual(target_response.SamplingTargetDocuments[0].Interval, None)
+            self.assertEqual(target_response.SamplingTargetDocuments[0].ReservoirQuota, None)
+            self.assertEqual(target_response.SamplingTargetDocuments[0].ReservoirQuotaTTL, None)
+            self.assertEqual(target_response.SamplingTargetDocuments[0].RuleName, "")