Support replacing whitespace in nested keys for JSON extractor (#3623)

Fixes #3622
Graylog2 · Mar 17, 2017 · 4955d39 · 4955d39
1 parent 8ec556f
commit 4955d39
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 15 deletions.
diff --git a/graylog2-server/src/main/java/org/graylog2/inputs/extractors/JsonExtractor.java b/graylog2-server/src/main/java/org/graylog2/inputs/extractors/JsonExtractor.java
@@ -38,6 +38,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
@@ -51,6 +52,7 @@ public class JsonExtractor extends Extractor {
     private static final String CK_REPLACE_KEY_WHITESPACE = "replace_key_whitespace";
     private static final String CK_KEY_WHITESPACE_REPLACEMENT = "key_whitespace_replacement";
     private static final String CK_KEY_PREFIX = "key_prefix";
+    private static final Pattern WHITE_SPACE_PATTERN = Pattern.compile("\\s");
     private static final RemoveNullPredicate REMOVE_NULL_PREDICATE = new RemoveNullPredicate();
 
     private final ObjectMapper mapper = new ObjectMapper();
@@ -115,48 +117,53 @@ public Map<String, Object> extractJson(String value) {
 
         final Map<String, Object> results = new HashMap<>(json.size());
         for (Map.Entry<String, Object> mapEntry : json.entrySet()) {
-            String key = keyPrefix + mapEntry.getKey();
-            if (replaceKeyWhitespace && key.contains(" ")) {
-                key = key.replace(" ", keyWhitespaceReplacement);
-            } else {
-                if (LOG.isDebugEnabled() && key.contains(" ")) {
-                    LOG.debug("Invalid key \"{}\" in JSON object!", key);
-                }
-            }
-            for (Entry entry : parseValue(key, mapEntry.getValue())) {
+            for (Entry entry : parseValue(keyPrefix + mapEntry.getKey(), mapEntry.getValue())) {
                 results.put(entry.key(), entry.value());
             }
         }
 
         return results;
     }
 
+    private String parseKey(String key) {
+        if (replaceKeyWhitespace && key.contains(" ")) {
+            return WHITE_SPACE_PATTERN.matcher(key).replaceAll(keyWhitespaceReplacement);
+        } else {
+            if (LOG.isDebugEnabled() && key.contains(" ")) {
+                LOG.debug("Invalid key \"{}\" in JSON object!", key);
+            }
+
+            return key;
+        }
+    }
+
     private Collection<Entry> parseValue(String key, Object value) {
+        final String processedKey = parseKey(key);
         if (value instanceof Boolean) {
-            return Collections.singleton(Entry.create(key, value));
+            return Collections.singleton(Entry.create(processedKey, value));
         } else if (value instanceof Number) {
-            return Collections.singleton(Entry.create(key, value));
+            return Collections.singleton(Entry.create(processedKey, value));
         } else if (value instanceof String) {
-            return Collections.singleton(Entry.create(key, value));
+            return Collections.singleton(Entry.create(processedKey, value));
         } else if (value instanceof Map) {
             @SuppressWarnings("unchecked")
             final Map<String, Object> map = (Map<String, Object>) value;
             final Map<String, Object> withoutNull = Maps.filterEntries(map, REMOVE_NULL_PREDICATE);
             if (flatten) {
                 final Joiner.MapJoiner joiner = Joiner.on(listSeparator).withKeyValueSeparator(kvSeparator);
-                return Collections.singleton(Entry.create(key, joiner.join(withoutNull)));
+                return Collections.singleton(Entry.create(processedKey, joiner.join(withoutNull)));
             } else {
                 final List<Entry> result = new ArrayList<>(map.size());
                 for (Map.Entry<String, Object> entry : map.entrySet()) {
-                    result.addAll(parseValue(key + keySeparator + entry.getKey(), entry.getValue()));
+                    result.addAll(parseValue(processedKey + keySeparator + entry.getKey(), entry.getValue()));
                 }
 
                 return result;
             }
         } else if (value instanceof List) {
             final List values = (List) value;
             final Joiner joiner = Joiner.on(listSeparator).skipNulls();
-            return Collections.singleton(Entry.create(key, joiner.join(values)));
+            return Collections.singleton(Entry.create(processedKey, joiner.join(values)));
         } else if (value == null) {
             // Ignore null values so we don't try to create fields for that in the message.
             return Collections.emptySet();

diff --git a/graylog2-server/src/test/java/org/graylog2/inputs/extractors/JsonExtractorTest.java b/graylog2-server/src/test/java/org/graylog2/inputs/extractors/JsonExtractorTest.java
@@ -195,6 +195,28 @@ public void testRunWithWhitespaceInKey() throws Exception {
         );
     }
 
+    @Test
+    public void testRunWithWhitespaceInNestedKey() throws Exception {
+        final String value = "{\"foo\":{\"b a r\":{\"b a z\": 42}}}";
+        final JsonExtractor jsonExtractor = new JsonExtractor(
+                new MetricRegistry(),
+                "json",
+                "title",
+                0L,
+                Extractor.CursorStrategy.COPY,
+                "source",
+                "target",
+                ImmutableMap.of("replace_key_whitespace", true, "key_whitespace_replacement", "-"),
+                "user",
+                Collections.emptyList(),
+                Extractor.ConditionType.NONE,
+                "");
+
+        assertThat(jsonExtractor.run(value)).containsOnly(
+                new Extractor.Result(42, "foo_b-a-r_b-a-z", -1, -1)
+        );
+    }
+
     @Test
     public void testRunWithKeyPrefix() throws Exception {
         final String value = "{\"text string\": \"foobar\", \"num   b er\": 1234.5678, \"bool\": true, \"null\": null}";