[instance] support instance tags in list format

Support instance tags defined in an array/list format.
DataDog · Mar 3, 2017 · 94cbee9 · 94cbee9
1 parent 2f58003
commit 94cbee9
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 7 deletions.
diff --git a/src/main/java/org/datadog/jmxfetch/Instance.java b/src/main/java/org/datadog/jmxfetch/Instance.java
@@ -1,6 +1,7 @@
 package org.datadog.jmxfetch;
 
 import java.io.IOException;
+import java.lang.ClassCastException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -69,7 +70,7 @@ public Instance(LinkedHashMap<String, Object> yamlInstance, LinkedHashMap<String
         this.yaml = yamlInstance != null ? new LinkedHashMap<String, Object>(yamlInstance) : null;
         this.initConfig = initConfig != null ? new LinkedHashMap<String, Object>(initConfig) : null;
         this.instanceName = (String) yaml.get("name");
-        this.tags = (LinkedHashMap<String, String>) yaml.get("tags");
+        this.tags = getTagsMap(yaml.get("tags"));
         this.checkName = checkName;
         this.matchingAttributes = new LinkedList<JMXAttribute>();
         this.failingAttributes = new HashSet<JMXAttribute>();
@@ -126,6 +127,27 @@ public Instance(LinkedHashMap<String, Object> yamlInstance, LinkedHashMap<String
         configurationList.add(new Configuration((LinkedHashMap<String, Object>) new YamlParser(this.getClass().getResourceAsStream("/jmx-2.yaml")).getParsedYaml()));
     }
 
+    /**
+     * Format the instance tags defined in the YAML configuration file to a `LinkedHashMap`.
+     * Supported inputs: `List`, `Map`.
+     */
+    private static LinkedHashMap<String, String> getTagsMap(Object yamlTags){
+        try {
+            // Input has `Map` format
+            return (LinkedHashMap<String, String>) yamlTags;
+        }
+        catch (ClassCastException e){
+            // Input has `List` format
+            LinkedHashMap<String, String> tags = new LinkedHashMap<String, String>();
+
+            for (String tag: (List<String>)yamlTags) {
+                tags.put(tag, null);
+            }
+
+            return tags;
+        }
+    }
+
     public Connection getConnection(LinkedHashMap<String, Object> connectionParams, boolean forceNewConnection) throws IOException {
         if (connection == null || !connection.isAlive()) {
             LOGGER.info("Connection closed or does not exist. Creating a new connection!");
@@ -335,9 +357,13 @@ public String[] getServiceCheckTags() {
             tags.add("jmx_server:" + this.yaml.get("host"));
         }
         if (this.tags != null) {
-        	for (Entry<String, String> e : this.tags.entrySet()) {
-        		tags.add(e.getKey() + ":" + e.getValue());
-        	}
+            for (Entry<String, String> e : this.tags.entrySet()) {
+                if (e.getValue()!=null){
+                    tags.add(e.getKey() + ":" + e.getValue());
+                } else {
+                    tags.add(e.getKey());
+                }
+            }
         }
         tags.add("instance:" + this.instanceName);
         return tags.toArray(new String[tags.size()]);

diff --git a/src/main/java/org/datadog/jmxfetch/JMXAttribute.java b/src/main/java/org/datadog/jmxfetch/JMXAttribute.java
@@ -134,7 +134,12 @@ private LinkedList<String> getBeanParametersList(String instanceName, Map<String
 
         if (instanceTags != null) {
             for (Map.Entry<String, String> tag : instanceTags.entrySet()) {
-                beanTags.add(tag.getKey() + ":" + tag.getValue());
+                if (tag.getValue() != null) {
+                    beanTags.add(tag.getKey() + ":" + tag.getValue());
+                }
+                else {
+                    beanTags.add(tag.getKey());
+                }
             }
         }
 

diff --git a/src/test/resources/jmx.yaml b/src/test/resources/jmx.yaml
@@ -5,8 +5,8 @@ instances:
         refresh_beans: 4
         name: jmx_test_instance
         tags:
-            env: stage
-            newTag: test
+            - "env:stage"
+            - "newTag:test"
         conf:
             - include:
                domain: org.datadog.jmxfetch.test