add option to normalize tag values in TagManager (DataDog#21474)

* add option to normalize tag values in TagManager

* add changelog

* fix docstr

* rename params

Author: lu-zhengda

datadog_checks_base/changelog.d/21474.added

@@ -0,0 +1 @@
+Add tag normalization support to TagManager with optional tag_normalizer parameter and normalize flags for set_tag/set_tags_from_list/delete_tag methods.

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -471,12 +471,13 @@ class TagManager:
     multiple times.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, normalizer: Optional[Callable[[Union[str, bytes]], str]] = None) -> None:
         self._tags: Dict[Union[str, TagType], List[str]] = {}
         self._cached_tag_list: Optional[tuple[str, ...]] = None
         self._keyless: TagType = TagType.KEYLESS
+        self._normalizer = normalizer
 
-    def set_tag(self, key: Optional[str], value: str, replace: bool = False) -> None:
+    def set_tag(self, key: Optional[str], value: str, replace: bool = False, normalize: bool = False) -> None:
         """
         Set a tag with the given key and value.
         If key is None or empty, the value is stored as a keyless tag.
@@ -485,7 +486,11 @@ def set_tag(self, key: Optional[str], value: str, replace: bool = False) -> None
             value (str): The tag value
             replace (bool): If True, replaces all existing values for this key
                            If False, appends the value if it doesn't exist
+            normalize (bool): If True, applies tag normalization using the configured normalizer
         """
+        if normalize and self._normalizer:
+            value = self._normalizer(value)
+
         if not key:
             key = self._keyless
 
@@ -498,13 +503,14 @@ def set_tag(self, key: Optional[str], value: str, replace: bool = False) -> None
             # Invalidate the cache since tags have changed
             self._cached_tag_list = None
 
-    def set_tags_from_list(self, tag_list: List[str], replace: bool = False) -> None:
+    def set_tags_from_list(self, tag_list: List[str], replace: bool = False, normalize: bool = False) -> None:
         """
         Set multiple tags from a list of strings.
         Strings can be in "key:value" format or just "value" format.
         Args:
             tag_list (List[str]): List of tags in "key:value" format or just "value"
             replace (bool): If True, replaces all existing tags with the new tags list
+            normalize (bool): If True, applies tag normalization using the configured normalizer
         """
         if replace:
             self._tags.clear()
@@ -513,21 +519,25 @@ def set_tags_from_list(self, tag_list: List[str], replace: bool = False) -> None
         for tag in tag_list:
             if ':' in tag:
                 key, value = tag.split(':', 1)
-                self.set_tag(key, value)
+                self.set_tag(key, value, normalize=normalize)
             else:
-                self.set_tag(None, tag)
+                self.set_tag(None, tag, normalize=normalize)
 
-    def delete_tag(self, key: Optional[str], value: Optional[str] = None) -> bool:
+    def delete_tag(self, key: Optional[str], value: Optional[str] = None, normalize: bool = False) -> bool:
         """
         Delete a tag or specific value for a tag.
         For keyless tags, use None or empty string as the key.
         Args:
             key (str): The tag key to delete, or None/empty for keyless tags
             value (str, optional): If provided, only deletes this specific value for the key.
                                  If None, deletes all values for the key.
+            normalize (bool): If True, applies tag normalization to the value for lookup
         Returns:
             bool: True if something was deleted, False otherwise
         """
+        if normalize and self._normalizer and value:
+            value = self._normalizer(value)
+
         if not key:
             key = self._keyless
 

datadog_checks_base/tests/base/utils/db/test_util.py

@@ -620,3 +620,183 @@ def test_get_tags_mutation(self):
         new_tags = tag_manager.get_tags()
         assert new_tags == ['test_key:test_value']
         assert new_tags != tags  # The lists should be different objects
