Implement Python API for setting check metadata

datadog_checks_base/datadog_checks/base/checks/base.py

@@ -21,6 +21,7 @@
 from ..utils.common import ensure_bytes, ensure_unicode, to_string
 from ..utils.http import RequestsWrapper
 from ..utils.limiter import Limiter
+from ..utils.metadata import MetadataManager
 from ..utils.proxy import config_proxy_skip
 
 try:
@@ -73,7 +74,23 @@ class __AgentCheck(object):
 
     OK, WARNING, CRITICAL, UNKNOWN = ServiceCheck
 
-    HTTP_CONFIG_REMAPPER = None  # Used by `self.http` RequestsWrapper
+    # Used by `self.http` for an instance of RequestsWrapper
+    HTTP_CONFIG_REMAPPER = None
+
+    # Used by `self.set_metadata` for an instance of MetadataManager
+    #
+    # This is a mapping of metadata names to functions. When you call `self.set_metadata(name, value, **options)`,
+    # if `name` is in this mapping then the corresponding function will be called with the `value`, and the
+    # return value(s) will be sent instead.
+    #
+    # Transformer functions must satisfy the following signature:
+    #
+    #    def transform_<NAME>(value: Any, options: dict) -> Union[str, Dict[str, str]]:
+    #
+    # If the return type is a string, then it will be sent as the value for `name`. If the return type is
+    # a mapping type, then each key will be considered a `name` and will be sent with its (str) value.
+    METADATA_TRANSFORMERS = None
+
     FIRST_CAP_RE = re.compile(br'(.)([A-Z][a-z]+)')
     ALL_CAP_RE = re.compile(br'([a-z0-9])([A-Z])')
     METRIC_REPLACEMENT = re.compile(br'([^a-zA-Z0-9_.]+)|(^[^a-zA-Z]+)')
@@ -135,6 +152,9 @@ class except the :py:meth:`check` method but sometimes it might be useful for a
         # Only new checks or checks on Agent 6.13+ can and should use this for HTTP requests.
         self._http = None
 
+        # Used for sending metadata via Go bindings
+        self._metadata_manager = None
+
         # Save the dynamically detected integration version
         self._check_version = None
 
@@ -206,6 +226,16 @@ def http(self):
 
         return self._http
 
+    @property
+    def metadata_manager(self):
+        if self._metadata_manager is None:
+            if not self.check_id:
+                raise RuntimeError('Attribute `check_id` must be set')
+
+            self._metadata_manager = MetadataManager(self.name, self.check_id, self.log, self.METADATA_TRANSFORMERS)
+
+        return self._metadata_manager
+
     @property
     def check_version(self):
         if self._check_version is None:
@@ -449,10 +479,20 @@ def _log_deprecation(self, deprecation_key):
             self.log.warning(self._deprecations[deprecation_key][1])
             self._deprecations[deprecation_key][0] = True
 
-    # TODO(olivier): implement service_metadata if it's worth it
+    # TODO: Remove once our checks stop calling it
     def service_metadata(self, meta_name, value):
         pass
 
+    def set_metadata(self, name, value, **options):
+        """Updates the cached metadata ``name`` with ``value``, which is then sent by the Agent at regular intervals.
+
+        :param str name: the name of the metadata
+        :param object value: the value for the metadata. if ``name`` has no transformer defined then the
+                             raw ``value`` will be submitted and therefore it must be a ``str``
+        :param options: keyword arguments to pass to any defined transformer
+        """
+        self.metadata_manager.submit(name, value, options)
+
     def set_external_tags(self, external_tags):
         # Example of external_tags format
         # [

datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py

@@ -31,6 +31,10 @@ def debug(msg, *args, **kwargs):
     pass
 
 
+def set_check_metadata(*args, **kwargs):
+    pass
+
+
 def set_external_tags(*args, **kwargs):
     pass
 

datadog_checks_base/datadog_checks/base/utils/common.py

@@ -5,7 +5,7 @@
 import re
 from decimal import ROUND_HALF_UP, Decimal
 
-from six import PY3, text_type
+from six import PY3, iteritems, text_type
 from six.moves.urllib.parse import urlparse
 
 
