Allow tag aliases for wmi tag_by and tag_queries parameters (#21792)

* adding alias to tag_queries

* Adding tag_by_prefix and test

* updating tag_by alias

* removing version changes

* fixing linter erors

* adding changelogs

* Updating tag_queries alias

* Updating changelogs

* fxing lint errors

* updating alias parsing

* updating alias parsing

* move alias parsing to init

* fixing lint errors

* changes to parse_alias

* removing .lower() calls

Author: mrafi97

datadog_checks_base/changelog.d/21792.added

@@ -0,0 +1 @@
+[wmi_check] Allow tag aliases for wmi `tag_by` and `tag_queries` parameters

datadog_checks_base/datadog_checks/base/checks/win/wmi/base.py

@@ -110,8 +110,8 @@ def _raise_on_invalid_tag_query_result(self, sampler, wmi_obj, tag_query):
             )
             raise TypeError
 
-    def _get_tag_query_tag(self, sampler, wmi_obj, tag_query):
-        # type: (WMISampler, WMIObject, TagQuery) -> str
+    def _get_tag_query_tag(self, sampler, wmi_obj, tag_query, alias):
+        # type: (WMISampler, WMIObject, TagQuery, str) -> str
         """
         Design a query based on the given WMIObject to extract a tag.
 
@@ -133,13 +133,14 @@ def _get_tag_query_tag(self, sampler, wmi_obj, tag_query):
 
             link_value = str(tag_query_sampler[0][target_property]).lower()
 
-        tag = "{tag_name}:{tag_value}".format(tag_name=target_property.lower(), tag_value="_".join(link_value.split()))
+        tag_name = alias.lower()
+        tag = "{tag_name}:{tag_value}".format(tag_name=tag_name, tag_value="_".join(link_value.split()))
 
         self.log.debug("Extracted `tag_queries` tag: '%s'", tag)
         return tag
 
-    def _extract_metrics(self, wmi_sampler, tag_by, tag_queries, constant_tags):
-        # type: (WMISampler, str, List[List[str]], List[str]) -> List[WMIMetric]
+    def _extract_metrics(self, wmi_sampler, tag_by, tag_queries, constant_tags, tag_by_aliases, tag_queries_aliases):
+        # type: (WMISampler, str, List[List[str]], List[str], Dict[str, str], List[str]) -> List[WMIMetric]
         """
         Extract and tag metrics from the WMISampler.
 
@@ -170,9 +171,9 @@ def _extract_metrics(self, wmi_sampler, tag_by, tag_queries, constant_tags):
             tags = list(constant_tags) if constant_tags else []
 
             # Tag with `tag_queries` parameter
-            for query in tag_queries:
+            for index, query in enumerate(tag_queries):
                 try:
-                    tags.append(self._get_tag_query_tag(wmi_sampler, wmi_obj, query))
+                    tags.append(self._get_tag_query_tag(wmi_sampler, wmi_obj, query, tag_queries_aliases[index]))
                 except TagQueryUniquenessFailure:
                     continue
 
@@ -193,17 +194,17 @@ def _extract_metrics(self, wmi_sampler, tag_by, tag_queries, constant_tags):
                     continue
                 # Tag with `tag_by` parameter
                 for t in tag_by.split(','):
-                    t = t.strip()
-                    if wmi_property == t:
+                    if normalized_wmi_property == t:
                         tag_value = str(wmi_value).lower()
                         if tag_queries and tag_value.find("#") > 0:
                             tag_value = tag_value[: tag_value.find("#")]
 
-                        tags.append("{name}:{value}".format(name=t, value=tag_value))
+                        alias = tag_by_aliases.get(t)
+                        tags.append("{name}:{value}".format(name=alias, value=tag_value))
                         continue
 
                 # No metric extraction on 'Name' and properties in tag_by
-                if wmi_property == 'name' or normalized_wmi_property in tag_by.lower():
+                if normalized_wmi_property == 'name' or normalized_wmi_property in tag_by.lower():
                     continue
 
                 try:
