Support functions in Data Prepper expressions #2626

data-prepper-expression/build.gradle

@@ -28,6 +28,7 @@ dependencies {
     implementation 'org.apache.logging.log4j:log4j-core'
     implementation 'org.apache.logging.log4j:log4j-slf4j2-impl'
     testImplementation testLibs.spring.test
+    testImplementation "org.apache.commons:commons-lang3:3.12.0"
     testImplementation 'com.fasterxml.jackson.core:jackson-databind'
 }
 

data-prepper-expression/src/main/antlr/DataPrepperExpression.g4

@@ -86,14 +86,14 @@ setInitializer
     : LBRACE primary (SET_DELIMITER primary)* RBRACE
     ;
 
-
 unaryOperator
     : NOT
     | SUBTRACT
     ;
 
 primary
     : jsonPointer
+    | function
     | variableIdentifier
     | setInitializer
     | literal
@@ -104,6 +104,25 @@ jsonPointer
     | EscapedJsonPointer
     ;
 
+function
+    : Function
+    ;
+
+Function
+    : JsonPointerCharacters LPAREN FunctionArgs RPAREN
+    ;
+
+fragment
+FunctionArgs
+    : (FunctionArg SPACE* COMMA)* SPACE* FunctionArg
+    ;
+
+fragment
+FunctionArg
+    : JsonPointer
+    | String
+    ;
+
 variableIdentifier
     : variableName
     ;
