Handle ExtendedAttributes in BoundedAttributes

Author: xrmx

opentelemetry-api/src/opentelemetry/attributes/__init__.py

@@ -118,6 +118,96 @@ def _clean_attribute(
     return None
 
 
+def _clean_extended_attribute_value(
+    value: types.AttributeValue, max_len: Optional[int]
+) -> Optional[Union[types.AttributeValue, Tuple[Union[str, int, float], ...]]]:
+    # for primitive types just return the value and eventually shorten the string length
+    if value is None or isinstance(value, _VALID_ATTR_VALUE_TYPES):
+        if max_len is not None and isinstance(value, str):
+            value = value[:max_len]
+        return value
+
+    if isinstance(value, Mapping):
+        cleaned_dict = {}
+        for key, element in value.items():
+            # skip invalid keys
+            if not (key and isinstance(key, str)):
+                _logger.warning(
+                    "invalid key `%s`. must be non-empty string.", key
+                )
+                continue
+
+            cleaned_dict[key] = _clean_extended_attribute(
+                key=key, value=element, max_len=max_len
+            )
+
+        return cleaned_dict
+
+    if isinstance(value, Sequence):
+        sequence_first_valid_type = None
+        cleaned_seq = []
+
+        for element in value:
+            if element is None:
+                cleaned_seq.append(element)
+                continue
+
+            if max_len is not None and isinstance(element, str):
+                element = element[:max_len]
+
+            element_type = type(element)
+            if element_type not in _VALID_ATTR_VALUE_TYPES:
+                return _clean_extended_attribute(element, max_len=max_len)
+
+            # The type of the sequence must be homogeneous. The first non-None
+            # element determines the type of the sequence
+            if sequence_first_valid_type is None:
+                sequence_first_valid_type = element_type
+            # use equality instead of isinstance as isinstance(True, int) evaluates to True
+            elif element_type != sequence_first_valid_type:
+                _logger.warning(
+                    "Mixed types %s and %s in attribute value sequence",
+                    sequence_first_valid_type.__name__,
+                    type(element).__name__,
+                )
+                return None
+
+            cleaned_seq.append(element)
+
+        # Freeze mutable sequences defensively
+        return tuple(cleaned_seq)
+
+    raise TypeError(
+        "Invalid type %s for attribute value. Expected one of %s or a "
+        "sequence of those types",
+        type(value).__name__,
+        [valid_type.__name__ for valid_type in _VALID_ANY_VALUE_TYPES],
+    )
+
+
+def _clean_extended_attribute(
+    key: str, value: types.AttributeValue, max_len: Optional[int]
+) -> Optional[Union[types.AttributeValue, Tuple[Union[str, int, float], ...]]]:
+    """Checks if attribute value is valid and cleans it if required.
+
+    The function returns the cleaned value or None if the value is not valid.
+
+    An attribute value is valid if it is an AnyValue.
+    An attribute needs cleansing if:
+        - Its length is greater than the maximum allowed length.
+    """
+
+    if not (key and isinstance(key, str)):
+        _logger.warning("invalid key `%s`. must be non-empty string.", key)
+        return None
+
+    try:
+        return _clean_extended_attribute_value(value, max_len=max_len)
+    except TypeError as exception:
+        _logger.warning(f"Attribute {key}: {exception}")
+        return None
+
+
 def _clean_attribute_value(
     value: types.AttributeValue, limit: Optional[int]
 ) -> Optional[types.AttributeValue]:
@@ -149,6 +239,7 @@ def __init__(
         attributes: types.Attributes = None,
         immutable: bool = True,
         max_value_len: Optional[int] = None,
+        extended_attributes: bool = False,
     ):
         if maxlen is not None:
             if not isinstance(maxlen, int) or maxlen < 0:
@@ -158,6 +249,7 @@ def __init__(
         self.maxlen = maxlen
         self.dropped = 0
         self.max_value_len = max_value_len
+        self._extended_attributes = extended_attributes
         # OrderedDict is not used until the maxlen is reached for efficiency.
 
         self._dict: Union[
@@ -184,19 +276,24 @@ def __setitem__(self, key: str, value: types.AttributeValue) -> None:
                 self.dropped += 1
                 return
 
-            value = _clean_attribute(key, value, self.max_value_len)  # type: ignore
-            if value is not None:
-                if key in self._dict:
-                    del self._dict[key]
-                elif (
-                    self.maxlen is not None and len(self._dict) == self.maxlen
-                ):
-                    if not isinstance(self._dict, OrderedDict):
-                        self._dict = OrderedDict(self._dict)
-                    self._dict.popitem(last=False)  # type: ignore
-                    self.dropped += 1
-
-                self._dict[key] = value  # type: ignore
+            if self._extended_attributes:
+                value = _clean_extended_attribute(
+                    key, value, self.max_value_len
+                )  # type: ignore
+            else:
+                value = _clean_attribute(key, value, self.max_value_len)  # type: ignore
+                if value is None:
+                    return
+
+            if key in self._dict:
+                del self._dict[key]
+            elif self.maxlen is not None and len(self._dict) == self.maxlen:
+                if not isinstance(self._dict, OrderedDict):
+                    self._dict = OrderedDict(self._dict)
+                self._dict.popitem(last=False)  # type: ignore
+                self.dropped += 1
+
+            self._dict[key] = value  # type: ignore
 
     def __delitem__(self, key: str) -> None:
         if getattr(self, "_immutable", False):  # type: ignore

opentelemetry-api/tests/attributes/test_attributes.py

@@ -17,7 +17,11 @@
 import unittest
 from typing import MutableSequence
 
-from opentelemetry.attributes import BoundedAttributes, _clean_attribute
+from opentelemetry.attributes import (
+    BoundedAttributes,
+    _clean_attribute,
+    _clean_extended_attribute,
+)
 
 
 class TestAttributes(unittest.TestCase):
@@ -89,6 +93,94 @@ def test_sequence_attr_decode(self):
         )
 
 
+class TestExtendedAttributes(unittest.TestCase):
+    # pylint: disable=invalid-name
+    def assertValid(self, value, key="k"):
+        expected = value
+        if isinstance(value, MutableSequence):
+            expected = tuple(value)
+        self.assertEqual(_clean_extended_attribute(key, value, None), expected)
+
+    def assertInvalid(self, value, key="k"):
+        self.assertIsNone(_clean_extended_attribute(key, value, None))
+
+    def test_attribute_key_validation(self):
+        # only non-empty strings are valid keys
+        self.assertInvalid(1, "")
+        self.assertInvalid(1, 1)
+        self.assertInvalid(1, {})
+        self.assertInvalid(1, [])
+        self.assertInvalid(1, b"1")
+        self.assertValid(1, "k")
+        self.assertValid(1, "1")
+
+    def test_clean_extended_attribute(self):
+        self.assertInvalid([1, 2, 3.4, "ss", 4])
+        self.assertInvalid([{}, 1, 2, 3.4, 4])
+        self.assertInvalid(["sw", "lf", 3.4, "ss"])
+        self.assertInvalid([1, 2, 3.4, 5])
+        self.assertInvalid([1, True])
+        self.assertValid(None)
+        self.assertValid(True)
+        self.assertValid("hi")
+        self.assertValid(3.4)
+        self.assertValid(15)
+        self.assertValid([1, 2, 3, 5])
+        self.assertValid([1.2, 2.3, 3.4, 4.5])
+        self.assertValid([True, False])
+        self.assertValid(["ss", "dw", "fw"])
+        self.assertValid([])
+        # None in sequences are valid
+        self.assertValid(["A", None, None])
+        self.assertValid(["A", None, None, "B"])
+        self.assertValid([None, None])
+        self.assertInvalid(["A", None, 1])
+        self.assertInvalid([None, "A", None, 1])
+        # mappings
+        self.assertValid({})
+        self.assertValid({"k": "v"})
+
+        # test keys
+        self.assertValid("value", "key")
+        self.assertInvalid("value", "")
+        self.assertInvalid("value", None)
+
+    def test_sequence_attr_decode(self):
+        seq = [
+            None,
+            b"Content-Disposition",
+            b"Content-Type",
+            b"\x81",
+            b"Keep-Alive",
+        ]
+        self.assertEqual(
+            _clean_extended_attribute("headers", seq, None), tuple(seq)
+        )
+
+    def test_mapping(self):
+        mapping = {
+            "": "invalid",
+            b"bytes": "invalid",
+            "none": {"": "invalid"},
+            "valid_primitive": "str",
+            "valid_sequence": ["str"],
+            "invalid_sequence": ["str", 1],
+            "valid_mapping": {"str": 1},
+            "invalid_mapping": {"": 1},
+        }
+        expected = {
+            "none": {},
+            "valid_primitive": "str",
+            "valid_sequence": ("str",),
+            "invalid_sequence": None,
+            "valid_mapping": {"str": 1},
+            "invalid_mapping": {},
+        }
+        self.assertEqual(
+            _clean_extended_attribute("headers", mapping, None), expected
+        )
+
+
 class TestBoundedAttributes(unittest.TestCase):
     # pylint: disable=consider-using-dict-items
     base = {
@@ -196,3 +288,13 @@ def test_locking(self):
 
         for num in range(100):
             self.assertEqual(bdict[str(num)], num)
+
+    def test_extended_attributes(self):
+        bdict = BoundedAttributes(extended_attributes=True, immutable=False)
+        with unittest.mock.patch(
+            "opentelemetry.attributes._clean_extended_attribute",
+            return_value="mock_value",
+        ) as clean_extended_attribute_mock:
+            bdict["key"] = "value"
+
+        clean_extended_attribute_mock.assert_called_once()