Merge pull request #2011 from DataDog/yann/wmi-new-collector

[core][wmi] new WMI metric collection core

checks.d/wmi_check.py

@@ -1,132 +1,285 @@
-'''
-Windows Only.
-
-Generic WMI check. This check allows you to specify particular metrics that you
-want from WMI in your configuration. Check wmi_check.yaml.example in your conf.d
-directory for more details on configuration.
-'''
-# 3rd party
-import wmi
+# stdlib
+from collections import namedtuple
 
 # project
 from checks import AgentCheck
+from checks.libs.wmi.sampler import WMISampler
+
+WMIMetric = namedtuple('WMIMetric', ['name', 'value', 'tags'])
+
+
+class InvalidWMIQuery(Exception):
+    """
+    Invalid WMI Query.
+    """
+    pass
+
+
+class MissingTagBy(Exception):
+    """
+    WMI query returned multiple rows but no `tag_by` value was given.
+    """
+    pass
 
-UP_METRIC = 'Up'
-SEARCH_WILDCARD = '*'
+
+class TagQueryUniquenessFailure(Exception):
+    """
+    'Tagging query' did not return or returned multiple results.
+    """
+    pass
 
 
 class WMICheck(AgentCheck):
+    """
+    WMI check.
+
+    Windows only.
+    """
     def __init__(self, name, init_config, agentConfig, instances):
         AgentCheck.__init__(self, name, init_config, agentConfig, instances)
-        self.wmi_conns = {}
-
-    def _get_wmi_conn(self, host, **kwargs):
-        key = "{0}:".format(host)
-        key += ":".join(str(v) for v in kwargs.values())
-        if key not in self.wmi_conns:
-            self.wmi_conns[key] = wmi.WMI(host, **kwargs)
-        return self.wmi_conns[key]
+        self.wmi_samplers = {}
+        self.wmi_props = {}
 
     def check(self, instance):
-        host = instance.get('host', None)
-        namespace = instance.get('namespace', None)
-        user = instance.get('username', None)
-        password = instance.get('password', None)
-        w = self._get_wmi_conn(host, namespace=namespace, user=user, password=password)
+        """
+        Fetch WMI metrics.
+        """
+        # Connection information
+        host = instance.get('host', "localhost")
+        namespace = instance.get('namespace', "root\\cimv2")
+        username = instance.get('username', "")
+        password = instance.get('password', "")
+
+        # WMI instance
         wmi_class = instance.get('class')
         metrics = instance.get('metrics')
         filters = instance.get('filters')
-        tag_by = instance.get('tag_by')
-        tag_queries = instance.get('tag_queries')
+        tag_by = instance.get('tag_by', "").lower()
+        tag_queries = instance.get('tag_queries', [])
         constant_tags = instance.get('constant_tags')
 