@@ -227,7 +246,10 @@ EscapeSequence
     : '\\' [btnfr"'\\$]
     ;
 
-SET_DELIMITER : ',';
+SET_DELIMITER
+    : COMMA
+    ;
+COMMA : ',';
 EQUAL : '==';
 NOT_EQUAL : '!=';
 LT : '<';

data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/ExpressionFunction.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.expression;
+
+import java.util.List;
+
+interface ExpressionFunction {
+    /**
+     * @return function name
+     * @since 2.3
+     */
+    String getFunctionName();
+
+    /**
+     * evaluates the function and returns the result
+     * @param args list of arguments to the function
+     * @return the result of function evaluation
+     * @since 2.3
+     */
+    Object evaluate(final List<Object> args);
+}

data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/ExpressionFunctionProvider.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.expression;
+
+import javax.inject.Named;
+import javax.inject.Inject;
+import java.util.Map;
+import java.util.List;
+import java.util.stream.Collectors;
+
+@Named
+public class ExpressionFunctionProvider {
+    private Map<String, ExpressionFunction> expressionFunctionsMap;
+
+    @Inject
+    public ExpressionFunctionProvider(final List<ExpressionFunction> expressionFunctions) {
+        expressionFunctionsMap = expressionFunctions.stream().collect(Collectors.toMap(e -> e.getFunctionName(), e -> e));
+    }
+
+    public Object provideFunction(final String functionName, final List<Object> argList) {
+        if (!expressionFunctionsMap.containsKey(functionName)) {
+            return null;
+        }
+        return expressionFunctionsMap.get(functionName).evaluate(argList);
+    }
+
+}

data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/LengthExpressionFunction.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.expression;
+
+import java.util.List;
+import javax.inject.Named;
+
+@Named
+public class LengthExpressionFunction implements ExpressionFunction {
+    public String getFunctionName() {
+        return "length";
+    }
+
+    public Object evaluate(final List<Object> args) {
+        if (args.size() != 1) {
+            throw new RuntimeException("length() takes only one argument");
+        }
+        Object arg = args.get(0);
+        if (!(arg instanceof String)) {
+            throw new RuntimeException("length() takes only String type arguments");
+        }
+        String argStr = ((String)arg).trim();
+        if (argStr.charAt(0) == '\"') {
+            if (argStr.charAt(argStr.length()-1) != '\"') {
+                throw new RuntimeException("Invalid string passed to length() "+argStr);
+            }
+            return Integer.valueOf(argStr.length()-2);
+        } else {
+            return Integer.valueOf(argStr.length());
+        }
+    }
+}
+

data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/ParseTreeCoercionService.java

@@ -13,21 +13,43 @@
 import javax.inject.Named;
 import java.io.Serializable;
 import java.util.Map;
+import java.util.List;
+import java.util.ArrayList;
 import java.util.function.Function;
 
 @Named
 class ParseTreeCoercionService {
     private final Map<Class<? extends Serializable>, Function<Object, Object>> literalTypeConversions;
+    private ExpressionFunctionProvider expressionFunctionProvider;
 
     @Inject
-    public ParseTreeCoercionService(final Map<Class<? extends Serializable>, Function<Object, Object>> literalTypeConversions) {
+    public ParseTreeCoercionService(final Map<Class<? extends Serializable>, Function<Object, Object>> literalTypeConversions, ExpressionFunctionProvider expressionFunctionProvider) {
         this.literalTypeConversions = literalTypeConversions;
+        this.expressionFunctionProvider = expressionFunctionProvider;
     }
 
     public Object coercePrimaryTerminalNode(final TerminalNode node, final Event event) {
         final int nodeType = node.getSymbol().getType();
         final String nodeStringValue = node.getText();
         switch (nodeType) {
+            case DataPrepperExpressionParser.Function:
+                final int funcNameIndex = nodeStringValue.indexOf("(");
+                final String functionName = nodeStringValue.substring(0, funcNameIndex);
+                final int argsEndIndex = nodeStringValue.indexOf(")", funcNameIndex);
+                final String argsStr = nodeStringValue.substring(funcNameIndex+1, argsEndIndex);
+                final String[] args = argsStr.split(",");
+                List<Object> argList = new ArrayList<>();
+                for (final String arg: args) {
+                    String trimmedArg = arg.trim();
+                    if (trimmedArg.charAt(0) == '/') {
+                        argList.add(resolveJsonPointerValue(arg, event));
+                    } else if (trimmedArg.charAt(0) == '"') {
+                        argList.add(trimmedArg);
+                    } else {
+                        throw new RuntimeException("Unsupported type passed as function argument");
+                    }
+                }
+                return expressionFunctionProvider.provideFunction(functionName, argList);
             case DataPrepperExpressionParser.EscapedJsonPointer:
                 final String jsonPointerWithoutQuotes = nodeStringValue.substring(1, nodeStringValue.length() - 1);
                 return resolveJsonPointerValue(jsonPointerWithoutQuotes, event);

data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ExpressionFunctionProviderTest.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.expression;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.apache.commons.lang3.RandomStringUtils;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.lenient;
+
+import java.util.List;
+
+@ExtendWith(MockitoExtension.class)
+class ExpressionFunctionProviderTest {
+    private ExpressionFunctionProvider objectUnderTest;
+    private String testFunctionName;
+    private Object testResultObject;
+    private ExpressionFunction expressionFunction;
+
+    public ExpressionFunctionProvider createObjectUnderTest() {
+        expressionFunction = mock(ExpressionFunction.class);
+        testFunctionName = RandomStringUtils.randomAlphabetic(8);
+        testResultObject = mock(Object.class);
+        lenient().when(expressionFunction.evaluate(any(List.class))).thenReturn(testResultObject);
+        lenient().when(expressionFunction.getFunctionName()).thenReturn(testFunctionName);
+
+        return new ExpressionFunctionProvider(List.of(expressionFunction));
+    }
+
+    @Test
+    void testUnknownFunction() {
+        objectUnderTest = createObjectUnderTest();
+        String unknownFunctionName = RandomStringUtils.randomAlphabetic(8);
+        assertThat(objectUnderTest.provideFunction(unknownFunctionName, List.of()), equalTo(null));
+    }
+
+    @Test
+    void testFunctionBasic() {
+        objectUnderTest = createObjectUnderTest();
+        assertThat(objectUnderTest.provideFunction(testFunctionName, List.of()), equalTo(testResultObject));
+    }
+
+}

data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/LengthExpressionFunctionTest.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.expression;
+
+import org.junit.jupiter.api.Test;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import org.apache.commons.lang3.RandomStringUtils;
+
+import java.util.List;
+
+class LengthExpressionFunctionTest {
+    private LengthExpressionFunction lengthExpressionFunction;
+
+    public LengthExpressionFunction createObjectUnderTest() {
+        return new LengthExpressionFunction();
+    }
+
+    @Test
+    void testWithOneStringArgumentWithOutQuotes() {
+        lengthExpressionFunction = createObjectUnderTest();
+        String testString = RandomStringUtils.randomAlphabetic(5);
+        assertThat(lengthExpressionFunction.evaluate(List.of(testString)), equalTo(testString.length()));
+    }
+
+    @Test
+    void testWithOneStringArgumentWithQuotes() {
+        lengthExpressionFunction = createObjectUnderTest();
+        String testString = RandomStringUtils.randomAlphabetic(5);
+        assertThat(lengthExpressionFunction.evaluate(List.of("\""+testString + "\"")), equalTo(testString.length()));
+    }
+
+    @Test
+    void testWithTwoArgs() {
+        lengthExpressionFunction = createObjectUnderTest();
+        assertThrows(RuntimeException.class, () -> lengthExpressionFunction.evaluate(List.of("arg1", "arg2")));
+    }
+
+    @Test
+    void testWithNonStringArgument() {
+        lengthExpressionFunction = createObjectUnderTest();
+        assertThrows(RuntimeException.class, () -> lengthExpressionFunction.evaluate(List.of(10)));
+    }
+
+    @Test
+    void testWithInvalidStringArgument() {
+        lengthExpressionFunction = createObjectUnderTest();
+        assertThrows(RuntimeException.class, () -> lengthExpressionFunction.evaluate(List.of("\"arg1")));
+    }
+}

data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ParseTreeCoercionServiceTest.java

@@ -19,9 +19,11 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.dataprepper.expression.antlr.DataPrepperExpressionParser;
 import org.opensearch.dataprepper.expression.util.TestObject;
+import org.apache.commons.lang3.RandomStringUtils;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.List;
 import java.util.Random;
 import java.util.stream.Stream;
 
@@ -31,6 +33,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mock;
@@ -48,8 +51,9 @@ class ParseTreeCoercionServiceTest {
 
     private final LiteralTypeConversionsConfiguration literalTypeConversionsConfiguration =
             new LiteralTypeConversionsConfiguration();
+    private final ExpressionFunctionProvider expressionFunctionProvider = mock(ExpressionFunctionProvider.class);
     private final ParseTreeCoercionService objectUnderTest = new ParseTreeCoercionService(
-            literalTypeConversionsConfiguration.literalTypeConversions());
+            literalTypeConversionsConfiguration.literalTypeConversions(), expressionFunctionProvider);
 
     @Test
     void testCoerceTerminalNodeStringType() {
@@ -229,6 +233,42 @@ void testCoerceFailure() {
         assertThrows(ExpressionCoercionException.class, () -> objectUnderTest.coerce(testObj, String.class));
     }
 
+    @Test
+    void testCoerceTerminalNodeLengthFunction() {
+        final String key = RandomStringUtils.randomAlphabetic(5);
+        final String value = RandomStringUtils.randomAlphabetic(10);
+        final Event testEvent = createTestEvent(Map.of(key, value));
+        when(terminalNode.getSymbol()).thenReturn(token);
+        when(terminalNode.getText()).thenReturn("length(/"+key+")");
+        when(expressionFunctionProvider.provideFunction(eq("length"), any(List.class))).thenReturn(value.length());
+        when(token.getType()).thenReturn(DataPrepperExpressionParser.Function);
+        assertThat(objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent), equalTo(value.length()));
+        final String testString = RandomStringUtils.randomAlphabetic(10);
+        when(terminalNode.getText()).thenReturn("length(\""+testString+"\")");
+        when(expressionFunctionProvider.provideFunction(eq("length"), any(List.class))).thenReturn(testString.length());
+        assertThat(objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent), equalTo(testString.length()));
+    }
+
+    @Test
+    void testCoerceTerminalNodeLengthFunctionKeyNotInEvent() {
+        final String key = RandomStringUtils.randomAlphabetic(5);
+        final String value = RandomStringUtils.randomAlphabetic(10);
+        final String key2 = RandomStringUtils.randomAlphabetic(5);
+        final Event testEvent = createTestEvent(Map.of(key, value));
+        when(terminalNode.getSymbol()).thenReturn(token);
+        when(terminalNode.getText()).thenReturn("length(/"+key2+")");
+        when(token.getType()).thenReturn(DataPrepperExpressionParser.Function);
+        assertThat(objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent), equalTo(null));
+    }
+
+    @Test
+    void testCoerceTerminalNodeWithUnknownFunction() {
+        when(terminalNode.getSymbol()).thenReturn(token);
+        when(token.getType()).thenReturn(DataPrepperExpressionParser.Function);
+        when(terminalNode.getText()).thenReturn("xyz(arg1)");
+        assertThrows(RuntimeException.class, () -> objectUnderTest.coercePrimaryTerminalNode(terminalNode, null));
+    }
+
     private Event createTestEvent(final Object data) {
         final Event event = mock(Event.class);
         final JsonNode node = mapper.valueToTree(data);