First prototype of json_extract

Signed-off-by: Hendrik Saly <[email protected]>
eliatra · Jun 28, 2024 · 7b50f3d · 7b50f3d
1 parent 8eae36f
commit 7b50f3d
Show file tree

Hide file tree

Showing 10 changed files with 183 additions and 0 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/DSL.java b/core/src/main/java/org/opensearch/sql/expression/DSL.java
@@ -404,6 +404,14 @@ public static FunctionExpression get_format(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.GET_FORMAT, expressions);
   }
 
+  public static FunctionExpression json_extract(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.JSON_EXTRACT, expressions);
+  }
+
+  public static FunctionExpression json_keys(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.JSON_KEYS, expressions);
+  }
+
   public static FunctionExpression hour(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.HOUR, expressions);
   }

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
@@ -201,6 +201,10 @@ public enum BuiltinFunctionName {
   TRIM(FunctionName.of("trim")),
   UPPER(FunctionName.of("upper")),
 
+  /** Json Functions. */
+  JSON_EXTRACT(FunctionName.of("json_extract")),
+  JSON_KEYS(FunctionName.of("json_keys")),
+
   /** NULL Test. */
   IS_NULL(FunctionName.of("is null")),
   IS_NOT_NULL(FunctionName.of("is not null")),

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionRepository.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionRepository.java
@@ -27,6 +27,7 @@
 import org.opensearch.sql.expression.aggregation.AggregatorFunction;
 import org.opensearch.sql.expression.datetime.DateTimeFunction;
 import org.opensearch.sql.expression.datetime.IntervalClause;
+import org.opensearch.sql.expression.json.JsonFunction;
 import org.opensearch.sql.expression.operator.arthmetic.ArithmeticFunction;
 import org.opensearch.sql.expression.operator.arthmetic.MathematicalFunction;
 import org.opensearch.sql.expression.operator.convert.TypeCastOperator;
@@ -81,6 +82,7 @@ public static synchronized BuiltinFunctionRepository getInstance() {
       TypeCastOperator.register(instance);
       SystemFunctions.register(instance);
       OpenSearchFunctions.register(instance);
+      JsonFunction.register(instance);
     }
     return instance;
   }

diff --git a/core/src/main/java/org/opensearch/sql/expression/json/JsonFunction.java b/core/src/main/java/org/opensearch/sql/expression/json/JsonFunction.java
@@ -0,0 +1,118 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.json;
+
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.expression.function.FunctionDSL.define;
+import static org.opensearch.sql.expression.function.FunctionDSL.impl;
+import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandling;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.CharMatcher;
+import com.jayway.jsonpath.Configuration;
+import com.jayway.jsonpath.JsonPath;
+import com.jayway.jsonpath.Option;
+import com.jayway.jsonpath.spi.json.JacksonJsonProvider;
+import com.jayway.jsonpath.spi.json.JsonProvider;
+import com.jayway.jsonpath.spi.mapper.JacksonMappingProvider;
+import com.jayway.jsonpath.spi.mapper.MappingProvider;
+import java.util.Set;
+import lombok.experimental.UtilityClass;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.model.ExprValueUtils;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
+import org.opensearch.sql.expression.function.DefaultFunctionResolver;
+import org.opensearch.sql.expression.function.SerializableBiFunction;
+
+/**
+ * The definition of json functions. 1) have the clear interface for function define. 2) the
+ * implementation should rely on ExprValue.
+ */
+@UtilityClass
+public class JsonFunction {
+
+  private static final ObjectMapper MAPPER = new ObjectMapper();
+
+  static {
+    Configuration.setDefaults(
+        new Configuration.Defaults() {
+
+          private final JsonProvider jsonProvider = new JacksonJsonProvider();
+          private final MappingProvider mappingProvider = new JacksonMappingProvider();
+
+          @Override
+          public JsonProvider jsonProvider() {
+            return jsonProvider;
+          }
+
+          @Override
+          public MappingProvider mappingProvider() {
+            return mappingProvider;
+          }
+
+          @Override
+          public Set<Option> options() {
+            return Set.of(Option.SUPPRESS_EXCEPTIONS);
+          }
+        });
+  }
+
+  /**
+   * Register Json Functions.
+   *
+   * @param repository {@link BuiltinFunctionRepository}.
+   */
+  public void register(BuiltinFunctionRepository repository) {
+    repository.register(extract());
+    // repository.register(keys());
+  }
+
+  private DefaultFunctionResolver extract() {
+    return define(
+        BuiltinFunctionName.JSON_EXTRACT.getName(),
+        impl(
+            nullMissingHandling(
+                (SerializableBiFunction<ExprValue, ExprValue, ExprValue>)
+                    (jsonExpr, pathExpr) -> {
+                      String jsonInput = jsonExpr.value().toString();
+                      String path = CharMatcher.is('\"').trimFrom(pathExpr.toString());
+
+                      try {
+
+                        // when path start with $ we interpret it as json path
+                        // everything else is interpreted as json pointer
+                        if (path.startsWith("$")) {
+                          Object retVal = JsonPath.parse(jsonInput).read(path);
+
+                          if (retVal instanceof String) {
+                            return ExprValueUtils.fromObjectValue(retVal);
+                          }
+
+                          String retVal2 = MAPPER.writeValueAsString(retVal);
+                          return ExprValueUtils.fromObjectValue(retVal2);
+                        } else {
+                          JsonNode retVal = MAPPER.readTree(jsonInput).at(path);
+                          if (retVal.isValueNode()) {
+                            return ExprValueUtils.fromObjectValue(retVal.asText());
+                          }
+
+                          return ExprValueUtils.fromObjectValue(retVal.toString());
+                        }
+                      } catch (Exception e) {
+                        throw new RuntimeException(e);
+                      }
+                    }),
+            STRING,
+            STRING,
+            STRING));
+  }
+
+  private DefaultFunctionResolver keys() {
+    return null;
+  }
+}
diff --git a/sql/src/main/antlr/OpenSearchSQLLexer.g4 b/sql/src/main/antlr/OpenSearchSQLLexer.g4
@@ -349,6 +349,10 @@ WILDCARD_QUERY:                     'WILDCARD_QUERY';
 SUBSTR:                             'SUBSTR';
 STRCMP:                             'STRCMP';
 