@@ -24,6 +24,10 @@ def ensure_unicode(s):
 to_string = ensure_unicode if PY3 else ensure_bytes
 
 
+def exclude_undefined_keys(mapping):
+    return {key: value for key, value in iteritems(mapping) if value is not None}
+
+
 def round_value(value, precision=0, rounding_method=ROUND_HALF_UP):
     precision = '0.{}'.format('0' * precision)
     return float(Decimal(str(value)).quantize(Decimal(precision), rounding=rounding_method))

datadog_checks_base/datadog_checks/base/utils/metadata/__init__.py

@@ -0,0 +1,4 @@
+# (C) Datadog, Inc. 2019
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from .core import MetadataManager

datadog_checks_base/datadog_checks/base/utils/metadata/core.py

@@ -0,0 +1,138 @@
+# (C) Datadog, Inc. 2019
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+import json
+import logging
+
+from six import iteritems
+
+from .utils import is_primitive
+from .version import parse_version
+
+try:
+    import datadog_agent
+except ImportError:
+    from ...stubs import datadog_agent
+
+LOGGER = logging.getLogger(__file__)
+
+
+class MetadataManager(object):
+    __slots__ = ('check_id', 'check_name', 'logger', 'metadata_transformers')
+
+    def __init__(self, check_name, check_id, logger=None, metadata_transformers=None):
+        self.check_name = check_name
+        self.check_id = check_id
+        self.logger = logger or LOGGER
+        self.metadata_transformers = {'config': self.transform_config, 'version': self.transform_version}
+
+        if metadata_transformers:
+            self.metadata_transformers.update(metadata_transformers)
+
+    def submit_raw(self, name, value):
+        datadog_agent.set_check_metadata(self.check_id, name, value)
+
+    def submit(self, name, value, options):
+        transformer = self.metadata_transformers.get(name)
+        if transformer:
+            try:
+                transformed = transformer(value, options)
+            except Exception as e:
+                if is_primitive(value):
+                    self.logger.error('Unable to transform `%s` metadata value `%s`: %s', name, value, e)
+                else:
+                    self.logger.error('Unable to transform `%s` metadata: %s', name, e)
+
+                return
+
+            if isinstance(transformed, str):
+                self.submit_raw(name, transformed)
+            else:
+                for transformed_name, transformed_value in iteritems(transformed):
+                    self.submit_raw(transformed_name, transformed_value)
+        else:
+            self.submit_raw(name, value)
+
+    def transform_version(self, version, options):
+        """Transforms a version like ``1.2.3-rc.4+5`` to its constituent parts. In all cases,
+        the metadata names ``version.raw`` and ``version.scheme`` will be sent.
+
+        If a ``scheme`` is defined then it will be looked up from our known schemes. If no
+        scheme is defined then it will default to semver.
+
+        The scheme may be set to ``regex`` in which case a ``pattern`` must also be defined. Any matching named
+        subgroups will then be sent as ``version.<GROUP_NAME>``. In this case, the check name will be used as
+        the value of ``version.scheme`` unless ``final_scheme`` is also set, which will take precedence.
+        """
+        scheme, version_parts = parse_version(version, options)
+        if scheme == 'regex':
+            scheme = options.get('final_scheme', self.check_name)
+
+        data = {'version.{}'.format(part_name): part_value for part_name, part_value in iteritems(version_parts)}
+        data['version.raw'] = version
+        data['version.scheme'] = scheme
+
+        return data
+
+    def transform_config(self, config, options):
+        """This transforms a ``dict`` of arbitrary user configuration. A ``section`` must be defined indicating
+        what the configuration represents e.g. ``init_config``.
+
+        The metadata name submitted will become ``config.<section>``.
+
+        The value will be a JSON ``str`` with the root being an array. There will be one map element for every
+        allowed field. Every map may have 2 entries:
+
+        1. ``is_set`` - a boolean indicating whether or not the field exists
+        2. ``value`` - the value of the field. this is only set if the field exists and the value is a
+                       primitive type (``None`` | ``bool`` | ``float`` | ``int`` | ``str``)
+
+        The allowed fields are derived from the optional ``whitelist`` and ``blacklist``. By default, nothing
+        will be sent.
+
+        User configuration can override defaults allowing complete, granular control of metadata submissions. In
+        any section, one may set ``metadata_whitelist`` and/or ``metadata_blacklist`` which will override their
+        keyword argument counterparts. In following our standard, blacklists take precedence over whitelists.
+        """
+        section = options.get('section')
+        if section is None:
+            raise ValueError('The `section` option is required')
+
+        # Although we define the default fields to send in code i.e. the default whitelist, there
+        # may be cases where a subclass (for example of OpenMetricsBaseCheck) would want to ignore
+        # just a few fields, hence for convenience we have the ability to also pass a blacklist.
+        whitelist = config.get('metadata_whitelist', options.get('whitelist', []))
+        blacklist = config.get('metadata_blacklist', options.get('blacklist', []))
+        allowed_fields = set(whitelist).difference(blacklist)
+
+        transformed_data = {}
+
+        data = []
+        for field in allowed_fields:
+            field_data = {}
+
+            if field in config:
+                field_data['is_set'] = True
+
+                value = config[field]
+                if is_primitive(value):
+                    field_data['value'] = value
+                else:
+                    self.logger.warning(
+                        'Skipping metadata submission of non-primitive type `%s` for field `%s` in section `%s`',
+                        type(value).__name__,
+                        field,
+                        section,
+                    )
+            else:
+                field_data['is_set'] = False
+
+            data.append(field_data)
+
+        if data:
+            # To avoid the backend having to parse a potentially unbounded number of unique keys, we
+            # send `config.<SECTION_NAME>` rather than `config.<SECTION_NAME>.<OPTION_NAME>` since
+            # the number of sections is finite (currently only `instance` and `init_config`).
+            transformed_data['config.{}'.format(section)] = json.dumps(data)
+
+        return transformed_data

