NIFI-6675 - This closes #3741. Added support for JsonPath put operation

apache · Sep 20, 2019 · 79e4f36 · 79e4f36
1 parent 3411251
commit 79e4f36
Show file tree

Hide file tree

Showing 7 changed files with 184 additions and 5 deletions.
diff --git a/...ain/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionLexer.g b/...ain/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionLexer.g
@@ -190,6 +190,7 @@ REPLACE_ALL : 'replaceAll';
 IF_ELSE : 'ifElse';
 JSON_PATH_SET : 'jsonPathSet';
 JSON_PATH_ADD : 'jsonPathAdd';
+JSON_PATH_PUT : 'jsonPathPut';
 PAD_LEFT : 'padLeft';
 PAD_RIGHT : 'padRight';
 

diff --git a/...in/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionParser.g b/...in/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionParser.g
@@ -80,6 +80,7 @@ oneArgString : ((SUBSTRING_BEFORE | SUBSTRING_BEFORE_LAST | SUBSTRING_AFTER | SU
 			   (TO_RADIX LPAREN! anyArg (COMMA! anyArg)? RPAREN!);
 twoArgString : ((REPLACE | REPLACE_FIRST | REPLACE_ALL | IF_ELSE | JSON_PATH_SET | JSON_PATH_ADD) LPAREN! anyArg COMMA! anyArg RPAREN!) |
 			   ((SUBSTRING | FORMAT | PAD_LEFT | PAD_RIGHT) LPAREN! anyArg (COMMA! anyArg)? RPAREN!);
+threeArgString: ((JSON_PATH_PUT) LPAREN! anyArg COMMA! anyArg COMMA! anyArg RPAREN!);
 fiveArgString : GET_DELIMITED_FIELD LPAREN! anyArg (COMMA! anyArg (COMMA! anyArg (COMMA! anyArg (COMMA! anyArg)?)?)?)? RPAREN!;
 
 // functions that return Booleans
@@ -98,7 +99,7 @@ oneArgNum	: ((INDEX_OF | LAST_INDEX_OF) LPAREN! anyArg RPAREN!) |
 oneOrTwoArgNum : MATH LPAREN! anyArg (COMMA! anyArg)? RPAREN!;
 zeroOrOneOrTwoArgNum : TO_DATE LPAREN! anyArg? (COMMA! anyArg)? RPAREN!;
 
-stringFunctionRef : zeroArgString | oneArgString | twoArgString | fiveArgString;
+stringFunctionRef : zeroArgString | oneArgString | twoArgString | threeArgString | fiveArgString;
 booleanFunctionRef : zeroArgBool | oneArgBool | multiArgBool;
 numberFunctionRef : zeroArgNum | oneArgNum | oneOrTwoArgNum | zeroOrOneOrTwoArgNum;
 

diff --git a/...c/main/java/org/apache/nifi/attribute/expression/language/compile/ExpressionCompiler.java b/...c/main/java/org/apache/nifi/attribute/expression/language/compile/ExpressionCompiler.java
@@ -68,6 +68,7 @@
 import org.apache.nifi.attribute.expression.language.evaluation.functions.JsonPathDeleteEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.JsonPathEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.JsonPathSetEvaluator;
+import org.apache.nifi.attribute.expression.language.evaluation.functions.JsonPathPutEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.LastIndexOfEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.LengthEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.LessThanEvaluator;
@@ -182,6 +183,7 @@
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JOIN;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JSON_PATH;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JSON_PATH_ADD;
+import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JSON_PATH_PUT;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JSON_PATH_DELETE;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.JSON_PATH_SET;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.LAST_INDEX_OF;
@@ -970,6 +972,20 @@ private Evaluator<?> buildFunctionEvaluator(final Tree tree, final Evaluator<?>
                         toStringEvaluator(argEvaluators.get(0), "first argument to jsonPathAdd"),
                         valueEvaluator), "jsonPathAdd");
             }