@@ -260,6 +261,7 @@ def _get_instance_key(self, host, namespace, wmi_class, other=None):
     def get_running_wmi_sampler(self, properties, filters, **kwargs):
         # type: (List[str], List[Dict[str, WMIFilter]], **Any) -> WMISampler
         tag_by = kwargs.pop('tag_by', "")
+
         return self._get_running_wmi_sampler(
             instance_key=None,
             wmi_class=self.wmi_class,
@@ -303,6 +305,49 @@ def _get_wmi_properties(self, instance_key, metrics, tag_queries):
 
         return self._wmi_props
 
+    def parse_tag_queries_aliases(self, tag_queries):
+        # type: (List[TagQuery]) -> None
+        """
+        Validate tag_queries configuration to ensure aliases are provided when 'AS' is used.
+        return parsed_tag_queries and aliases
+        """
+        aliases = []  # type: list[str]
+        parsed_tag_queries = []  # type: list[TagQuery]
+        for tag_query in tag_queries:
+            if len(tag_query) < 4:
+                continue
+            target_property_str = tag_query[3]
+            property, alias = self.parse_alias(target_property_str)
+            aliases.append(alias)
+            tag_query[3] = property
+            parsed_tag_queries.append(tag_query)
+        return parsed_tag_queries, aliases
+
+    def parse_alias(self, property):
+        # type: (str) -> Tuple[str, Optional[str]]
+        """
+        Parse an alias from a string.
+        """
+        property = property.strip().lower()
+        if ' as ' in property:
+            # Valid format: property AS alias (with spaces around AS)
+            property_split = property.split(' as ')
+            property = property_split[0].strip()
+            alias = property_split[1].strip()
+            self.log.debug("Parsed alias: {%s} for property: {%s}", alias, property)
+            if alias == "":
+                self.log.warning("No alias provided after 'AS' for property: %s. Using property for tag", property)
+                alias = property
+        elif ' as' in property:
+            # Invalid format: AS found but without proper spacing
+            raise InvalidWMIQuery(
+                "Invalid alias syntax in property '{}'. "
+                "Expected format: 'property AS alias' with spaces around 'AS'".format(property)
+            )
+        else:
+            alias = property
+        return property, alias
+
 
 def from_time(
     year=None, month=None, day=None, hours=None, minutes=None, seconds=None, microseconds=None, timezone=None

wmi_check/assets/configuration/spec.yaml

@@ -143,21 +143,24 @@ files:
         The `tag_by` parameter lets you tag each metric with a property from the WMI class you're using.
         This is only useful when you will have multiple values for your WMI query.
         Comma-separated list of property names
+        Add aliases to the property tags with by appending AS <ALIAS_NAME> to the property name.
+        For example: Name AS wmi_name will be tagged as wmi_name:value instead of Name:value.
       value:
         type: string
         display_default: null
-        example: Name,Label
+        example: Name AS wmi_name,Label
     - name: tag_queries
       description: |
         The `tag_queries`  parameter lets you specify a list of queries, to tag metrics with a target class property.
         Each item in the list is a set of :
 
-        `[<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY>]`
+        `[<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY> AS <TAG_ALIAS>]`
 
           * `<LINK_SOURCE_PROPERTY>` contains the link value
           * `<TARGET_CLASS>` is the class to link to
           * `<LINK_TARGET_CLASS_PROPERTY>` is the target class property to link to
           * `<TARGET_PROPERTY>` contains the value to tag with
+          * `<TAG_ALIAS>`is the alias to use for the tag. If not provided, the target property's name will be used.
 
         It translates to a WMI query:
 
@@ -170,6 +173,6 @@ files:
             type: string
         compact_example: true
         example:
-          - [<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY>]
+          - [<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY> AS <TAG_ALIAS>]
 
     - template: instances/default

wmi_check/changelog.d/21792.added

@@ -0,0 +1 @@
+Allow tag aliases for wmi `tag_by` and `tag_queries` parameters

wmi_check/datadog_checks/wmi_check/data/conf.yaml.example

@@ -132,26 +132,29 @@ instances:
     ## The `tag_by` parameter lets you tag each metric with a property from the WMI class you're using.
     ## This is only useful when you will have multiple values for your WMI query.
     ## Comma-separated list of property names
+    ## Add aliases to the property tags with by appending AS <ALIAS_NAME> to the property name.
+    ## For example: Name AS wmi_name will be tagged as wmi_name:value instead of Name:value.
     #
-    # tag_by: Name,Label
+    # tag_by: Name AS wmi_name,Label
 
     ## @param tag_queries - list of lists - optional
     ## The `tag_queries`  parameter lets you specify a list of queries, to tag metrics with a target class property.
     ## Each item in the list is a set of :
     ##
-    ## `[<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY>]`
+    ## `[<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY> AS <TAG_ALIAS>]`
     ##
     ##   * `<LINK_SOURCE_PROPERTY>` contains the link value
     ##   * `<TARGET_CLASS>` is the class to link to
     ##   * `<LINK_TARGET_CLASS_PROPERTY>` is the target class property to link to
     ##   * `<TARGET_PROPERTY>` contains the value to tag with
+    ##   * `<TAG_ALIAS>`is the alias to use for the tag. If not provided, the target property's name will be used.
     ##
     ## It translates to a WMI query:
     ##
     ## SELECT '<TARGET_PROPERTY>' FROM '<TARGET_CLASS>' WHERE '<LINK_TARGET_CLASS_PROPERTY>' = '<LINK_SOURCE_PROPERTY>'
     #
     # tag_queries:
-    #   - [<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY>]
+    #   - [<LINK_SOURCE_PROPERTY>, <TARGET_CLASS>, <LINK_TARGET_CLASS_PROPERTY>, <TARGET_PROPERTY> AS <TAG_ALIAS>]
 
     ## @param tags - list of strings - optional
     ## A list of tags to attach to every metric and service check emitted by this instance.

wmi_check/datadog_checks/wmi_check/wmi_check.py

@@ -23,12 +23,23 @@ def __init__(self, name, init_config, instances):
         self.tag_by = self.instance.get('tag_by', "")  # type: str
         self.tag_queries = self.instance.get('tag_queries', [])  # type: List[TagQuery]
 
+        # Parse the tag_queries and validate the aliases
+        self.parsed_tag_queries, self.tag_queries_aliases = self.parse_tag_queries_aliases(self.tag_queries)
+
         custom_tags = self.instance.get('tags', [])  # type: List[str]
         self.constant_tags = self.instance.get('constant_tags', [])  # type: List[str]
         if self.constant_tags:
             self.log.warning("`constant_tags` is being deprecated, please use `tags`")
         self.constant_tags.extend(custom_tags)
 
+        self.tag_by_properties = ""
+        self.tag_by_aliases = {}  # type: Dict[str, str]
+        for t in self.tag_by.split(','):
+            property, alias = self.parse_alias(t)
+            self.tag_by_properties += property + ","
+            self.tag_by_aliases[property.lower()] = alias.lower()
+        self.tag_by_properties = self.tag_by_properties.rstrip(',')
+
     def check(self, _):
         # type: (Any) -> None
         """
@@ -37,7 +48,7 @@ def check(self, _):
         # Create or retrieve an existing WMISampler
         metric_name_and_type_by_property, properties = self.get_wmi_properties()
 
-        wmi_sampler = self.get_running_wmi_sampler(properties, self.filters, tag_by=self.tag_by)
+        wmi_sampler = self.get_running_wmi_sampler(properties, self.filters, tag_by=self.tag_by_properties)
 
         # Sample, extract & submit metrics
         try:
@@ -56,8 +67,15 @@ def check(self, _):
 
     def extract_metrics(self, wmi_sampler):
         # type: (WMISampler) -> List[WMIMetric]
-        return self._extract_metrics(wmi_sampler, self.tag_by, self.tag_queries, self.constant_tags)
+        return self._extract_metrics(
+            wmi_sampler,
+            self.tag_by_properties,
+            self.parsed_tag_queries,
+            self.constant_tags,
+            self.tag_by_aliases,
+            self.tag_queries_aliases,
+        )
 
     def get_wmi_properties(self):
         # type: () -> WMIProperties
-        return self._get_wmi_properties(None, self.metrics_to_capture, self.tag_queries)
+        return self._get_wmi_properties(None, self.metrics_to_capture, self.parsed_tag_queries)

wmi_check/tests/conftest.py

@@ -2,9 +2,12 @@
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
 
+from unittest.mock import MagicMock
+
 import pytest
 from mock import patch
 
+from datadog_checks.base.checks.win.wmi.sampler import CaseInsensitiveDict
 from datadog_checks.wmi_check import WMICheck
 
 
@@ -21,7 +24,15 @@ def __init__(self, wmi_objects=None, properties=None, filters=None):
             properties = []
 
         self._wmi_objects = []
-        self._mock_wmi_objects = wmi_objects
+        # Convert regular dicts to CaseInsensitiveDict to match real WMISampler behavior
+        # Must iterate through items to ensure keys are properly lowercased
+        converted_objects = []
+        for obj in wmi_objects:
+            case_insensitive_obj = CaseInsensitiveDict()
+            for key, value in obj.items():
+                case_insensitive_obj[key] = value
+            converted_objects.append(case_insensitive_obj)
+        self._mock_wmi_objects = converted_objects
         self.property_names = properties
         self._filters = []
 
@@ -73,3 +84,63 @@ def mock_disk_sampler():
 
     with patch("datadog_checks.wmi_check.WMICheck._get_running_wmi_sampler", return_value=sampler):
         yield
+
+
+@pytest.fixture
+def mock_sampler_with_tag_queries():
+    # Main sampler with IDProcess for tag queries
+    main_wmi_objects = [
+        {
+            "IOReadBytesPerSec": 20455,
+            "IDProcess": 1234,
+            "ThreadCount": 4,
+            "VirtualBytes": 3811,
+            "PercentProcessorTime": 5,
+        }
+    ]
+    main_property_names = ["ThreadCount", "IOReadBytesPerSec", "VirtualBytes", "PercentProcessorTime", "IDProcess"]
+    main_sampler = MockSampler(main_wmi_objects, main_property_names)
+    main_sampler.class_name = 'Win32_PerfFormattedData_PerfProc_Process'
+
+    # Tag query sampler for process names
+    tag_wmi_objects = [{'Name': 'chrome.exe'}]
+    tag_property_names = ['Name']
+    tag_sampler = MockSampler(tag_wmi_objects, tag_property_names)
+    tag_sampler.class_name = 'Win32_Process'
+    tag_sampler.sample()  # Populate the mock data
+
+    with patch("datadog_checks.wmi_check.WMICheck._get_running_wmi_sampler", return_value=main_sampler):
+        with patch("datadog_checks.base.checks.win.wmi.base.WMISampler") as mock_wmi_sampler:
+            # Setup context manager to return tag_sampler for tag queries
+            mock_wmi_sampler.return_value.__enter__ = MagicMock(return_value=tag_sampler)
+            mock_wmi_sampler.return_value.__exit__ = MagicMock(return_value=False)
+            yield
+
+
+@pytest.fixture
+def mock_sampler_with_tag_by_alias():
+    main_wmi_objects = [
+        {
+            "IOReadBytesPerSec": 20455,
+            "IDProcess": 1234,
+            "Name": "foo",
+            "ThreadCount": 4,
+            "VirtualBytes": 3811,
+            "PercentProcessorTime": 5,
+            "Label": "bar",
+        }
+    ]
+    main_property_names = [
+        "ThreadCount",
+        "IOReadBytesPerSec",
+        "VirtualBytes",
+        "PercentProcessorTime",
+        "IDProcess",
+        "Name",
+        "Label",
+    ]
+    main_sampler = MockSampler(main_wmi_objects, main_property_names)
+    main_sampler.class_name = 'Win32_PerfFormattedData_PerfProc_Process'
+
+    with patch("datadog_checks.wmi_check.WMICheck._get_running_wmi_sampler", return_value=main_sampler):
+        yield