[wmi] friendly com_errors 💑

Maintain a list of "known" WMI `com_errors`, and translate them to user
intelligible errors.
Add a `mute` parameter to `WMISampler`.
* If set to `True` (default for system checks), `com_error` are
  "translated" (when possible), logged and do not cause any
  interruption.
* If not (user checks), "translated" exceptions raise.

Author: Yann MAHE

checks/libs/wmi/exceptions.py

@@ -0,0 +1,49 @@
+"""
+List WMI "known" `com_errors` errors.
+
+Translate to user intelligible exceptions.
+"""
+
+_user_exception_by_com_errors = {}
+
+
+def com_error(error_id):
+    """
+    A decorator that assigns an `error_id` to an intelligible exception.
+    """
+    def set_exception(exception):
+        _user_exception_by_com_errors[error_id] = exception
+        return exception
+    return set_exception
+
+
+def raise_on_com_error(error, silent=True):
+    """
+    Raise the user exception associated with the given `com_error` or
+    fall back to the original exception.
+    """
+    raise _user_exception_by_com_errors.get(error[0], error)
+
+
+# List of user exceptions
+class WMIException(Exception):
+    """
+    Abtract exception for WMI.
+    """
+    def __init__(self):
+        """
+        Use the class docstring as an exception message.
+        """
+        super(WMIException, self).__init__(self.__doc__)
+
+
+@com_error(-2147352567)
+class WMIProviderNotSupported(WMIException):
+    """WMI `provider` option is not supported on the system."""
+    pass
+
+
+@com_error(-2147217392)
+class WMIInvalidClass(WMIException):
+    """WMI class is invalid."""
+    pass

checks/libs/wmi/sampler.py

@@ -30,6 +30,7 @@
 
 # project
 from checks.libs.wmi.counter_type import get_calculator, get_raw, UndefinedCalculator
+from checks.libs.wmi.exceptions import raise_on_com_error, WMIException
 from utils.timeout import timeout, TimeoutException
 
 
@@ -91,8 +92,10 @@ class WMISampler(object):
 
     def __init__(self, logger, class_name, property_names, filters="", host="localhost",
                  namespace="root\\cimv2", provider=None,
-                 username="", password="", and_props=[], timeout_duration=10):
+                 username="", password="", and_props=[],
+                 mute=True, timeout_duration=10):
         self.logger = logger
+        self.mute = mute
 
         # Connection information
         self.host = host
@@ -424,13 +427,12 @@ def build_where_clause(fltr):
                 more=build_where_clause(fltr)
             )
 
-
         if not filters:
             return ""
 
         return " WHERE {clause}".format(clause=build_where_clause(filters))
 
-    def _query(self): # pylint: disable=E0202
+    def _query(self):  # pylint: disable=E0202
         """
         Query WMI using WMI Query Language (WQL) & parse the results.
 
@@ -444,31 +446,50 @@ def _query(self): # pylint: disable=E0202
         )
         self.logger.debug(u"Querying WMI: {0}".format(wql))
 
-        try:
-            # From: https://msdn.microsoft.com/en-us/library/aa393866(v=vs.85).aspx
-            flag_return_immediately = 0x10  # Default flag.
-            flag_forward_only = 0x20
-            flag_use_amended_qualifiers = 0x20000
+        # From: https://msdn.microsoft.com/en-us/library/aa393866(v=vs.85).aspx
+        flag_return_immediately = 0x10  # Default flag.
+        flag_forward_only = 0x20
+        flag_use_amended_qualifiers = 0x20000
 
-            query_flags = flag_return_immediately | flag_forward_only
+        query_flags = flag_return_immediately | flag_forward_only
 
-            # For the first query, cache the qualifiers to determine each
-            # propertie's "CounterType"
-            includes_qualifiers = self.is_raw_perf_class and self._property_counter_types is None
-            if includes_qualifiers:
-                self._property_counter_types = CaseInsensitiveDict()
-                query_flags |= flag_use_amended_qualifiers
+        # For the first query, cache the qualifiers to determine each
+        # propertie's "CounterType"
+        includes_qualifiers = self.is_raw_perf_class and self._property_counter_types is None
+        if includes_qualifiers:
+            self._property_counter_types = CaseInsensitiveDict()
+            query_flags |= flag_use_amended_qualifiers
 
+        try:
             raw_results = self.get_connection().ExecQuery(wql, "WQL", query_flags)
-
-            results = self._parse_results(raw_results, includes_qualifiers=includes_qualifiers)
-
-        except pywintypes.com_error:
-            self.logger.warning(u"Failed to execute WMI query (%s)", wql, exc_info=True)
+        except pywintypes.com_error as e:
+            self._handle_com_error(e, wql)
             results = []
+        else:
+            results = self._parse_results(raw_results, includes_qualifiers=includes_qualifiers)
 
         return results
 
+    def _handle_com_error(self, error, wql):
+        """
+        Attempt to translate the WMI `com_error` to something intelligible.
+        Raise when needed or log a warning.
+        """
+        warning_template = u"Failed to execute WMI query ({}).".format(wql)
+
+        try:
+            raise_on_com_error(error)
+
+        # Translate to user exceptions
+        except WMIException as e:
+            if not self.mute:
+                raise
+            self.logger.warning(u"%s Reason:%s", warning_template, e.message)
+
+        # Unknown exceptions
+        except Exception:
+            self.logger.exception(warning_template)
+
     def _parse_results(self, raw_results, includes_qualifiers):
         """
         Parse WMI query results in a more comprehensive form.