+
+    # Normalization tests
+    def mock_tag_normalizer(self, tag):
+        """Mock normalizer that replaces spaces and hyphens with underscores and lowercases"""
+        return tag.replace(' ', '_').replace('-', '_').lower()
+
+    def test_init_with_normalizer(self):
+        """Test initialization of TagManager with normalizer"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+        assert tag_manager._normalizer == self.mock_tag_normalizer
+        assert tag_manager._tags == {}
+        assert tag_manager._cached_tag_list is None
+
+    def test_set_tag_with_normalization(self):
+        """Test setting tags with normalization enabled"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Test with normalize=True - only value should be normalized, not key
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=True)
+        assert tag_manager._tags == {'test-key': ['value_with_spaces']}
+
+        # Test keyless tag normalization
+        tag_manager.set_tag(None, 'keyless value', normalize=True)
+        assert TagType.KEYLESS in tag_manager._tags
+        assert tag_manager._tags[TagType.KEYLESS] == ['keyless_value']
+
+    def test_set_tag_without_normalization(self):
+        """Test setting tags without normalization"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Test with normalize=False
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=False)
+        assert tag_manager._tags == {'test-key': ['value with spaces']}
+
+        # Test with no normalize parameter (should default to False)
+        tag_manager.set_tag('another-key', 'another value')
+        assert tag_manager._tags == {'test-key': ['value with spaces'], 'another-key': ['another value']}
+
+    def test_set_tags_from_list_with_normalization(self):
+        """Test setting tags from list with normalization enabled"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        tag_list = ['env:prod-test', 'service:web app', 'keyless-tag']
+        tag_manager.set_tags_from_list(tag_list, normalize=True)
+
+        expected_tags = sorted(['env:prod_test', 'service:web_app', 'keyless_tag'])
+        assert sorted(tag_manager.get_tags()) == expected_tags
+
+    def test_set_tags_from_list_without_normalization(self):
+        """Test setting tags from list without normalization"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        tag_list = ['env:prod-test', 'service:web app', 'keyless-tag']
+        tag_manager.set_tags_from_list(tag_list, normalize=False)
+
+        expected_tags = sorted(['env:prod-test', 'service:web app', 'keyless-tag'])
+        assert sorted(tag_manager.get_tags()) == expected_tags
+
+    def test_delete_tag_with_normalization(self):
+        """Test deleting tags with normalization enabled"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Set tag with normalization (only value normalized)
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=True)
+        assert 'test-key' in tag_manager._tags
+        assert tag_manager._tags['test-key'] == ['value_with_spaces']
+
+        # Delete with normalization - should find the normalized value
+        result = tag_manager.delete_tag('test-key', 'value with spaces', normalize=True)
+        assert result is True
+        assert tag_manager._tags == {}
+
+    def test_delete_tag_without_normalization(self):
+        """Test deleting tags without normalization"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Set tag without normalization
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=False)
+        assert tag_manager._tags == {'test-key': ['value with spaces']}
+
+        # Delete without normalization
+        result = tag_manager.delete_tag('test-key', 'value with spaces', normalize=False)
+        assert result is True
+        assert tag_manager._tags == {}
+
+    def test_delete_tag_normalization_mismatch(self):
+        """Test that deletion fails when normalization doesn't match storage"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Set tag with normalization (only value normalized)
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=True)
+        assert tag_manager._tags == {'test-key': ['value_with_spaces']}
+
+        # Try to delete without normalization - should fail
+        result = tag_manager.delete_tag('test-key', 'value with spaces', normalize=False)
+        assert result is False
+        assert tag_manager._tags == {'test-key': ['value_with_spaces']}
+
+        # Delete with normalization - should succeed
+        result = tag_manager.delete_tag('test-key', 'value with spaces', normalize=True)
+        assert result is True
+        assert tag_manager._tags == {}
+
+    def test_delete_entire_key_with_normalization(self):
+        """Test deleting entire key with normalization"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Set multiple values for a key with normalization (only values normalized)
+        tag_manager.set_tag('test-key', 'value1', normalize=True)
+        tag_manager.set_tag('test-key', 'value2', normalize=True)
+        assert tag_manager._tags == {'test-key': ['value1', 'value2']}
+
+        # Delete entire key - key doesn't need normalization
+        result = tag_manager.delete_tag('test-key', normalize=True)
+        assert result is True
+        assert tag_manager._tags == {}
+
+    def test_normalization_with_no_normalizer(self):
+        """Test that normalize parameter is ignored when no normalizer is provided"""
+        tag_manager = TagManager()  # No normalizer
+
+        # Should work the same regardless of normalize parameter
+        tag_manager.set_tag('test-key', 'value with spaces', normalize=True)
+        assert tag_manager._tags == {'test-key': ['value with spaces']}
+
+        tag_manager.set_tags_from_list(['env:prod-test'], normalize=True)
+        expected_tags = sorted(['test-key:value with spaces', 'env:prod-test'])
+        assert sorted(tag_manager.get_tags()) == expected_tags
+
+        # Delete should also work
+        result = tag_manager.delete_tag('test-key', 'value with spaces', normalize=True)
+        assert result is True
+
+    def test_case_sensitivity_normalization(self):
+        """Test that normalization handles case sensitivity correctly (only values normalized)"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Test only value gets normalized, key assumed to be lowercase already
+        tag_manager.set_tag('test_key', 'UPPERCASE-VALUE', normalize=True)
+        assert tag_manager._tags == {'test_key': ['uppercase_value']}
+
+        # Test mixed case value normalization
+        tag_manager.set_tag('another_key', 'SoMe VaLuE', normalize=True)
+        assert 'another_key' in tag_manager._tags
+        assert tag_manager._tags['another_key'] == ['some_value']
+
+        # Verify final tags - keys lowercase, values normalized
+        expected_tags = sorted(['test_key:uppercase_value', 'another_key:some_value'])
+        assert sorted(tag_manager.get_tags()) == expected_tags
+
+        # Test deletion with case sensitivity
+        result = tag_manager.delete_tag('test_key', 'UPPERCASE-VALUE', normalize=True)
+        assert result is True
+        assert tag_manager._tags == {'another_key': ['some_value']}
+
+    def test_case_sensitivity_in_tag_list(self):
+        """Test case sensitivity normalization in tag lists"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        tag_list = ['env:PRODUCTION', 'service:WEB-APP', 'datacenter:US-EAST-1', 'KEYLESS-TAG-UPPERCASE']
+        tag_manager.set_tags_from_list(tag_list, normalize=True)
+
+        # When normalize_tag is applied to the full tag string, it normalizes everything
+        expected_tags = sorted(['env:production', 'service:web_app', 'datacenter:us_east_1', 'keyless_tag_uppercase'])
+        assert sorted(tag_manager.get_tags()) == expected_tags
+
+    def test_case_sensitivity_without_normalization(self):
+        """Test that case sensitivity is preserved without normalization"""
+        tag_manager = TagManager(normalizer=self.mock_tag_normalizer)
+
+        # Set tags without normalization - case should be preserved
+        tag_manager.set_tag('test_key', 'UPPERCASE-VALUE', normalize=False)
+        assert tag_manager._tags == {'test_key': ['UPPERCASE-VALUE']}
+
+        # Set tag list without normalization - case should be preserved
+        tag_list = ['env:PRODUCTION', 'KEYLESS-TAG-UPPERCASE']
+        tag_manager.set_tags_from_list(tag_list, normalize=False)
+
+        expected_tags = sorted(['test_key:UPPERCASE-VALUE', 'env:PRODUCTION', 'KEYLESS-TAG-UPPERCASE'])
+        assert sorted(tag_manager.get_tags()) == expected_tags