Replace buggy image name extraction with get_identifier.

hkaj · gmmeyer · commit 5ed54ff5eb0d · 2016-06-28T12:27:40.000-04:00
get_identifier looks for a datadog-id label in the container
and falls back to the image name if it can't find it.
diff --git a/tests/core/test_service_discovery.py b/tests/core/test_service_discovery.py
@@ -60,30 +60,42 @@ def client_read(path):
     return TestServiceDiscovery.mock_tpls.get(image)[0][config_parts.index(config_part)]
 
 
+def issue_read(identifier):
+    return TestServiceDiscovery.mock_tpls.get(identifier)
+
+
 class TestServiceDiscovery(unittest.TestCase):
     docker_container_inspect = {
         u'Id': u'69ff25598b2314d1cdb7752cc3a659fb1c1352b32546af4f1454321550e842c0',
-        u'Image': u'6ffc02088cb870652eca9ccd4c4fb582f75b29af2879792ed09bb46fd1c898ef',
+        u'Image': u'nginx',
         u'Name': u'/nginx',
         u'NetworkSettings': {u'IPAddress': u'172.17.0.21', u'Ports': {u'443/tcp': None, u'80/tcp': None}}
     }
+    docker_container_inspect_with_label = {
+        u'Id': u'69ff25598b2314d1cdb7752cc3a659fb1c1352b32546af4f1454321550e842c0',
+        u'Image': u'nginx',
+        u'Name': u'/nginx',
+        u'NetworkSettings': {u'IPAddress': u'172.17.0.21', u'Ports': {u'443/tcp': None, u'80/tcp': None}},
+        u'Labels': {'com.datadoghq.sd.check.id': 'custom-nginx'}
+    }
     kubernetes_container_inspect = {
         u'Id': u'389dc8a4361f3d6c866e9e9a7b6972b26a31c589c4e2f097375d55656a070bc9',
-        u'Image': u'de309495e6c7b2071bc60c0b7e4405b0d65e33e3a4b732ad77615d90452dd827',
+        u'Image': u'foo',
         u'Name': u'/k8s_sentinel.38057ab9_redis-master_default_27b84e1e-a81c-11e5-8347-42010af00002_f70875a1',
         u'Config': {u'ExposedPorts': {u'6379/tcp': {}}},
         u'NetworkSettings': {u'IPAddress': u'', u'Ports': None}
     }
     malformed_container_inspect = {
         u'Id': u'69ff25598b2314d1cdb7752cc3a659fb1c1352b32546af4f1454321550e842c0',
-        u'Image': u'6ffc02088cb870652eca9ccd4c4fb582f75b29af2879792ed09bb46fd1c898ef',
+        u'Image': u'foo',
         u'Name': u'/nginx'
     }
     container_inspects = [
-        # (inspect_dict, expected_ip, expected_port)
-        (docker_container_inspect, '172.17.0.21', 'port', '443'),
-        (kubernetes_container_inspect, None, 'port', '6379'),  # arbitrarily defined in the mocked pod_list
-        (malformed_container_inspect, None, 'port', KeyError)
+        # (inspect_dict, expected_ip, expected_port, expected_ident)
+        (docker_container_inspect, '172.17.0.21', ['80', '443'], 'nginx'),
+        (docker_container_inspect_with_label, '172.17.0.21', ['80', '443'], 'custom-nginx'),
+        (kubernetes_container_inspect, None, ['6379'], 'foo'),  # arbitrarily defined in the mocked pod_list
+        (malformed_container_inspect, None, KeyError, 'foo')
     ]
 
     # templates with variables already extracted
@@ -114,7 +126,12 @@ class TestServiceDiscovery(unittest.TestCase):
             [('check_2', {}, {"host": "%%host%%", "port": "%%port%%"})]),
         'bad_image_0': ((['invalid template']), []),
         'bad_image_1': (('invalid template'), []),