datadog_checks_base/datadog_checks/base/utils/metadata/utils.py

@@ -0,0 +1,8 @@
+# (C) Datadog, Inc. 2019
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+
+
+def is_primitive(obj):
+    # https://github.com/python/cpython/blob/4f82a53c5d34df00bf2d563c2417f5e2638d1004/Lib/json/encoder.py#L357-L377
+    return obj is None or isinstance(obj, (bool, float, int, str))

datadog_checks_base/datadog_checks/base/utils/metadata/version.py

@@ -0,0 +1,64 @@
+# (C) Datadog, Inc. 2019
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+import re
+
+from ..common import exclude_undefined_keys
+
+# https://semver.org/#is-there-a-suggested-regular-expression-regex-to-check-a-semver-string
+SEMVER_PATTERN = re.compile(
+    r"""
+    (?P<major>0|[1-9]\d*)
+    \.
+    (?P<minor>0|[1-9]\d*)
+    \.
+    (?P<patch>0|[1-9]\d*)
+    (?:-(?P<release>
+        (?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)
+        (?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*
+    ))?
+    (?:\+(?P<build>
+        [0-9a-zA-Z-]+
+        (?:\.[0-9a-zA-Z-]+)*
+    ))?
+    """,
+    re.VERBOSE,
+)
+
+
+def parse_semver(version, options):
+    match = SEMVER_PATTERN.search(version)
+    if not match:
+        raise ValueError('Version does not adhere to semantic versioning')
+
+    return exclude_undefined_keys(match.groupdict())
+
+
+def parse_regex(version, options):
+    pattern = options.get('pattern')
+    if not pattern:
+        raise ValueError('Version scheme `regex` requires a `pattern` option')
+
+    match = re.search(pattern, version)
+    if not match:
+        raise ValueError('Version does not match the regular expression pattern')
+
+    parts = match.groupdict()
+    if not parts:
+        raise ValueError('Regular expression pattern has no named subgroups')
+
+    return exclude_undefined_keys(parts)
+
+
+def parse_version(version, options):
+    scheme = options.get('scheme')
+
+    if not scheme:
+        scheme = 'semver'
+    elif scheme not in SCHEMES:
+        raise ValueError('Unsupported version scheme `{}`'.format(scheme))
+
+    return scheme, SCHEMES[scheme](version, options)
+
+
+SCHEMES = {'semver': parse_semver, 'regex': parse_regex}