Add JSON Support to V2 Engine

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

@@ -0,0 +1,106 @@
+/*
+ * 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.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<Void, JsonSupportVisitorContext> {
+  @Override
+  public Void visit(Node node, JsonSupportVisitorContext context) {
+    visitChildren(node, context);
+    return null;
+  }
+
+  @Override
+  public Void visitChildren(Node node, JsonSupportVisitorContext context) {
+    for (Node child : node.getChild()) {
+      child.accept(this, context);
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitAggregation(Aggregation node, JsonSupportVisitorContext context) {
+    if (!node.getGroupExprList().isEmpty()) {
+      throw new UnsupportedOperationException(
+          "Queries with aggregation are not yet supported with json format in the new engine");
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitFunction(Function node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // queries with function calls are not supported.
+      throw new UnsupportedOperationException(
+          "Queries with functions are not yet supported with json format in the new engine");
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitLiteral(Literal node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // queries with literal values are not supported
+      throw new UnsupportedOperationException(
+          "Queries with literals are not yet supported with json format in the new engine");
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitCast(Cast node, JsonSupportVisitorContext context) {
+    // Supported if outside of Project
+    if (context.isVisitingProject()) {
+      // Queries with cast are not supported
+      throw new UnsupportedOperationException(
+          "Queries with casts are not yet supported with json format in the new engine");
+    }
+    return null;
+  }
+
+  @Override
+  public Void 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())) {
+        node.getDelegated().accept(this, context);
+      } else {
+        throw new UnsupportedOperationException(
+            "Queries with aliases are not yet supported with json format in the new engine");
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitProject(Project node, JsonSupportVisitorContext context) {
+    visit(node, context);
+
+    context.setVisitingProject(true);
+    node.getProjectList().forEach(e -> e.accept(this, context));
+    context.setVisitingProject(false);
+    return null;
+  }
+}

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import lombok.Getter;
+import lombok.Setter;
+
+/**
+ * The context used for JsonSupportVisitor.
+ */
+public class JsonSupportVisitorContext {
+  @Getter
+  @Setter
+  private boolean isVisitingProject = false;
+}

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

@@ -45,4 +45,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(r -> r != null && !r.isEmpty()).findFirst().orElse("");
+  }
 }

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

@@ -0,0 +1,125 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+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.UnresolvedPlan;
+
+class JsonSupportVisitorTest {
+  @Test
+  public void visitLiteralInProject() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.intLiteral(1));
+    assertThrows(UnsupportedOperationException.class,
+        () -> project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitLiteralOutsideProject() {
+    Literal intLiteral = AstDSL.intLiteral(1);
+    assertNull(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")));
+    assertThrows(UnsupportedOperationException.class,
+        () -> project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitCastOutsideProject() {
+    UnresolvedExpression intCast = AstDSL.cast(
+        AstDSL.intLiteral(1),
+        AstDSL.stringLiteral("INT"));
+    assertNull(intCast.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasInProject() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.alias("alias", AstDSL.intLiteral(1)));
+    assertThrows(UnsupportedOperationException.class,
+        () -> 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"));
+    assertThrows(UnsupportedOperationException.class,
+        () -> project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasInProjectWithSupportedDelegated() {
+    UnresolvedPlan project = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.alias("alias", AstDSL.field("field")));
+    assertNull(project.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitAliasOutsideProject() {
+    UnresolvedExpression alias = AstDSL.alias("alias", AstDSL.intLiteral(1));
+    assertNull(alias.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitFunctionInProject() {
+    UnresolvedPlan function = AstDSL.project(
+        AstDSL.relation("table", "table"),
+        AstDSL.function("abs", AstDSL.intLiteral(-1)));
+    assertThrows(UnsupportedOperationException.class,
+        () -> function.accept(new JsonSupportVisitor(), new JsonSupportVisitorContext()));
+  }
+
+  @Test
+  public void visitFunctionOutsideProject() {
+    UnresolvedExpression function = AstDSL.function("abs", AstDSL.intLiteral(-1));
+    assertNull(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()));
+    assertThrows(UnsupportedOperationException.class,
+        () -> 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());
+    assertNull(aggregation.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());
+  }
 }