-        'bad_image_2': (None, [])
+        'bad_image_2': (None, []),
+        'nginx': ('["nginx"]', '[{}]', '[{"host": "localhost"}]'),
+        'nginx:latest': ('["nginx"]', '[{}]', '[{"host": "localhost", "tags": ["foo"]}]'),
+        'custom-nginx': ('["nginx"]', '[{}]', '[{"host": "localhost"}]'),
+        'repo/custom-nginx': ('["nginx"]', '[{}]', '[{"host": "localhost", "tags": ["bar"]}]'),
+        'repo/dir:5000/custom-nginx:latest': ('["nginx"]', '[{}]', '[{"host": "local", "tags": ["foobar"]}]')
     }
 
     bad_mock_templates = {
@@ -207,22 +224,22 @@ def test_get_host_address(self, mock_check_yaml, mock_get):
         mock_check_yaml.return_value = kubernetes_config
         mock_get.return_value = Response(pod_list)
 
-        for c_ins, tpl_var, expected_ip in ip_address_inspects:
+        for c_ins, expected_ip, _, _ in self.container_inspects:
             with mock.patch.object(AbstractConfigStore, '__init__', return_value=None):
                 with mock.patch('utils.dockerutil.DockerUtil.client', return_value=None):
                     with mock.patch('utils.kubeutil.get_conf_path', return_value=None):
                         sd_backend = get_sd_backend(agentConfig=self.auto_conf_agentConfig)
-                        self.assertEquals(sd_backend._get_host_address(c_ins, tpl_var), expected_ip)
+                        self.assertEqual(sd_backend._get_host(c_ins), expected_ip)
                         clear_singletons(self.auto_conf_agentConfig)
 
     def test_get_port(self):
         with mock.patch('utils.dockerutil.DockerUtil.client', return_value=None):
-            for c_ins, _, var_tpl, expected_ports in self.container_inspects:
+            for c_ins, _, expected_ports, _ in self.container_inspects:
                 sd_backend = get_sd_backend(agentConfig=self.auto_conf_agentConfig)
-                if isinstance(expected_ports, str):
-                    self.assertEquals(sd_backend._get_port(c_ins, var_tpl), expected_ports)
+                if isinstance(expected_ports, list):
+                    self.assertEqual(sd_backend._get_ports(c_ins), expected_ports)
                 else:
-                    self.assertRaises(expected_ports, sd_backend._get_port, c_ins, var_tpl)
+                    self.assertRaises(expected_ports, sd_backend._get_ports, c_ins)
                 clear_singletons(self.auto_conf_agentConfig)
 
     @mock.patch('docker.Client.inspect_container', side_effect=_get_container_inspect)
@@ -498,3 +515,29 @@ def test_get_check_tpls(self, mock_client_read):
         for image in invalid_config:
             tpl = self.mock_tpls.get(image)[1]
             self.assertEquals(tpl, config_store.get_check_tpls(image))
+
+    def test_get_config_id(self):
+        """Test get_config_id"""
+        with mock.patch('utils.dockerutil.DockerUtil.client', return_value=None):
+            for c_ins, _, _, expected_ident in self.container_inspects:
+                sd_backend = get_sd_backend(agentConfig=self.auto_conf_agentConfig)
+                self.assertEqual(
+                    sd_backend.get_config_id(c_ins.get('Image'), c_ins.get('Labels', {})),
+                    expected_ident)
+                clear_singletons(self.auto_conf_agentConfig)
+
+    @mock.patch.object(AbstractConfigStore, '_issue_read', side_effect=issue_read)
+    def test_read_config_from_store(self, issue_read):
+        """Test read_config_from_store"""
+        valid_idents = [('nginx', 'nginx'), ('nginx:latest', 'nginx:latest'),
+                        ('custom-nginx', 'custom-nginx'), ('custom-nginx:latest', 'custom-nginx'),
+                        ('repo/custom-nginx:latest', 'custom-nginx'),
+                        ('repo/dir:5000/custom-nginx:latest', 'repo/dir:5000/custom-nginx:latest')]
+        invalid_idents = ['foo']
+        config_store = get_config_store(self.auto_conf_agentConfig)
+        for ident, expected_key in valid_idents:
+            tpl = config_store.read_config_from_store(ident)
+            # source is added after reading from the store
+            self.assertEquals(tpl, ('template',) + self.mock_tpls.get(expected_key))
+        for ident in invalid_idents:
+            self.assertEquals(config_store.read_config_from_store(ident), [])
diff --git a/utils/service_discovery/abstract_config_store.py b/utils/service_discovery/abstract_config_store.py
@@ -20,6 +20,9 @@
 CONFIG_FROM_FILE = 'YAML file'
 CONFIG_FROM_TEMPLATE = 'template'
 TRACE_CONFIG = 'trace_config'  # used for tracing config load by service discovery
+CHECK_NAMES = 'check_names'
+INIT_CONFIGS = 'init_configs'
+INSTANCES = 'instances'
 
 
 class KeyNotFound(Exception):
@@ -59,7 +62,8 @@ def dump_directory(self, path, **kwargs):
         raise NotImplementedError()
 
     def _get_auto_config(self, image_name):
-        if image_name in self.auto_conf_images:
+        # use the image name, ignore the tag
+        if image_name.split(':')[0] in self.auto_conf_images:
             check_name = self.auto_conf_images[image_name]
 
             # get the check class to verify it matches
@@ -75,35 +79,35 @@ def _get_auto_config(self, image_name):
 
         return None
 
-    def get_check_tpls(self, image, **kwargs):
-        """Retrieve template configs for an image from the config_store or auto configuration."""
-        # TODO: make mixing both sources possible
+    def get_check_tpls(self, identifier, **kwargs):
+        """Retrieve template configs for an identifier from the config_store or auto configuration."""
         templates = []
         trace_config = kwargs.get(TRACE_CONFIG, False)
 
         # this flag is used when no valid configuration store was provided
         # it makes the method skip directly to the auto_conf
         if kwargs.get('auto_conf') is True:
-            auto_config = self._get_auto_config(image)
+            # in auto config mode, identifier is the image name
+            auto_config = self._get_auto_config(identifier)
             if auto_config is not None:
                 source = CONFIG_FROM_AUTOCONF
                 if trace_config:
                     return [(source, auto_config)]
                 return [auto_config]
             else:
-                log.debug('No auto config was found for image %s, leaving it alone.' % image)
+                log.debug('No auto config was found for image %s, leaving it alone.' % identifier)
                 return []
         else:
-            config = self.read_config_from_store(image)
+            config = self.read_config_from_store(identifier)
             if config:
                 source, check_names, init_config_tpls, instance_tpls = config
             else:
                 return []
 
         if len(check_names) != len(init_config_tpls) or len(check_names) != len(instance_tpls):
             log.error('Malformed configuration template: check_names, init_configs '
-                      'and instances are not all the same length. Image {0} '
-                      'will not be configured by the service discovery'.format(image))
+                      'and instances are not all the same length. Container with identifier {} '
+                      'will not be configured by the service discovery'.format(identifier))
             return []
 
         for idx, c_name in enumerate(check_names):
