allow autoconf templates in docker labels

tests/core/test_service_discovery.py

@@ -14,7 +14,7 @@
 from utils.service_discovery.consul_config_store import ConsulStore
 from utils.service_discovery.etcd_config_store import EtcdStore
 from utils.service_discovery.abstract_config_store import AbstractConfigStore, \
-    _TemplateCache, CONFIG_FROM_KUBE, CONFIG_FROM_TEMPLATE, CONFIG_FROM_AUTOCONF
+    _TemplateCache, CONFIG_FROM_KUBE, CONFIG_FROM_TEMPLATE, CONFIG_FROM_AUTOCONF, CONFIG_FROM_LABELS
 from utils.service_discovery.sd_backend import get_sd_backend
 from utils.service_discovery.sd_docker_backend import SDDockerBackend, _SDDockerBackendConfigFetchState
 from utils.dockerutil import DockerUtil
@@ -603,6 +603,31 @@ def test_get_check_tpls_kube(self, *args):
                          'service-discovery.datadoghq.com/foo.instances'],
                         self.mock_raw_templates[image][0]))))
 
+    @mock.patch('config.get_auto_confd_path', return_value=os.path.join(
+        os.path.dirname(__file__), 'fixtures/auto_conf/'))
+    @mock.patch.object(AbstractConfigStore, 'client_read', side_effect=client_read)
+    def test_get_check_tpls_labels(self, *args):
+        """Test get_check_tpls for kubernetes annotations"""
+        valid_config = ['image_0', 'image_1', 'image_2', 'image_3', 'image_4']
+        invalid_config = ['bad_image_0']
+        config_store = get_config_store(self.auto_conf_agentConfig)
+        for image in valid_config + invalid_config:
+            tpl = self.mock_raw_templates.get(image)[1]
+            tpl = [(CONFIG_FROM_LABELS, t[1]) for t in tpl]
+            if tpl:
+                self.assertNotEquals(
+                    tpl,
+                    config_store.get_check_tpls(image, auto_conf=True))
+            self.assertEquals(
+                tpl,
+                config_store.get_check_tpls(
+                    image, auto_conf=True,
+                    docker_labels=dict(zip(
+                        ['service-discovery.datadoghq.com/check_names',
+                         'service-discovery.datadoghq.com/init_configs',
+                         'service-discovery.datadoghq.com/instances'],
+                        self.mock_raw_templates[image][0]))))
+
     @mock.patch('config.get_auto_confd_path', return_value=os.path.join(
         os.path.dirname(__file__), 'fixtures/auto_conf/'))
     def test_get_config_id(self, mock_get_auto_confd_path):

utils/service_discovery/abstract_config_store.py

@@ -23,12 +23,15 @@
 CONFIG_FROM_FILE = 'YAML file'
 CONFIG_FROM_TEMPLATE = 'template'
 CONFIG_FROM_KUBE = 'Kubernetes Pod Annotation'
+CONFIG_FROM_LABELS = 'Docker container labels'
+
 TRACE_CONFIG = 'trace_config'  # used for tracing config load by service discovery
 CHECK_NAMES = 'check_names'
 INIT_CONFIGS = 'init_configs'
 INSTANCES = 'instances'
 KUBE_ANNOTATIONS = 'kube_annotations'
 KUBE_CONTAINER_NAME = 'kube_container_name'
+DOCKER_LABELS = 'docker_labels'
 KUBE_ANNOTATION_PREFIX = 'service-discovery.datadoghq.com'
 
 
@@ -210,17 +213,28 @@ def dump_directory(self, path, **kwargs):
         raise NotImplementedError()
 
     def _get_kube_config(self, identifier, kube_annotations, kube_container_name):
+        prefix = '{}/{}.'.format(KUBE_ANNOTATION_PREFIX, kube_container_name)
+        return self._extract_template(identifier, prefix, kube_annotations)
+
+    def _get_docker_config(self, identifier, docker_labels):
+        prefix = '{}/'.format(KUBE_ANNOTATION_PREFIX)
+        return self._extract_template(identifier, prefix, docker_labels)
+
+    def _extract_template(self, identifier, key_prefix, source_dict):
+        """
+        Looks for autodiscovery configuration in a given source_dict (either docker labels
+        or kubernetes annotations) and returns it if found.
+        """
         try:
