Add JSON Support to V2 Engine (#217)

core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitor.java

@@ -0,0 +1,111 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import org.apache.commons.lang3.StringUtils;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
+import org.opensearch.sql.ast.expression.Alias;
+import org.opensearch.sql.ast.expression.Cast;
+import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.Literal;
+import org.opensearch.sql.ast.tree.Aggregation;
+import org.opensearch.sql.ast.tree.Limit;
+import org.opensearch.sql.ast.tree.Project;
+
+/**
+ * This visitor's sole purpose is to throw UnsupportedOperationExceptions
+ * for unsupported features in the new engine when JSON format is specified.
+ * Unsupported features in V2 are ones the produce results that differ from
+ * legacy results.
+ */
+public class JsonSupportVisitor extends AbstractNodeVisitor<Boolean, JsonSupportVisitorContext> {
+  @Override
+  public Boolean visit(Node node, JsonSupportVisitorContext context) {
+    // UnresolvedPlan can only have one child until joins are supported
+    return node.getChild().get(0).accept(this, context);
+  }
+
+  @Override
+  protected Boolean defaultResult() {
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitLimit(Limit node, JsonSupportVisitorContext context) {
+    context.addToUnsupportedNodes("limit");
+    return Boolean.FALSE;
+  }
+
+  @Override
+  public Boolean visitAggregation(Aggregation node, JsonSupportVisitorContext context) {
+    if (!node.getGroupExprList().isEmpty()) {
+      context.addToUnsupportedNodes("aggregation");
+      return Boolean.FALSE;
+    }
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitFunction(Function node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // queries with function calls are not supported.
+      context.addToUnsupportedNodes("functions");
+      return Boolean.FALSE;
+    }
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitLiteral(Literal node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // queries with literal values are not supported
+      context.addToUnsupportedNodes("literal");
+      return Boolean.FALSE;
+    }
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitCast(Cast node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // Queries with cast are not supported
+      context.addToUnsupportedNodes("cast");
+      return Boolean.FALSE;
+    }
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitAlias(Alias node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // Alias node is accepted if it does not have a user-defined alias
+      // and if the delegated expression is accepted.
+      if (StringUtils.isEmpty(node.getAlias())) {
+        return node.getDelegated().accept(this, context);
+      } else {
+        context.addToUnsupportedNodes("alias");
+        return Boolean.FALSE;
+      }
+    }
+    return Boolean.TRUE;
+  }
+
+  @Override
+  public Boolean visitProject(Project node, JsonSupportVisitorContext context) {
+    Boolean isSupported = visit(node, context);
+
+    context.setVisitingProject(true);
+    isSupported = node.getProjectList().stream()
+        .allMatch(e -> e.accept(this, context)) && isSupported;
+    context.setVisitingProject(false);
+    return isSupported;
+  }
+}

core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitorContext.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import java.util.ArrayList;
+import java.util.List;
+import lombok.Getter;
+import lombok.Setter;
+
+/**
+ * The context used for JsonSupportVisitor.
+ */
+public class JsonSupportVisitorContext {
+  @Getter
+  @Setter
+  private boolean isVisitingProject = false;
+
+  @Getter
+  @Setter
+  private List<String> unsupportedNodes = new ArrayList<>();
+
+  public void addToUnsupportedNodes(String unsupportedNode) {
+    unsupportedNodes.add(unsupportedNode);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -82,7 +82,7 @@ public T visitChildren(Node node, C context) {
     return result;
   }
 
-  private T defaultResult() {
+  protected T defaultResult() {
     return null;
   }
 

core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java

@@ -50,9 +50,23 @@ void execute(PhysicalPlan plan, ExecutionContext context,
    * Data class that encapsulates ExprValue.
    */
   @Data
+  @RequiredArgsConstructor
   class QueryResponse {
     private final Schema schema;
     private final List<ExprValue> results;
+    private final String rawResponse; // JSON response from the OpenSearch instance
+
+    /**
+     * Constructor for Query Response.
+     *
+     * @param schema  schema of the query
+     * @param results list of expressions
+     */
+    public QueryResponse(Schema schema, List<ExprValue> results) {
+      this.schema = schema;
+      this.results = results;
+      this.rawResponse = "";
+    }
   }
 
   @Data

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java

@@ -7,6 +7,7 @@
 package org.opensearch.sql.planner.physical;
 
 import java.util.Iterator;
+import org.apache.commons.lang3.StringUtils;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.planner.PlanNode;
@@ -45,4 +46,9 @@ public ExecutionEngine.Schema schema() {
     throw new IllegalStateException(String.format("[BUG] schema can been only applied to "
         + "ProjectOperator, instead of %s", toString()));
   }
+
+  public String getRawResponse() {
+    return getChild().stream().map(PhysicalPlan::getRawResponse)
+        .filter(StringUtils::isNotEmpty).findFirst().orElse("");
+  }
 }

core/src/test/java/org/opensearch/sql/analysis/JsonSupportVisitorTest.java

@@ -0,0 +1,135 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
+
+import com.google.common.collect.ImmutableList;
+import java.util.Collections;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.ast.dsl.AstDSL;
+import org.opensearch.sql.ast.expression.Literal;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.ast.tree.Limit;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+
+class JsonSupportVisitorTest {
+  @Test
+  public void visitLiteralInProject() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.intLiteral(1));
+    assertFalse(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitLiteralOutsideProject() {
+    Literal intLiteral = AstDSL.intLiteral(1);
+    assertTrue(intLiteral.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitCastInProject() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.cast(AstDSL.intLiteral(1), AstDSL.stringLiteral("INT")));
+    assertFalse(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitCastOutsideProject() {
+    UnresolvedExpression intCast = AstDSL.cast(
+        AstDSL.intLiteral(1),
+        AstDSL.stringLiteral("INT"));
+    assertTrue(intCast.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasInProject() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.alias("alias", AstDSL.intLiteral(1)));
+    assertFalse(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasInProjectWithUnsupportedDelegated() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.alias("alias", AstDSL.intLiteral(1), "alias"));
+    assertFalse(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasInProjectWithSupportedDelegated() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.alias("alias", AstDSL.field("field")));
+    assertTrue(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasOutsideProject() {
+    UnresolvedExpression alias = AstDSL.alias("alias", AstDSL.intLiteral(1));
+    assertTrue(alias.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitFunctionInProject() {
+    UnresolvedPlan function = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.function("abs", AstDSL.intLiteral(-1)));
+    assertFalse(function.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitFunctionOutsideProject() {
+    UnresolvedExpression function = AstDSL.function("abs", AstDSL.intLiteral(-1));
+    assertTrue(function.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAggregationWithGroupExprList() {
+    UnresolvedPlan projectAggr = AstDSL.project(AstDSL.agg(
+        AstDSL.relation("table", "table"),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        ImmutableList.of(AstDSL.alias("alias", qualifiedName("integer_value"))),
+        Collections.emptyList()));
+    assertFalse(projectAggr.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAggregationWithAggExprList() {
+    UnresolvedPlan aggregation = AstDSL.agg(
+        AstDSL.relation("table", "table"),
+        ImmutableList.of(
+            AstDSL.alias("AVG(alias)",
+                AstDSL.aggregate("AVG",
+                    qualifiedName("integer_value")))),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList());
+    assertTrue(aggregation.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitLimit() {
+    Limit limit = AstDSL.limit(AstDSL.relation("table", "table"), 10, 5);
+    assertFalse(limit.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitWithMultipleUnsupportedProjectNodes() {
+    UnresolvedPlan plan = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.function("abs", AstDSL.intLiteral(-1)),
+        AstDSL.alias("alias", AstDSL.intLiteral(1)));
+    assertFalse(plan.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+}

core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java

@@ -5,7 +5,9 @@
 
 package org.opensearch.sql.planner.physical;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.util.List;
 import org.junit.jupiter.api.Test;
@@ -50,4 +52,12 @@ void addSplitToChildByDefault() {
     testPlan.add(split);
     verify(child).add(split);
   }
+
+  @Test
+  void testGetRawResponse() {
+    when(child.getRawResponse()).thenReturn("not empty", "", null);
+    assertEquals("not empty", testPlan.getRawResponse());
+    assertEquals("", testPlan.getRawResponse());
+    assertEquals("", testPlan.getRawResponse());
+  }
 }