NIFI-1826 Expression Language: add function to check enumerator

Signed-off-by: Matt Burgess <mattyb149@apache.org>

nifi-commons/nifi-expression-language/src/main/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionLexer.g

@@ -161,6 +161,9 @@ REPLACE_ALL : 'replaceAll';
 // 4 arg functions
 GET_DELIMITED_FIELD	: 'getDelimitedField';
 
+// unlimited arg functions
+IN : 'in';
+
 // STRINGS
 STRING_LITERAL
 @init{StringBuilder lBuf = new StringBuilder();}

nifi-commons/nifi-expression-language/src/main/antlr3/org/apache/nifi/attribute/expression/language/antlr/AttributeExpressionParser.g

@@ -87,6 +87,7 @@ oneArgBool	: ((FIND | MATCHES | EQUALS_IGNORE_CASE) LPAREN! anyArg RPAREN!) |
 			  (GREATER_THAN | LESS_THAN | GREATER_THAN_OR_EQUAL | LESS_THAN_OR_EQUAL) LPAREN! anyArg RPAREN! |
 			  (EQUALS) LPAREN! anyArg RPAREN! |
 			  (AND | OR) LPAREN! anyArg RPAREN!;
+multiArgBool : (IN) LPAREN! anyArg (COMMA! anyArg)* RPAREN!;
 
 
 // functions that return Numbers
@@ -96,10 +97,10 @@ oneArgNum	: ((INDEX_OF | LAST_INDEX_OF) LPAREN! anyArg RPAREN!) |
 			  ((MOD | PLUS | MINUS | MULTIPLY | DIVIDE) LPAREN! anyArg RPAREN!);
 
 stringFunctionRef : zeroArgString | oneArgString | twoArgString | fiveArgString;
-booleanFunctionRef : zeroArgBool | oneArgBool;
+booleanFunctionRef : zeroArgBool | oneArgBool | multiArgBool;
 numberFunctionRef : zeroArgNum | oneArgNum;
 
-anyArg : NUMBER | numberFunctionRef | STRING_LITERAL | zeroArgString | oneArgString | twoArgString | fiveArgString | booleanLiteral | zeroArgBool | oneArgBool | expression;
+anyArg : NUMBER | numberFunctionRef | STRING_LITERAL | zeroArgString | oneArgString | twoArgString | fiveArgString | booleanLiteral | zeroArgBool | oneArgBool | multiArgBool | expression;
 stringArg : STRING_LITERAL | zeroArgString | oneArgString | twoArgString | expression;
 functionRef : stringFunctionRef | booleanFunctionRef | numberFunctionRef;
 

nifi-commons/nifi-expression-language/src/main/java/org/apache/nifi/attribute/expression/language/Query.java

@@ -55,6 +55,7 @@ import org.apache.nifi.attribute.expression.language.evaluation.functions.Greate
 import org.apache.nifi.attribute.expression.language.evaluation.functions.GreaterThanOrEqualEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.HostnameEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.IPEvaluator;
+import org.apache.nifi.attribute.expression.language.evaluation.functions.InEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.IndexOfEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.IsEmptyEvaluator;
 import org.apache.nifi.attribute.expression.language.evaluation.functions.IsNullEvaluator;
@@ -131,6 +132,7 @@ import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpre
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.ATTRIBUTE_REFERENCE;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.ATTR_NAME;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.CONTAINS;
+import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.IN;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.COUNT;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.DIVIDE;
 import static org.apache.nifi.attribute.expression.language.antlr.AttributeExpressionParser.ENDS_WITH;
@@ -1188,6 +1190,13 @@ public class Query {
                 return addToken(new ContainsEvaluator(toStringEvaluator(subjectEvaluator),
                     toStringEvaluator(argEvaluators.get(0), "first argument to contains")), "contains");
             }
+            case IN: {
+                List<Evaluator<String>> list = new ArrayList<Evaluator<String>>();
+                for(int i = 0; i < argEvaluators.size(); i++) {
+                    list.add(toStringEvaluator(argEvaluators.get(i), i + "th argument to in"));
+                }
+                return addToken(new InEvaluator(toStringEvaluator(subjectEvaluator), list), "in");
+            }
             case FIND: {
                 verifyArgCount(argEvaluators, 1, "find");
                 return addToken(new FindEvaluator(toStringEvaluator(subjectEvaluator),

nifi-commons/nifi-expression-language/src/main/java/org/apache/nifi/attribute/expression/language/evaluation/functions/InEvaluator.java

@@ -0,0 +1,61 @@
+/*
+ * 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 java.util.List;
+import java.util.Map;
+
+import org.apache.nifi.attribute.expression.language.evaluation.BooleanEvaluator;
+import org.apache.nifi.attribute.expression.language.evaluation.BooleanQueryResult;
+import org.apache.nifi.attribute.expression.language.evaluation.Evaluator;
+import org.apache.nifi.attribute.expression.language.evaluation.QueryResult;
+
+public class InEvaluator extends BooleanEvaluator {
+
+    private final Evaluator<String> subject;
+    private final List<Evaluator<String>> search;
+
+    public InEvaluator(final Evaluator<String> subject, final List<Evaluator<String>> list) {
+        this.subject = subject;
+        this.search = list;
+    }
+
+    @Override
+    public QueryResult<Boolean> evaluate(final Map<String, String> attributes) {
+        final String subjectValue = subject.evaluate(attributes).getValue();
+        if (subjectValue == null) {
+            return new BooleanQueryResult(false);
+        }
+
+        boolean isInList = false;
+        for (Evaluator<String> evaluator : search) {
+            final String searchString = evaluator.evaluate(attributes).getValue();
+            isInList = searchString == null ? false : subjectValue.equals(searchString);
+            if(isInList) {
+                break;
+            }
+        }
+
+        return new BooleanQueryResult(isInList);
+    }
+
+    @Override
+    public Evaluator<?> getSubjectEvaluator() {
+        return subject;
+    }
+
+}

nifi-commons/nifi-expression-language/src/test/java/org/apache/nifi/attribute/expression/language/TestQuery.java

@@ -1082,6 +1082,17 @@ public class TestQuery {
         verifyEquals("${ abc:toNumber():equals(123) }", attributes, true);
     }
 
+    @Test
+    public void testIn() {
+        final Map<String, String> attributes = new HashMap<>();
+        attributes.put("myEnum", "JOHN");
+        verifyEquals("${ myEnum:in('PAUL', 'JOHN', 'MIKE') }", attributes, true);
+        verifyEquals("${ myEnum:in('RED', 'BLUE', 'GREEN') }", attributes, false);
+
+        attributes.put("toReplace", "BLUE");
+        verifyEquals("${ myEnum:in('RED', ${ toReplace:replace('BLUE', 'JOHN') }, 'GREEN') }", attributes, true);
+    }
+
     @Test
     public void testSubjectAsEmbeddedExpressionWithSurroundChars() {
         final Map<String, String> attributes = new HashMap<>();