-            prefix = '{}/{}.'.format(KUBE_ANNOTATION_PREFIX, kube_container_name)
-            check_names = json.loads(kube_annotations[prefix + CHECK_NAMES])
-            init_config_tpls = json.loads(kube_annotations[prefix + INIT_CONFIGS])
-            instance_tpls = json.loads(kube_annotations[prefix + INSTANCES])
+            check_names = json.loads(source_dict[key_prefix + CHECK_NAMES])
+            init_config_tpls = json.loads(source_dict[key_prefix + INIT_CONFIGS])
+            instance_tpls = json.loads(source_dict[key_prefix + INSTANCES])
             return [check_names, init_config_tpls, instance_tpls]
         except KeyError:
             return None
         except json.JSONDecodeError:
             log.exception('Could not decode the JSON configuration template '
-                          'for the kubernetes pod with ident %s...' % identifier)
+                          'for the entity %s...' % identifier)
             return None
 
     def _get_auto_config(self, image_name):
@@ -254,13 +268,19 @@ def get_checks_to_refresh(self, identifier, **kwargs):
 
         kube_annotations = kwargs.get(KUBE_ANNOTATIONS)
         kube_container_name = kwargs.get(KUBE_CONTAINER_NAME)
+        docker_labels = kwargs.get(DOCKER_LABELS)
 
         # then from annotations
         if kube_annotations:
             kube_config = self._get_kube_config(identifier, kube_annotations, kube_container_name)
             if kube_config is not None:
                 to_check.update(kube_config[0])
 
+        # then from docker labels
+        if docker_labels:
+            kube_config = self._get_docker_config(identifier, docker_labels)
+            if kube_config is not None:
+                to_check.update(kube_config[0])
         # lastly, try with legacy name for auto-conf
         to_check.update(self.template_cache.get_check_names(self._get_image_ident(identifier)))
 
@@ -276,13 +296,21 @@ def get_check_tpls(self, identifier, **kwargs):
             # annotations for configs before falling back to autoconf.
             kube_annotations = kwargs.get(KUBE_ANNOTATIONS)
             kube_container_name = kwargs.get(KUBE_CONTAINER_NAME)
+            docker_labels = kwargs.get(DOCKER_LABELS)
+            source = ""
+
+            kube_config = None
             if kube_annotations:
                 kube_config = self._get_kube_config(identifier, kube_annotations, kube_container_name)
-                if kube_config is not None:
-                    check_names, init_config_tpls, instance_tpls = kube_config
-                    source = CONFIG_FROM_KUBE
-                    return [(source, vs)
-                            for vs in zip(check_names, init_config_tpls, instance_tpls)]
+                source = CONFIG_FROM_KUBE
+            if kube_config is None and docker_labels is not None:
+                kube_config = self._get_docker_config(identifier, docker_labels)
+                source = CONFIG_FROM_LABELS
+
+            if kube_config is not None:
+                check_names, init_config_tpls, instance_tpls = kube_config
+                return [(source, vs)
+                        for vs in zip(check_names, init_config_tpls, instance_tpls)]
 
             # in auto config mode, identifier is the image name
             auto_config = self._get_auto_config(identifier)

utils/service_discovery/sd_docker_backend.py

@@ -374,7 +374,7 @@ def get_configs(self):
             try:
                 # value of the DATADOG_ID tag or the image name if the label is missing
                 identifier = self.get_config_id(image, labels)
-                check_configs = self._get_check_configs(state, cid, identifier) or []
+                check_configs = self._get_check_configs(state, cid, identifier, labels) or []
                 for conf in check_configs:
                     source, (check_name, init_config, instance) = conf
 
@@ -403,7 +403,7 @@ def get_config_id(self, image, labels):
         """Look for a DATADOG_ID label, return its value or the image name if missing"""
         return labels.get(DATADOG_ID) or image
 
-    def _get_check_configs(self, state, c_id, identifier):
+    def _get_check_configs(self, state, c_id, identifier, labels=None):
         """Retrieve configuration templates and fill them with data pulled from docker and tags."""
         platform_kwargs = {}
         if Platform.is_k8s():
@@ -412,6 +412,9 @@ def _get_check_configs(self, state, c_id, identifier):
                 'kube_container_name': state.get_kube_container_name(c_id),
                 'kube_annotations': kube_metadata.get('annotations'),
             }
+        if labels:
+            platform_kwargs['docker_labels'] = labels
+
         config_templates = self._get_config_templates(identifier, **platform_kwargs)
         if not config_templates:
             return None