@@ -113,45 +117,72 @@ def get_check_tpls(self, image, **kwargs):
                 templates.append((c_name, init_config_tpls[idx], instance_tpls[idx]))
         return templates
 
-    def read_config_from_store(self, image):
+    def read_config_from_store(self, identifier):
         """Try to read from the config store, falls back to auto-config in case of failure."""
         try:
-            check_names = json.loads(
-                self.client_read(path.join(self.sd_template_dir, image, 'check_names').lstrip('/')))
-            init_config_tpls = json.loads(
-                self.client_read(path.join(self.sd_template_dir, image, 'init_configs').lstrip('/')))
-            instance_tpls = json.loads(
-                self.client_read(path.join(self.sd_template_dir, image, 'instances').lstrip('/')))
-            source = CONFIG_FROM_TEMPLATE
-        except (KeyNotFound, TimeoutError, json.JSONDecodeError) as ex:
-            # first case is kind of expected, it means that no template was provided for this container
-            if isinstance(ex, KeyNotFound):
-                log.debug("Could not find directory {0} in the config store, "
-                          "trying to auto-configure a check...".format(image))
+            try:
+                res = self._issue_read(identifier)
+            except KeyNotFound:
+                log.debug("Could not find directory {} in the config store, "
+                          "trying to convert to the old format.".format(identifier))
+                legacy_ident = identifier.split(':')[0].split('/')[-1]
+                res = self._issue_read(legacy_ident)
+
+            if res and len(res) == 3:
+                source = CONFIG_FROM_TEMPLATE
+                check_names, init_config_tpls, instance_tpls = res
+            else:
+                log.debug("Could not find directory {} in the config store, "
+                          "trying to convert to the old format...".format(identifier))
+                legacy_ident = identifier.split(':')[0].split('/')[-1]
+                res = self._issue_read(legacy_ident)
+                if res and len(res) == 3:
+                    source = CONFIG_FROM_TEMPLATE
+                    check_names, init_config_tpls, instance_tpls = res
+                else:
+                    raise KeyError
+        except (KeyError, KeyNotFound, TimeoutError, json.JSONDecodeError) as ex:
+            # this is kind of expected, it means that no template was provided for this container
+            if isinstance(ex, KeyError) or isinstance(ex, KeyNotFound):
+                log.debug("Could not find directory {} in the config store, "
+                          "trying to auto-configure a check...".format(identifier))
             # this case is not expected, the agent can't reach the config store
