feat: Add JsonUpdate and JsonQuery scalar udfs

dagger-functions/src/main/java/io/odpf/dagger/functions/udfs/factories/FunctionFactory.java

@@ -40,6 +40,8 @@
 import io.odpf.dagger.functions.udfs.scalar.StartOfWeek;
 import io.odpf.dagger.functions.udfs.scalar.TimeInDate;
 import io.odpf.dagger.functions.udfs.scalar.TimestampFromUnix;
+import io.odpf.dagger.functions.udfs.scalar.JsonQuery;
+import io.odpf.dagger.functions.udfs.scalar.JsonUpdate;
 import io.odpf.dagger.functions.udfs.table.HistogramBucket;
 import io.odpf.dagger.functions.udfs.table.OutlierMad;
 
@@ -105,6 +107,8 @@ public HashSet<ScalarUdf> getScalarUdfs() {
         scalarUdfs.add(new ArrayAggregate());
         scalarUdfs.add(new ArrayOperate());
         scalarUdfs.add(new ByteToString());
+        scalarUdfs.add(new JsonQuery());
+        scalarUdfs.add(new JsonUpdate());
         return scalarUdfs;
     }
 

dagger-functions/src/main/java/io/odpf/dagger/functions/udfs/scalar/JsonQuery.java

@@ -0,0 +1,35 @@
+package io.odpf.dagger.functions.udfs.scalar;
+
+import com.jayway.jsonpath.Configuration;
+import com.jayway.jsonpath.JsonPath;
+import com.jayway.jsonpath.Option;
+import io.odpf.dagger.common.udfs.ScalarUdf;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonInclude;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.core.JsonProcessingException;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.table.annotation.DataTypeHint;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/**
+ * The Json query udf.
+ */
+public class JsonQuery extends ScalarUdf implements Serializable {
+
+    /**
+     * Extracts JSON values from a JSON string.
+     * The result is always returned as a JSON STRING.
+     *
+     * @param jsonEvent the json String
+     * @param jPath     the jPath
+     * @return jsonString
+     */
+    public @DataTypeHint("STRING") String eval(@DataTypeHint("STRING") String jsonEvent, @DataTypeHint("STRING") String jPath) throws JsonProcessingException {
+        Configuration configuration = Configuration.defaultConfiguration().setOptions(Option.DEFAULT_PATH_LEAF_TO_NULL);
+        Object jChildObject = JsonPath.using(configuration).parse(jsonEvent).read(JsonPath.compile(jPath));
+        ObjectMapper mapper = new ObjectMapper();
+        mapper.setSerializationInclusion(JsonInclude.Include.USE_DEFAULTS);
+        return Objects.isNull(jChildObject) ? null : mapper.writeValueAsString(jChildObject);
+    }
+}

dagger-functions/src/main/java/io/odpf/dagger/functions/udfs/scalar/JsonUpdate.java

@@ -0,0 +1,31 @@
+package io.odpf.dagger.functions.udfs.scalar;
+
+import com.jayway.jsonpath.Configuration;
+import com.jayway.jsonpath.JsonPath;
+import com.jayway.jsonpath.Option;
+import com.jayway.jsonpath.PathNotFoundException;
+import io.odpf.dagger.common.udfs.ScalarUdf;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.InputGroup;
+
+import java.io.Serializable;
+
+/**
+ * The Json update udf.
+ */
+public class JsonUpdate extends ScalarUdf implements Serializable {
+
+    /**
+     * Updates JSON  from a JSON path.
+     * The result is always returned as a JSON STRING.
+     *
+     * @param jsonEvent the json String
+     * @param jPath     the jPath
+     * @param newValue       the new object value
+     * @return jsonString
+     */
+    public @DataTypeHint("STRING") String eval(String jsonEvent, String jPath, @DataTypeHint(inputGroup = InputGroup.ANY) Object newValue) throws PathNotFoundException {
+        Configuration configuration = Configuration.defaultConfiguration().setOptions(Option.DEFAULT_PATH_LEAF_TO_NULL);
+        return JsonPath.using(configuration).parse(jsonEvent).set(JsonPath.compile(jPath), newValue).jsonString();
+    }
+}

