From 81a72d2afac838034289b9181ec988be223ad885 Mon Sep 17 00:00:00 2001
From: Yann Mahe <yann@datadoghq.com>
Date: Thu, 14 May 2015 11:28:08 -0400
Subject: [PATCH] [fluentd] limit `tag_by` param to a specific scope

Limit `tag_by` parameter to `plugin_id` or `type` only, to prevent
empty or bad tagging. Minor cosmetic changes too.
---
 checks.d/fluentd.py         | 23 ++++++++++++++---------
 conf.d/fluentd.yaml.example | 15 ++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/checks.d/fluentd.py b/checks.d/fluentd.py
index 89227e3759..ac5ea7bd3f 100644
--- a/checks.d/fluentd.py
+++ b/checks.d/fluentd.py
@@ -1,18 +1,18 @@
 # stdlib
-from collections import defaultdict
 import urlparse
 
+# 3rd party
+import requests
+
 # project
 from util import headers
 from checks import AgentCheck
 
-# 3rd party
-import simplejson as json
-import requests
 
 class Fluentd(AgentCheck):
     SERVICE_CHECK_NAME = 'fluentd.is_ok'
     GAUGES = ['retry_count', 'buffer_total_queued_size', 'buffer_queue_length']
+    _AVAILABLE_TAGS = frozenset(['plugin_id', 'type'])
 
     """Tracks basic fluentd metrics via the monitor_agent plugin
     * number of retry_count
@@ -29,12 +29,16 @@ def check(self, instance):
         try:
             url = instance.get('monitor_agent_url')
             plugin_ids = instance.get('plugin_ids', [])
-            tag_by = instance.get('tag_by', 'plugin_id')
+
+            # Fallback  with `tag_by: plugin_id`
+            tag_by = instance.get('tag_by')
+            tag_by = tag_by if (tag_by in self._AVAILABLE_TAGS) else 'plugin_id'
 
             parsed_url = urlparse.urlparse(url)
             monitor_agent_host = parsed_url.hostname
             monitor_agent_port = parsed_url.port or 24220
-            service_check_tags = ['fluentd_host:%s' % monitor_agent_host, 'fluentd_port:%s' % monitor_agent_port]
+            service_check_tags = ['fluentd_host:%s' % monitor_agent_host, 'fluentd_port:%s'
+                                  % monitor_agent_port]
 
             r = requests.get(url, headers=headers(self.agentConfig))
             r.raise_for_status()
@@ -46,11 +50,12 @@ def check(self, instance):
                     if p.get(m) is None:
                         continue
                     # Filter unspecified plugins to keep backward compatibility.
-                    if tag_by == 'type' or len(plugin_ids) == 0 or p.get('plugin_id') in plugin_ids:
+                    if len(plugin_ids) == 0 or p.get('plugin_id') in plugin_ids:
                         self.gauge('fluentd.%s' % (m), p.get(m), [tag])
         except Exception, e:
             msg = "No stats could be retrieved from %s : %s" % (url, str(e))
-            self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.CRITICAL, tags=service_check_tags, message=msg)
-            raise e
+            self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.CRITICAL,
+                               tags=service_check_tags, message=msg)
+            raise
         else:
             self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.OK, tags=service_check_tags)
diff --git a/conf.d/fluentd.yaml.example b/conf.d/fluentd.yaml.example
index eac9092a2a..b49cfb6747 100644
--- a/conf.d/fluentd.yaml.example
+++ b/conf.d/fluentd.yaml.example
@@ -1,16 +1,13 @@
 init_config:
 
 instances:
-    # For every instance, you have an `monitor_agent_url` 
-    # Use fluend 'plugin_id' as a tag. 
-    # You need to specify plugin id in fluentd configuration file.
-    -  monitor_agent_url: http://example.com:24220/api/plugins.json
-       tag_by: plugin_id
-    # Use fluend plugin 'type' as a tag. 
-    -  monitor_agent_url: http://example.org:24220/api/plugins.json
-       tag_by: type
-    # For backward compatibilty, you can specify plugin_ids that are used as tags.
+    # Every instance requires a `monitor_agent_url`
+    # Optional, set `plugin_ids` to monitor a specific scope of plugins.
     -  monitor_agent_url: http://example.com:24220/api/plugins.json
        plugin_ids:
          - plg1
          - plg2
+    # Optional, set 'tag_by' to specify how to tag metrics. By default, metrics are tagged with `plugin_id`
+    -  monitor_agent_url: http://example.com:24220/api/plugins.json
+       tag_by: type
+