-            elif isinstance(ex, TimeoutError):
+            if isinstance(ex, TimeoutError):
                 log.warning("Connection to the config backend timed out. Is it reachable?\n"
-                            "Trying to auto-configure a check for the image %s..." % image)
+                            "Trying to auto-configure a check for the container with ident %s." % identifier)
             # the template is reachable but invalid
             elif isinstance(ex, json.JSONDecodeError):
-                log.error('Could not decode the JSON configuration template. '
-                          'Trying to auto-configure a check for the image %s...' % image)
-            # In any case cases we try to read from auto-config templates
-            auto_config = self._get_auto_config(image)
+                log.error('Could not decode the JSON configuration template '
+                          'for the container with ident %s...' % identifier)
+                return []
+            # try to read from auto-config templates
+            auto_config = self._get_auto_config(identifier)
             if auto_config is not None:
                 # create list-format config based on an autoconf template
                 check_names, init_config_tpls, instance_tpls = map(lambda x: [x], auto_config)
                 source = CONFIG_FROM_AUTOCONF
             else:
-                log.debug('No auto config was found for image %s, leaving it alone.' % image)
+                log.debug('No config was found for container with ident %s, leaving it alone.' % identifier)
                 return []
         except Exception as ex:
             log.warning(
                 'Fetching the value for {0} in the config store failed, this check '
-                'will not be configured by the service discovery. Error: {1}'.format(image, str(ex)))
+                'will not be configured by the service discovery. Error: {1}'.format(identifier, str(ex)))
             return []
         return source, check_names, init_config_tpls, instance_tpls
 
+    def _issue_read(self, identifier):
+        try:
+            check_names = json.loads(
+                self.client_read(path.join(self.sd_template_dir, identifier, CHECK_NAMES).lstrip('/')))
+            init_config_tpls = json.loads(
+                self.client_read(path.join(self.sd_template_dir, identifier, INIT_CONFIGS).lstrip('/')))
+            instance_tpls = json.loads(
+                self.client_read(path.join(self.sd_template_dir, identifier, INSTANCES).lstrip('/')))
+            return [check_names, init_config_tpls, instance_tpls]
+        except KeyError:
+            return None
+
     def crawl_config_template(self):
         """Return whether or not configuration templates have changed since the previous crawl"""
         try:
diff --git a/utils/service_discovery/sd_docker_backend.py b/utils/service_discovery/sd_docker_backend.py