+            case JSON_PATH_PUT: {
+                verifyArgCount(argEvaluators, 3, "jsonPathPut");
+                Evaluator<?> argValueEvaluator = argEvaluators.get(1);
+                String valueLocation = "second argument to jsonPathPut";
+                Evaluator<?> valueEvaluator = getJsonPathUpdateEvaluator(argValueEvaluator, valueLocation);
+                Evaluator<?> argKeyEvaluator = argEvaluators.get(2);
+                String keyLocation = "third argument to jsonPathPut";
+                Evaluator<?> keyEvaluator = getJsonPathUpdateEvaluator(argKeyEvaluator, keyLocation);
+
+                return addToken(new JsonPathPutEvaluator(toStringEvaluator(subjectEvaluator),
+                        toStringEvaluator(argEvaluators.get(0), "first argument to jsonPathPut"),
+                        toStringEvaluator(keyEvaluator, keyLocation),
+                        valueEvaluator), "jsonPathPut");
+            }
             case IF_ELSE: {
                 verifyArgCount(argEvaluators, 2, "ifElse");
                 return addToken(new IfElseEvaluator(toBooleanEvaluator(subjectEvaluator),

diff --git a/.../apache/nifi/attribute/expression/language/evaluation/functions/JsonPathPutEvaluator.java b/.../apache/nifi/attribute/expression/language/evaluation/functions/JsonPathPutEvaluator.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.attribute.expression.language.evaluation.functions;
+
+import com.jayway.jsonpath.DocumentContext;
+import com.jayway.jsonpath.JsonPath;
+import org.apache.nifi.attribute.expression.language.EvaluationContext;
+import org.apache.nifi.attribute.expression.language.evaluation.Evaluator;
+import org.apache.nifi.attribute.expression.language.evaluation.QueryResult;
+import org.apache.nifi.attribute.expression.language.evaluation.StringQueryResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * JsonPathPutEvaluator allows setting or adding a key and scalar value at the specified existing path
+ */
+public class JsonPathPutEvaluator extends JsonPathUpdateEvaluator {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(JsonPathPutEvaluator.class);
+
+    protected Evaluator<?> keyEvaluator;
+
+    public JsonPathPutEvaluator(final Evaluator<String> subject, final Evaluator<String> jsonPathExp, final Evaluator<?> valueEvaluator, final Evaluator<?> keyEvaluator) {
+        super(subject, jsonPathExp, valueEvaluator);
+        this.keyEvaluator = keyEvaluator;
+    }
+
+    @Override
+    public QueryResult<String> evaluate(EvaluationContext context) {
+        DocumentContext documentContext = getDocumentContext(context);
+        final JsonPath compiledJsonPath = getJsonPath(context);
+
+        final Object value = valueEvaluator.evaluate(context).getValue();
+        final String key = keyEvaluator.evaluate(context).getValue().toString();
+
+        String result;
+        try {
+            result = documentContext.put(compiledJsonPath, key, value).jsonString();
+        } catch (Exception e) {
+            LOGGER.error("Failed to put value " + value + " at key " + key + " at path " + compiledJsonPath + " with error " + e.getLocalizedMessage(), e);
+            // assume the path did not match anything in the document
+            return EMPTY_RESULT;
+        }
+
+        return new StringQueryResult(getResultRepresentation(result, EMPTY_RESULT.getValue()));
+    }
+
+}
+
diff --git a/...ache/nifi/attribute/expression/language/evaluation/functions/JsonPathUpdateEvaluator.java b/...ache/nifi/attribute/expression/language/evaluation/functions/JsonPathUpdateEvaluator.java
@@ -18,6 +18,7 @@
 
 import com.jayway.jsonpath.DocumentContext;
 import com.jayway.jsonpath.JsonPath;
+import org.apache.commons.lang3.NotImplementedException;
 import org.apache.nifi.attribute.expression.language.EvaluationContext;
 import org.apache.nifi.attribute.expression.language.evaluation.Evaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.QueryResult;
@@ -27,6 +28,8 @@
 
 /**
  * JsonPathUpdateEvaluator is base class for updating attributes
+ *
+ * Subclasses need to implement {@link #updateAttribute} method otherwise it throws {@link NotImplementedException}
  */
 public abstract class JsonPathUpdateEvaluator extends JsonPathBaseEvaluator {
 
@@ -59,12 +62,15 @@ public QueryResult<String> evaluate(EvaluationContext context) {
     }
 
     /**
-     * Update the attribute at the specified path
+     * Update the attribute at the specified path.  The subclasses will need to implement this method.
      * @param documentContext the document to be updated
      * @param jsonPath the path to update
      * @param value the value to be applied at the specified path
      * @return the updated DocumentContext
+     * @throws NotImplementedException if operation is not implemented
      */
-    public abstract DocumentContext updateAttribute(DocumentContext documentContext, JsonPath jsonPath, Object value);
+    public DocumentContext updateAttribute(DocumentContext documentContext, JsonPath jsonPath, Object value) {
+        throw new NotImplementedException("Please implement updateAttribute method in the implementation class");
+    }
 }
 
diff --git a/...ssion-language/src/test/java/org/apache/nifi/attribute/expression/language/TestQuery.java b/...ssion-language/src/test/java/org/apache/nifi/attribute/expression/language/TestQuery.java
@@ -375,9 +375,9 @@ private Map<String,String> verifyJsonPathExpressions(String targetAttribute, Obj
             verifyEquals(targetAttribute, attributes, originalValue);
         }
 
-        String addressBookAfterDelete = Query.evaluateExpressions(updateExpression, attributes, ParameterLookup.EMPTY);
+        String addressBookAfterUpdate = Query.evaluateExpressions(updateExpression, attributes, ParameterLookup.EMPTY);
         attributes.clear();
-        attributes.put("json", addressBookAfterDelete);
+        attributes.put("json", addressBookAfterUpdate);
 
         verifyAddressBookAttributes(addressBook, attributes, targetAttribute, updatedValue);
 
@@ -486,6 +486,45 @@ public void testJsonPathAddNicknameJimmyAtNonArray() throws IOException {
                 "");
     }
 
+    @Test
+    public void testJsonPathPutRootLevelMiddlenameTuron() throws IOException {
+        Map<String,String> attributes = verifyJsonPathExpressions(
+                ADDRESS_BOOK_JSON_PATH_EMPTY,
+                "",
+                "${json:jsonPathPut('$','middlename','Turon')}",
+                "");
+        verifyEquals("${json:jsonPath('$.middlename')}", attributes, "Turon");
+    }
+
+    @Test
+    public void testJsonPathPutCountryToMap() throws IOException {
+        Map<String,String> attributes = verifyJsonPathExpressions(
+                ADDRESS_BOOK_JSON_PATH_EMPTY,
+                "",
+                "${json:jsonPathPut('$.address','country','US')}",
+                "");
+        verifyEquals("${json:jsonPath('$.address.country')}", attributes, "US");
+    }
+
+    @Test
+    public void testJsonPathPutElementToArray() throws IOException {
+        Map<String,String> attributes = verifyJsonPathExpressions(
+                ADDRESS_BOOK_JSON_PATH_EMPTY,
+                "",
+                "${json:jsonPathPut('$.phoneNumbers[1]', 'backup', '212-555-1212')}",
+                "");
+        verifyEquals("${json:jsonPath('$.phoneNumbers[1].backup')}", attributes, "212-555-1212");
+    }
+
+    @Test
+    public void testJsonPathPutOverwriteFirstNameToJimmy() throws IOException {
+        Map<String,String> attributes = verifyJsonPathExpressions(
+                ADDRESS_BOOK_JSON_PATH_FIRST_NAME,
+                "John",
+                "${json:jsonPathPut('$','firstName','Jimmy')}",
+                "Jimmy");
+    }
+
     @Test
     public void testEmbeddedExpressionsAndQuotesWithProperties() {
         final Map<String, String> attributes = new HashMap<>();

diff --git a/nifi-docs/src/main/asciidoc/expression-language-guide.adoc b/nifi-docs/src/main/asciidoc/expression-language-guide.adoc
@@ -1806,6 +1806,59 @@ form of the updated JSON.#
 
 An empty subject value or a subject value with an invalid JSON document results in an exception bulletin.
 
+[.function]
+=== jsonPathPut
+
+*Description*: [.description]#The `jsonPathPut` function puts the key and scalar value at the specified JsonPath on a Subject JSON and returns string
+form of the updated JSON.#
+
+*Subject Type*: [.subject]#String#
+
+*Arguments*:
+
+- [.argName]#_jsonPath_# : [.argDesc]#the JSON path expression to set value on the Subject.#
+- [.argName]#_value_# : [.argDesc]#the value expression to be set on the specified path on Subject.#
+- [.argName]#_key_# : [.argDesc]#the key expression with the associated value the specified path on Subject.#
+
+*Return Type*: [.returnType]#String#
+
+*Examples*: If the "myJson" attribute is
+
+..........
+{
+     "firstName": "John",
+     "lastName": "Smith",
+     "age": 25,
+     "voter" : true,
+     "height" : 6.1,
+     "address" : {
+         "streetAddress": "21 2nd Street",
+         "city": "New York",
+         "state": "NY",
+         "postalCode": "10021-3100"
+     },
+     "phoneNumbers": [
+         {
+             "type": "home",
+             "number": "212 555-1234"
+         },
+         {
+             "type": "office",
+             "number": "646 555-4567"
+         }
+     ],
+     "nicknames" : []
+ }
+..........
+
+.jsonPathPut Examples
+|===================================================================
+| Expression | Value
+| `${myJson:jsonPathPut('$','middlename','Turon')}` | `{"firstName":"James", lastName":"Smith", "middlename": "Turon", "age":25, "voter":true, "height":6.1, "address":{"streetAddress":"21 2nd Street", "city":"New York", "state":"NY", "postalCode":"10021-3100"}, "phoneNumbers":[{"type":"home", "number":"212 555-1234"}, {"type":"office", "number":"646 555-4567"}]}`
+|===================================================================
+
+An empty subject value or a subject value with an invalid JSON document results in an exception bulletin.
+
 [[numbers]]
 == Mathematical Operations and Numeric Manipulation