dagger-functions/src/test/java/io/odpf/dagger/functions/udfs/factories/FunctionFactoryTest.java

@@ -1,5 +1,6 @@
 package io.odpf.dagger.functions.udfs.factories;
 
+import io.odpf.dagger.functions.udfs.scalar.*;
 import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
 
 import io.odpf.dagger.common.configuration.Configuration;
@@ -11,32 +12,6 @@
 import io.odpf.dagger.functions.udfs.aggregate.Features;
 import io.odpf.dagger.functions.udfs.aggregate.FeaturesWithType;
 import io.odpf.dagger.functions.udfs.aggregate.PercentileAggregator;
-import io.odpf.dagger.functions.udfs.scalar.ArrayAggregate;
-import io.odpf.dagger.functions.udfs.scalar.ArrayOperate;
-import io.odpf.dagger.functions.udfs.scalar.ByteToString;
-import io.odpf.dagger.functions.udfs.scalar.CondEq;
-import io.odpf.dagger.functions.udfs.scalar.DartContains;
-import io.odpf.dagger.functions.udfs.scalar.DartGet;
-import io.odpf.dagger.functions.udfs.scalar.Distance;
-import io.odpf.dagger.functions.udfs.scalar.ElementAt;
-import io.odpf.dagger.functions.udfs.scalar.EndOfMonth;
-import io.odpf.dagger.functions.udfs.scalar.EndOfWeek;
-import io.odpf.dagger.functions.udfs.scalar.ExponentialMovingAverage;
-import io.odpf.dagger.functions.udfs.scalar.Filters;
-import io.odpf.dagger.functions.udfs.scalar.FormatTimeInZone;
-import io.odpf.dagger.functions.udfs.scalar.GeoHash;
-import io.odpf.dagger.functions.udfs.scalar.LinearTrend;
-import io.odpf.dagger.functions.udfs.scalar.ListContains;
-import io.odpf.dagger.functions.udfs.scalar.MapGet;
-import io.odpf.dagger.functions.udfs.scalar.S2AreaInKm2;
-import io.odpf.dagger.functions.udfs.scalar.S2Id;
-import io.odpf.dagger.functions.udfs.scalar.SelectFields;
-import io.odpf.dagger.functions.udfs.scalar.SingleFeatureWithType;
-import io.odpf.dagger.functions.udfs.scalar.Split;
-import io.odpf.dagger.functions.udfs.scalar.StartOfMonth;
-import io.odpf.dagger.functions.udfs.scalar.StartOfWeek;
-import io.odpf.dagger.functions.udfs.scalar.TimeInDate;
-import io.odpf.dagger.functions.udfs.scalar.TimestampFromUnix;
 import io.odpf.dagger.functions.udfs.table.HistogramBucket;
 import io.odpf.dagger.functions.udfs.table.OutlierMad;
 import org.junit.Assert;
@@ -117,6 +92,8 @@ public void shouldRegisterScalarUdfs() {
         verify(streamTableEnvironment, times(1)).createTemporaryFunction(eq("ArrayAggregate"), any(ArrayAggregate.class));
         verify(streamTableEnvironment, times(1)).createTemporaryFunction(eq("ArrayOperate"), any(ArrayOperate.class));
         verify(streamTableEnvironment, times(1)).createTemporaryFunction(eq("ByteToString"), any(ByteToString.class));
+        verify(streamTableEnvironment, times(1)).createTemporaryFunction(eq("JsonUpdate"), any(JsonUpdate.class));
+        verify(streamTableEnvironment, times(1)).createTemporaryFunction(eq("JsonQuery"), any(JsonQuery.class));
     }
 
     @Test

dagger-functions/src/test/java/io/odpf/dagger/functions/udfs/scalar/JsonQueryTest.java