+// JSON FUNCTIONS
+JSON_EXTRACT:                       'JSON_EXTRACT';
+JSON_KEYS:                          'JSON_KEYS';
+
 // DATE AND TIME FUNCTIONS
 ADDDATE:                            'ADDDATE';
 YEARWEEK:                           'YEARWEEK';

diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4
@@ -335,6 +335,7 @@ functionCall
    | extractFunction                                            # extractFunctionCall
    | getFormatFunction                                          # getFormatFunctionCall
    | timestampFunction                                          # timestampFunctionCall
+   | jsonFunctionName LR_BRACKET functionArgs RR_BRACKET        # jsonFunctionCall
    ;
 
 timestampFunction
@@ -639,6 +640,11 @@ textFunctionName
    | REVERSE
    ;
 
+jsonFunctionName
+   : JSON_EXTRACT
+   | JSON_KEYS
+   ;
+
 flowControlFunctionName
    : IF
    | IFNULL
@@ -825,6 +831,7 @@ ident
    | BACKTICK_QUOTE_ID
    | keywordsCanBeId
    | scalarFunctionName
+   | jsonFunctionName
    ;
 
 keywordsCanBeId

diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java
@@ -143,6 +143,12 @@ public UnresolvedExpression visitGetFormatFunctionCall(GetFormatFunctionCallCont
         ctx.getFormatFunction().GET_FORMAT().toString(), getFormatFunctionArguments(ctx));
   }
 
+  @Override
+  public UnresolvedExpression visitJsonFunctionCall(
+      OpenSearchSQLParser.JsonFunctionCallContext ctx) {
+    return buildFunction(ctx.jsonFunctionName().getText(), ctx.functionArgs().functionArg());
+  }
+
   @Override
   public UnresolvedExpression visitHighlightFunctionCall(HighlightFunctionCallContext ctx) {
     ImmutableMap.Builder<String, Literal> builder = ImmutableMap.builder();

diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java
@@ -75,6 +75,23 @@ public void onFailure(Exception e) {
         });
   }
 
+  @Test
+  public void can_execute_sql_json_query() throws InterruptedException {
+    sqlService.execute(
+        new SQLQueryRequest(new JSONObject(), "SELECT json_extract('123','$.name')", QUERY, "jdbc"),
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(QueryResponse response) {
+            assertNotNull(response);
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            fail(e);
+          }
+        });
+  }
+
   @Test
   public void can_execute_cursor_query() {
     sqlService.execute(

diff --git a/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java b/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java
@@ -270,6 +270,12 @@ public void cannot_parse_get_format_function_with_bad_arg() {
         SyntaxCheckException.class, () -> parser.parse("GET_FORMAT(NONSENSE_ARG,'INTERNAL')"));
   }
 
+  @Test
+  public void can_parse_json_function() {
+    assertNotNull(parser.parse("SELECT JSON_EXTRACT(abc, '$.name')"));
+    assertNotNull(parser.parse("SELECT JSON_KEYS('$.name')"));
+  }
+
   @Test
   public void can_parse_hour_functions() {
     assertNotNull(parser.parse("SELECT hour('2022-11-18 12:23:34')"));

diff --git a/sql/src/test/java/org/opensearch/sql/sql/parser/AstExpressionBuilderTest.java b/sql/src/test/java/org/opensearch/sql/sql/parser/AstExpressionBuilderTest.java
@@ -153,6 +153,17 @@ public void canBuildGetFormatFunctionCall() {
         buildExprAst("get_format(DATE,\"USA\")"));
   }
 
+  @Test
+  public void canBuildJsonFunctionCall() {
+    assertEquals(function("json_keys", stringLiteral("abc")), buildExprAst("json_keys(\"abc\")"));
+
+    assertEquals(function("json_extract"), buildExprAst("json_extract()"));
+
+    assertEquals(
+        function("json_extract", stringLiteral("abc"), stringLiteral("name")),
+        buildExprAst("json_extract(\"abc\",\"name\")"));
+  }
+
   @Test
   public void canBuildNestedFunctionCall() {
     assertEquals(