-        if not wmi_class:
-            raise Exception('WMI instance is missing a value for `class` in wmi_check.yaml')
-
-        # If there are filters, we need one query per filter.
-        if filters:
-            for f in filters:
-                prop = f.keys()[0]
-                search = f.values()[0]
-                if SEARCH_WILDCARD in search:
-                    search = search.replace(SEARCH_WILDCARD, '%')
-                    wql = "SELECT * FROM %s WHERE %s LIKE '%s'" \
-                        % (wmi_class, prop, search)
-                    results = w.query(wql)
-                else:
-                    results = getattr(w, wmi_class)(**f)
-                self._extract_metrics(results, metrics, tag_by, w, tag_queries, constant_tags)
-        else:
-            results = getattr(w, wmi_class)()
-            self._extract_metrics(results, metrics, tag_by, w, tag_queries, constant_tags)
-
-    def _extract_metrics(self, results, metrics, tag_by, wmi, tag_queries, constant_tags):
-        if len(results) > 1 and tag_by is None:
-            raise Exception('WMI query returned multiple rows but no `tag_by` value was given. '
-                            'metrics=%s' % metrics)
-
-        for res in results:
-            tags = []
-
-            # include any constant tags...
-            if constant_tags:
-                tags.extend(constant_tags)
-
-            # if tag_queries is specified then get attributes from other classes and use as a tags
-            if tag_queries:
-                for query in tag_queries:
-                    link_source_property = int(getattr(res, query[0]))
-                    target_class = query[1]
-                    link_target_class_property = query[2]
-                    target_property = query[3]
-
-                    link_results = \
-                        wmi.query("SELECT {0} FROM {1} WHERE {2} = {3}"
-                                  .format(target_property, target_class,
-                                          link_target_class_property, link_source_property))
-
-                    if len(link_results) != 1:
-                        self.log.warning("Failed to find {0} for {1} {2}. No metrics gathered"
-                                         .format(target_class, link_target_class_property,
-                                                 link_source_property))
-                        continue
-
-                    link_value = str(getattr(link_results[0], target_property)).lower()
-                    tags.append("{0}:{1}".format(target_property.lower(),
-                                "_".join(link_value.split())))
-
-            # Grab the tag from the result if there's a `tag_by` value (e.g.: "name:jenkins")
-            # Strip any #instance off the value when `tag_queries` is set (gives us unique tags)
-            if tag_by:
-                tag_value = str(getattr(res, tag_by)).lower()
-                if tag_queries and tag_value.find("#") > 0:
-                    tag_value = tag_value[:tag_value.find("#")]
-                tags.append('%s:%s' % (tag_by.lower(), tag_value))
-
-            if len(tags) == 0:
-                tags = None
-
-            for wmi_property, name, mtype in metrics:
-                if wmi_property == UP_METRIC:
-                    # Special-case metric will just submit 1 for every value
-                    # returned in the result.
-                    val = 1
-                else:
-                    try:
-                        val = float(getattr(res, wmi_property))
-                    except (ValueError, TypeError):
-                        self.log.warning("When extracting metrics with WMI, found a non digit value"
-                                         " for property '{0}'.".format(wmi_property))
-                        continue
-                    except AttributeError:
-                        self.log.warning("'{0}' WMI class has no property '{1}'."
-                                         .format(res.__class__.__name__, wmi_property))
-                        continue
-
-                # Submit the metric to Datadog
+        # Create or retrieve an existing WMISampler
+        instance_key = self._get_instance_key(host, namespace, wmi_class)
+
+        metric_name_and_type_by_property, properties = \
+            self._get_wmi_properties(instance_key, metrics, tag_queries)
+
+        wmi_sampler = self._get_wmi_sampler(
+            instance_key,
+            wmi_class, properties,
+            filters=filters,
+            host=host, namespace=namespace,
+            username=username, password=password
+        )
+
+        # Sample, extract & submit metrics
+        wmi_sampler.sample()
+        metrics = self._extract_metrics(wmi_sampler, tag_by, tag_queries, constant_tags)
+        self._submit_metrics(metrics, metric_name_and_type_by_property)
+
+    def _format_tag_query(self, sampler, wmi_obj, tag_query):
+        """
+        Format `tag_query` or raise on incorrect parameters.
+        """
+        try:
+            link_source_property = int(wmi_obj[tag_query[0]])
+            target_class = tag_query[1]
+            link_target_class_property = tag_query[2]
+            target_property = tag_query[3]
+        except IndexError:
+            self.log.error(
+                u"Wrong `tag_queries` parameter format. "
+                "Please refer to the configuration file for more information.")
+            raise
+        except TypeError:
+            self.log.error(
+                u"Incorrect 'link source property' in `tag_queries` parameter:"
+                " `{wmi_property}` is not a property of `{wmi_class}`".format(
+                    wmi_property=tag_query[0],
+                    wmi_class=sampler.class_name,
+                )
+            )
+            raise
+
+        return target_class, target_property, [{link_target_class_property: link_source_property}]
+
+    def _raise_on_invalid_tag_query_result(self, sampler, wmi_obj, tag_query):
+        """
+        """
+        target_property = sampler.property_names[0]
+        target_class = sampler.class_name
+
+        if len(sampler) != 1:
+            message = "no result was returned"
+            if len(sampler):
+                message = "multiple results returned (one expected)"
+
+            self.log.warning(
+                u"Failed to extract a tag from `tag_queries` parameter: {reason}."
+                " wmi_object={wmi_obj} - query={tag_query}".format(
+                    reason=message,
+                    wmi_obj=wmi_obj, tag_query=tag_query,
+                )
+            )
+            raise TagQueryUniquenessFailure
+
+        if sampler[0][target_property] is None:
+            self.log.error(
+                u"Incorrect 'target property' in `tag_queries` parameter:"
+                " `{wmi_property}` is not a property of `{wmi_class}`".format(
+                    wmi_property=target_property,
+                    wmi_class=target_class,
+                )
+            )
+            raise TypeError
+
+    def _get_tag_query_tag(self, sampler, wmi_obj, tag_query):
+        """
+        Design a query based on the given WMIObject to extract a tag.
+
+        Returns: tag or TagQueryUniquenessFailure exception.
+        """
+        self.log.debug(
+            u"`tag_queries` parameter found."
+            " wmi_object={wmi_obj} - query={tag_query}".format(
+                wmi_obj=wmi_obj, tag_query=tag_query,
+            )
+        )
+
+        # Extract query information
+        target_class, target_property, filters = \
+            self._format_tag_query(sampler, wmi_obj, tag_query)
+
+        # Create a specific sampler
+        connection = sampler.get_connection()
+        tag_query_sampler = WMISampler(
+            self.log,
+            target_class, [target_property],
+            filters=filters,
+            **connection
+        )
+
+        tag_query_sampler.sample()
+
+        # Extract tag
+        self._raise_on_invalid_tag_query_result(tag_query_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())
+        )
+
+        self.log.debug(u"Extracted `tag_queries` tag: '{tag}'".format(tag=tag))
+        return tag
+
+    def _extract_metrics(self, wmi_sampler, tag_by, tag_queries, constant_tags):
+        """
+        Extract and tag metrics from the WMISampler.
+
+        Raise when multiple WMIObject were returned by the sampler with no `tag_by` specified.
+
+        Returns: List of WMIMetric
+        ```
+        [
+            WMIMetric("freemegabytes", 19742, ["name:_total"]),
+            WMIMetric("avgdiskbytesperwrite", 1536, ["name:c:"]),
+        ]
+        ```
+        """
+        if len(wmi_sampler) > 1 and not tag_by:
+            raise MissingTagBy(
+                u"WMI query returned multiple rows but no `tag_by` value was given."
+                " class={wmi_class} - properties={wmi_properties} - filters={filters}".format(
+                    wmi_class=wmi_sampler.class_name, wmi_properties=wmi_sampler.property_names,
+                    filters=wmi_sampler.filters,
+                )
+            )
+
+        metrics = []
+
+        for wmi_obj in wmi_sampler:
+            tags = list(constant_tags) if constant_tags else []
+
+            # Tag with `tag_queries` parameter
+            for query in tag_queries:
+                try:
+                    tags.append(self._get_tag_query_tag(wmi_sampler, wmi_obj, query))
+                except TagQueryUniquenessFailure:
+                    continue
+
+            for wmi_property, wmi_value in wmi_obj.iteritems():
+                # Tag with `tag_by` parameter
+                if wmi_property == tag_by:
+                    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=tag_by.lower(), value=tag_value
+                        )
+                    )
+                    continue
                 try:
-                    func = getattr(self, mtype)
-                except AttributeError:
-                    raise Exception('Invalid metric type: {0}'.format(mtype))
+                    metrics.append(WMIMetric(wmi_property, float(wmi_value), tags))
+                except ValueError:
+                    self.log.warning(u"When extracting metrics with WMI, found a non digit value"
+                                     " for property '{0}'.".format(wmi_property))
+                    continue
+                except TypeError:
+                    self.log.warning(u"When extracting metrics with WMI, found a missing property"
+                                     " '{0}'".format(wmi_property))
+                    continue
+        return metrics
+
+    def _submit_metrics(self, metrics, metric_name_and_type_by_property):
+        """
+        Resolve metric names and types and submit it.
+        """
+        for metric in metrics:
+            if metric.name not in metric_name_and_type_by_property:
+                # Only report the metrics that were specified in the configration
+                # Ignore added properties like 'Timestamp_Sys100NS', `Frequency_Sys100NS`, etc ...
+                continue
+
+            metric_name, metric_type = metric_name_and_type_by_property[metric.name]
+            try:
+                func = getattr(self, metric_type)
+            except AttributeError:
+                raise Exception(u"Invalid metric type: {0}".format(metric_type))
+
+            func(metric_name, metric.value, metric.tags)
+
+    def _get_instance_key(self, host, namespace, wmi_class):
+        """
+        Return an index key for a given instance. Usefull for caching.
+        """
+        return "{host}:{namespace}:{wmi_class}".format(
+            host=host, namespace=namespace, wmi_class=wmi_class,
+        )
+
+    def _get_wmi_sampler(self, instance_key, wmi_class, properties, **kwargs):
+        """
+        Create and cache a WMISampler for the given (class, properties)
+        """
+        if instance_key not in self.wmi_samplers:
+            wmi_sampler = WMISampler(self.log, wmi_class, properties, **kwargs)
+            self.wmi_samplers[instance_key] = wmi_sampler
+
+        return self.wmi_samplers[instance_key]
+
+    def _get_wmi_properties(self, instance_key, metrics, tag_queries):
+        """
+        Create and cache a (metric name, metric type) by WMI property map and a property list.
+        """
+        if instance_key not in self.wmi_props:
+            metric_name_by_property = dict(
+                (wmi_property.lower(), (metric_name, metric_type))
+                for wmi_property, metric_name, metric_type in metrics
+            )
+            properties = map(lambda x: x[0], metrics + tag_queries)
+            self.wmi_props[instance_key] = (metric_name_by_property, properties)
 
-                func(name, val, tags=tags)
+        return self.wmi_props[instance_key]