tests/core/test_wmi.py

@@ -9,7 +9,9 @@
 from mock import Mock, patch
 
 # project
+from checks.libs.wmi.exceptions import WMIInvalidClass
 from tests.checks.common import Fixtures
+from tests.core.test_wmi_exceptions import MockedWMICOMError
 from utils.timeout import TimeoutException
 
 
@@ -349,7 +351,34 @@ def assertInPartial(self, first, second):
 
         Note: needs to be defined for Python 2.6
         """
-        self.assertTrue(any(key for key in second if key.startswith(first)), "{0} not in {1}".format(first, second))
+        self.assertTrue(
+            any(key for key in second if key.startswith(first)),
+            "{0} not in {1}".format(first, second)
+        )
+
+    def _assertLogger(self, logger, message, level):
+        """
+        Assert that the logger was called with the given level and submessage.
+        """
+        logger_method = getattr(logger, level)
+
+        # Logger was called with the given level
+        self.assertTrue(logger_method.called)
+
+        # Submessage was logged
+        self.assertTrue([s for s in logger_method.call_args[0] if message in s])
+
+    def assertWarning(self, *args):
+        """
+        Assert logging with the WARNING level.
+        """
+        self._assertLogger(*args, level="warning")
+
+    def assertException(self, *args):
+        """
+        Assert logging with the EXCEPTION level.
+        """
+        self._assertLogger(*args, level="exception")
 
     def getProp(self, dict, prefix):
         """
@@ -788,6 +817,54 @@ def test_missing_property(self):
 
         self.assertWMISampler(wmi_raw_sampler, ["MissingProperty"], count=1)
 
+    def test_warnings_on_com_errors(self):
+        """
+        Log a warning on WMI `com_errors`.
+        """
+        # Mute to WMI Sampler
+        from checks.libs.wmi.sampler import WMISampler
+        logger = Mock()
+        wmi_sampler = WMISampler(logger, "WMI_Class", ["Property"], mute=True)
+
+        # Samples
+        invalid_class_com_error = MockedWMICOMError(-2147217392)
+        unknown_com_error = MockedWMICOMError(123456)
+
+        # Method to test
+        log_on_com_errors = partial(wmi_sampler._handle_com_error, wql="")
+
+        # Log WARNING an intelligible when possible
+        log_on_com_errors(invalid_class_com_error)
+        self.assertWarning(logger, "class is invalid")
+
+        # Or log ERROR the exception trace
+        log_on_com_errors(unknown_com_error)
+        self.assertException(logger, "Failed to execute WMI query")
+
+    def test_raise_on_com_errors(self):
+        """
+        Log a warning on WMI `com_errors`.
+        """
+        # Do not mute to WMI Sampler
+        from checks.libs.wmi.sampler import WMISampler
+        logger = Mock()
+        wmi_sampler = WMISampler(logger, "WMI_Class", ["Property"], mute=False)
+
+        # Samples
+        invalid_class_com_error = MockedWMICOMError(-2147217392)
+        unknown_com_error = MockedWMICOMError(123456)
+
+        # Method to test
+        raise_on_com_errors = partial(wmi_sampler._handle_com_error, wql="")
+
+        # Raise known exceptions
+        with self.assertRaises(WMIInvalidClass):
+            raise_on_com_errors(invalid_class_com_error)
+
+        # Log exceptions the others
+        raise_on_com_errors(unknown_com_error)
+        self.assertException(logger, "Failed to execute WMI query")
+
 
 class TestIntegrationWMI(unittest.TestCase):
     """

tests/core/test_wmi_exceptions.py

@@ -0,0 +1,55 @@
+# stdlib
+import unittest
+
+# datadog
+from checks.libs.wmi.exceptions import com_error, raise_on_com_error
+
+
+class MockedWMICOMError(Exception):
+    """
+    Mocking a WMI `com_error` error.
+    """
+    def __init__(self, error_id):
+        super(MockedWMICOMError, self).__init__()
+        self.error_id = error_id
+
+    def __getitem__(self, index):
+        if index == 0:
+            return self.error_id
+
+        raise NotImplementedError
+
+
+class TestWMIExceptions(unittest.TestCase):
+    """
+    Unit testing for WMI `com_error` errors and user exceptions.
+    """
+    def test_register_com_error(self):
+        """
+        `com_error` decorator register and map a WMI `com_error` to an user exception.
+        """
+        # Mock a WMI `com_error`
+        sample_error = MockedWMICOMError(123456)
+
+        # Map it to an user exception
+        @com_error(sample_error.error_id)
+        class WMISampleException(Exception):
+            """
+            Sample WMI exception.
+            """
+            pass
+
+        # Assert that the exception is raised
+        with self.assertRaises(WMISampleException):
+            raise_on_com_error(sample_error)
+
+    def test_unregistred_com_error(self):
+        """
+        Unknown `com_error` errors remain intact.
+        """
+        # Mock a WMI `com_error`, do not register it
+        sample_error = MockedWMICOMError(456789)
+
+        # Assert that the original exception is raised
+        with self.assertRaises(MockedWMICOMError):
+            raise_on_com_error(sample_error)