(Backport 5.2) - Pipeline function json handling improvements (#17895)

* Update select_jsonpath to accept strings of JSON in addition to JsonNode objects (#17683) * Update select_jsonpath to accept strings of JSON in addition to JsonNode objects * Add unit test * Add changelog entry * Update changelog with correct issue/pr Co-authored-by: Zack King <[email protected]> --------- Co-authored-by: Zack King <[email protected]> * Add handling for json arrays in lookup_all pipeline function (#17820) * Add handling for json arrays in lookup_all pipeline function * Add changelog entry * Update failing tests * Revert unneeded non string functionality * Update changelog * Cleanup test --------- Co-authored-by: Zack King <[email protected]>
Graylog2 · Jan 11, 2024 · 782a5d6 · 782a5d6
1 parent 31583a0
commit 782a5d6
Show file tree

Hide file tree

Showing 7 changed files with 148 additions and 21 deletions.
diff --git a/changelog/unreleased/issue-17647.toml b/changelog/unreleased/issue-17647.toml
@@ -0,0 +1,5 @@
+type = "c"
+message = "Updated the select_jsonpath pipeline function to accept JSON strings as the `json` parameter in addition to parsed JsonNode objects."
+
+issues = ["17647"]
+pulls = ["17683"]
diff --git a/changelog/unreleased/pr-17820.toml b/changelog/unreleased/pr-17820.toml
@@ -0,0 +1,5 @@
+type = "fixed"
+message = "Fix issue preventing the lookup_all pipeline function from working with json arrays"
+
+issues = ["graylog-plugin-enterprise#6363"]
+pulls = ["17820"]
diff --git a/...er/src/main/java/org/graylog/plugins/pipelineprocessor/functions/json/SelectJsonPath.java b/...er/src/main/java/org/graylog/plugins/pipelineprocessor/functions/json/SelectJsonPath.java
@@ -16,6 +16,7 @@
  */
 package org.graylog.plugins.pipelineprocessor.functions.json;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
@@ -44,18 +45,20 @@ public class SelectJsonPath extends AbstractFunction<Map<String, Object>> {
 
     public static final String NAME = "select_jsonpath";
 
+    private final ObjectMapper objectMapper;
     private final Configuration configuration;
-    private final ParameterDescriptor<JsonNode, JsonNode> jsonParam;
+    private final ParameterDescriptor<Object, Object> jsonParam;
     private final ParameterDescriptor<Map<String, String>, Map<String, JsonPath>> pathsParam;
 
     @Inject
     public SelectJsonPath(ObjectMapper objectMapper) {
+        this.objectMapper = objectMapper;
         configuration = Configuration.builder()
                 .options(Option.SUPPRESS_EXCEPTIONS)
                 .jsonProvider(new JacksonJsonNodeJsonProvider(objectMapper))
                 .build();
 
-        jsonParam = ParameterDescriptor.type("json", JsonNode.class).description("A parsed JSON tree").build();
+        jsonParam = ParameterDescriptor.type("json", Object.class).description("A parsed JSON tree or String representation of a JSON tree").build();
         // sigh generics and type erasure
         //noinspection unchecked
         pathsParam = ParameterDescriptor.type("paths",
@@ -70,7 +73,21 @@ public SelectJsonPath(ObjectMapper objectMapper) {
 
     @Override
     public Map<String, Object> evaluate(FunctionArgs args, EvaluationContext context) {
-        final JsonNode json = jsonParam.required(args, context);
+        final Object jsonObj = jsonParam.required(args, context);
+        JsonNode json = null;
+        if (jsonObj instanceof JsonNode jsonNode) {
+            json = jsonNode;
+        } else if (jsonObj instanceof String jsonString) {
+            try {
+                json = objectMapper.readTree(jsonString);
+            } catch (JsonProcessingException e) {
+                log.warn(context.pipelineErrorMessage("Unable to parse JSON"), e);
+            }
+        } else {
+            throw new IllegalArgumentException(context.pipelineErrorMessage(
+                    "`json` parameter must be a parsed JSON tree or String representation of a JSON tree"));
+        }
+
         final Map<String, JsonPath> paths = pathsParam.required(args, context);
         if (json == null || paths == null) {
             return Collections.emptyMap();

diff --git a/...erver/src/main/java/org/graylog/plugins/pipelineprocessor/functions/lookup/LookupAll.java b/...erver/src/main/java/org/graylog/plugins/pipelineprocessor/functions/lookup/LookupAll.java
@@ -16,8 +16,7 @@
  */
 package org.graylog.plugins.pipelineprocessor.functions.lookup;
 
-import com.google.common.base.Functions;
-import com.google.common.collect.Lists;
+import com.fasterxml.jackson.databind.node.ValueNode;
 import com.google.common.reflect.TypeToken;
 import org.graylog.plugins.pipelineprocessor.EvaluationContext;
 import org.graylog.plugins.pipelineprocessor.ast.functions.AbstractFunction;
@@ -30,8 +29,11 @@
 
 import javax.inject.Inject;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 import static org.graylog.plugins.pipelineprocessor.ast.functions.ParameterDescriptor.object;
 import static org.graylog.plugins.pipelineprocessor.ast.functions.ParameterDescriptor.string;
@@ -55,7 +57,7 @@ public LookupAll(LookupTableService lookupTableService) {
                 .build();
         keysParam = object("keys", LIST_RETURN_TYPE)
                 .description("The keys to lookup in the table")
-                .transform(LookupAll::transformValueToList)
+                .transform(this::transformToList)
                 .build();
     }
 
@@ -93,11 +95,20 @@ public FunctionDescriptor<List<Object>> descriptor() {
                 .build();
     }
 
-    private static List<Object> transformValueToList(Object value) {
-        if (value instanceof List) {
-            return Lists.transform((List) value, Functions.toStringFunction());
-        } else {
-            return Collections.singletonList(value.toString());
+    private List<Object> transformToList(Object value) {
+        if (value instanceof Collection<?>) {
+            return ((Collection<?>) value).stream()
+                    .map(LookupAll::convertValue)
+                    .filter(Objects::nonNull)
+                    .collect(Collectors.toList());
         }
+        return Collections.singletonList(value.toString());
+    }
+
+    private static String convertValue(Object o) {
+        if (o instanceof ValueNode node) {
+            return node.textValue();
+        }
+        return o.toString();
     }
 }
diff --git a/.../src/test/java/org/graylog/plugins/pipelineprocessor/functions/FunctionsSnippetsTest.java b/.../src/test/java/org/graylog/plugins/pipelineprocessor/functions/FunctionsSnippetsTest.java
@@ -474,6 +474,55 @@ public void jsonpath() {
         assertThat(message.hasField("this_should_exist")).isTrue();
     }
 
+    @Test
+    public void jsonpathFromMessageField() {
+        final String json = "{\n" +
+                "    \"store\": {\n" +
+                "        \"book\": [\n" +
+                "            {\n" +
+                "                \"category\": \"reference\",\n" +
+                "                \"author\": \"Nigel Rees\",\n" +
+                "                \"title\": \"Sayings of the Century\",\n" +
+                "                \"price\": 8.95\n" +
+                "            },\n" +
+                "            {\n" +
+                "                \"category\": \"fiction\",\n" +
+                "                \"author\": \"Evelyn Waugh\",\n" +
+                "                \"title\": \"Sword of Honour\",\n" +
+                "                \"price\": 12.99\n" +
+                "            },\n" +
+                "            {\n" +
+                "                \"category\": \"fiction\",\n" +
+                "                \"author\": \"Herman Melville\",\n" +
+                "                \"title\": \"Moby Dick\",\n" +
+                "                \"isbn\": \"0-553-21311-3\",\n" +
+                "                \"price\": 8.99\n" +
+                "            },\n" +
+                "            {\n" +
+                "                \"category\": \"fiction\",\n" +
+                "                \"author\": \"J. R. R. Tolkien\",\n" +
+                "                \"title\": \"The Lord of the Rings\",\n" +
+                "                \"isbn\": \"0-395-19395-8\",\n" +
+                "                \"price\": 22.99\n" +
+                "            }\n" +
+                "        ],\n" +
+                "        \"bicycle\": {\n" +
+                "            \"color\": \"red\",\n" +
+                "            \"price\": 19.95\n" +
+                "        }\n" +
+                "    },\n" +
+                "    \"expensive\": 10\n" +
+                "}";
+
+        final Rule rule = parser.parseRule(ruleForTest(), false);
+        final Message message = evaluateRule(rule, new Message(json, "test", Tools.nowUTC()));
+
+        assertThat(message.hasField("author_first")).isTrue();
+        assertThat(message.getField("author_first")).isEqualTo("Nigel Rees");
+        assertThat(message.hasField("author_last")).isTrue();
+        assertThat(message.hasField("this_should_exist")).isTrue();
+    }
+
     @Test
     public void json() {
         final String flatJson = "{\"str\":\"foobar\",\"int\":42,\"float\":2.5,\"bool\":true,\"array\":[1,2,3]}";
@@ -1349,20 +1398,30 @@ public void lookupAssignTtl() {
     }
 
     @Test
-    public void lookupAll() {
-        doReturn(LookupResult.single("val1")).when(lookupTable).lookup("key1");
-        doReturn(LookupResult.single("val2")).when(lookupTable).lookup("key2");
-        doReturn(LookupResult.single("val3")).when(lookupTable).lookup("key3");
+    public void lookupAll() throws IOException {
+        doReturn(LookupResult.single("val1")).when(lookupTable).lookup("one");
+        doReturn(LookupResult.single("val2")).when(lookupTable).lookup("two");
+        doReturn(LookupResult.single("val3")).when(lookupTable).lookup("three");
 
         final Rule rule = parser.parseRule(ruleForTest(), false);
-        final Message message = evaluateRule(rule);
+        final Message message = new Message("message", "source", DateTime.now(DateTimeZone.UTC));
+
+        try (InputStream inputStream = getClass().getResourceAsStream("with-arrays.json")) {
+            String jsonString = IOUtils.toString(Objects.requireNonNull(inputStream), StandardCharsets.UTF_8);
+            message.addField("json_with_arrays", jsonString);
+            evaluateRule(rule, message);
+            assertThat(actionsTriggered.get()).isTrue();
+        }
+
+        verify(lookupTable, times(3)).lookup("one");
+        verify(lookupTable, times(2)).lookup("two");
+        verify(lookupTable, times(2)).lookup("three");
 
-        verify(lookupTable).lookup("key1");
-        verify(lookupTable).lookup("key2");
-        verify(lookupTable).lookup("key3");
         verifyNoMoreInteractions(lookupTable);
 
+        assertThat(message.getField("json_results")).isEqualTo(Arrays.asList("val1", "val2", "val3"));
         assertThat(message.getField("results")).isEqualTo(Arrays.asList("val1", "val2", "val3"));
+        assertThat(message.getField("single_result")).isEqualTo(Arrays.asList("val1"));
     }
 
     @Test

diff --git a/...st/resources/org/graylog/plugins/pipelineprocessor/functions/jsonpathFromMessageField.txt b/...st/resources/org/graylog/plugins/pipelineprocessor/functions/jsonpathFromMessageField.txt
@@ -0,0 +1,23 @@
+rule "jsonpathFromMessageField"
+when
+  is_json(parse_json("{}")) == true &&
+  is_json("foobar") == false &&
+  is_json(1234) == false &&
+  is_json(12.34) == false &&
+  is_json(true) == false
+then
+  let new_fields = select_jsonpath($message.message,
+            { author_first: "$['store']['book'][0]['author']",
+              author_last: "$['store']['book'][-1:]['author']"
+            });
+  set_fields(new_fields);
+
+  // Don't fail on empty input
+  let invalid_json = parse_json("#FOOBAR#");
+  let invalid_json_fields = select_jsonpath(invalid_json, { some_field: "$.message" });
+  set_fields(invalid_json_fields);
+
+  // Don't fail on missing field
+  let missing_fields = select_jsonpath($message.message, { some_field: "$.i_dont_exist", this_should_exist: "$['store']['book'][-1:]['author']" });
+  set_fields(missing_fields);
+end
diff --git a/...2-server/src/test/resources/org/graylog/plugins/pipelineprocessor/functions/lookupAll.txt b/...2-server/src/test/resources/org/graylog/plugins/pipelineprocessor/functions/lookupAll.txt
@@ -2,6 +2,13 @@ rule "lookup_all"
 when
     true
 then
-    set_field("results", lookup_all("table", ["key1", "key2", "key3"]));
+    let json = parse_json(to_string($message.json_with_arrays));
+            let selectedPath = select_jsonpath(json: json,
+                paths: {strings: "$.strings[*].value"});
+
+    set_field("json_results", lookup_all("table", selectedPath.strings));
+    set_field("results", lookup_all("table", ["one", "two", "three"]));
+    set_field("single_result", lookup_all("table", "one"));
+
     trigger_test();
-end
+end