@@ -0,0 +1,117 @@
+package io.odpf.dagger.functions.udfs.scalar;
+
+import org.apache.flink.metrics.Gauge;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.core.JsonProcessingException;
+import org.apache.flink.table.functions.FunctionContext;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.Mock;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.*;
+import static org.mockito.MockitoAnnotations.initMocks;
+
+public class JsonQueryTest {
+
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+    @Mock
+    private MetricGroup metricGroup;
+
+    @Mock
+    private FunctionContext functionContext;
+
+    @Before
+    public void setup() {
+        initMocks(this);
+        when(functionContext.getMetricGroup()).thenReturn(metricGroup);
+        when(metricGroup.addGroup("udf", "JsonQuery")).thenReturn(metricGroup);
+    }
+
+    @Test
+    public void shouldRegisterGauge() throws Exception {
+        JsonQuery jsonQuery = new JsonQuery();
+        jsonQuery.open(functionContext);
+        verify(metricGroup, times(1)).gauge(any(String.class), any(Gauge.class));
+    }
+
+    @Test
+    public void shouldReturnJsonString() throws JsonProcessingException {
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String expectedJsonEvent = "\"v2\"";
+        String jPath = "$.k2";
+        Assert.assertEquals(expectedJsonEvent, jsonQuery.eval(jsonEvent, jPath));
+    }
+
+    @Test
+    public void shouldReturnJsonStringForNestedJson() throws JsonProcessingException {
+        JsonQuery jsonQuery = new JsonQuery();
+        String expectedJsonEvent = "{\"key1\":\"value1\",\"key2\":\"value2\"}";
+        String jsonEvent = "{\"k1\":null,\"k2\":{\"key1\":\"value1\",\"key2\":\"value2\"}}";
+        String jPath = "$.k2";
+        Assert.assertEquals(expectedJsonEvent, jsonQuery.eval(jsonEvent, jPath));
+    }
+
+    @Test
+    public void shouldReturnJsonStringForNullValue() throws JsonProcessingException {
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String jPath = "$.k1";
+        String result = jsonQuery.eval(jsonEvent, jPath);
+        Assert.assertNull(result);
+    }
+
+    @Test
+    public void shouldThrowErrorWhenNullJsonEvent() throws JsonProcessingException {
+        thrown.expect(IllegalArgumentException.class);
+        thrown.expectMessage("json string can not be null or empty");
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = null;
+        String jPath = "$.k2";
+        jsonQuery.eval(jsonEvent, jPath);
+    }
+
+    @Test
+    public void shouldThrowErrorWhenNullJPath() throws JsonProcessingException {
+        thrown.expect(IllegalArgumentException.class);
+        thrown.expectMessage("json can not be null or empty");
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String jPath = null;
+        jsonQuery.eval(jsonEvent, jPath);
+    }
+
+    @Test
+    public void shouldThrowErrorWhenEmptyJPath() throws JsonProcessingException {
+        thrown.expect(IllegalArgumentException.class);
+        thrown.expectMessage("json can not be null or empty");
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String jPath = "";
+        jsonQuery.eval(jsonEvent, jPath);
+    }
+
+    @Test
+    public void shouldThrowErrorWhenImproperJPath() throws JsonProcessingException {
+        thrown.expect(com.jayway.jsonpath.InvalidPathException.class);
+        thrown.expectMessage("Illegal character at position 1 expected '.' or '[");
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String jPath = "$ .k1";
+        jsonQuery.eval(jsonEvent, jPath);
+    }
+
+    @Test
+    public void shouldThrowErrorWhenInvalidJPath() throws JsonProcessingException {
+        thrown.expect(com.jayway.jsonpath.PathNotFoundException.class);
+        JsonQuery jsonQuery = new JsonQuery();
+        String jsonEvent = "{\"k1\":null,\"k2\":\"v2\"}";
+        String jPath = "$.k2.k4";
+        jsonQuery.eval(jsonEvent, jPath);
+    }
+}