diff --git a/core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitor.java b/core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitor.java new file mode 100644 index 00000000000..c611f988ce7 --- /dev/null +++ b/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 { + @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; + } +} diff --git a/core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitorContext.java b/core/src/main/java/org/opensearch/sql/analysis/JsonSupportVisitorContext.java new file mode 100644 index 00000000000..f50596fdd58 --- /dev/null +++ b/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 unsupportedNodes = new ArrayList<>(); + + public void addToUnsupportedNodes(String unsupportedNode) { + unsupportedNodes.add(unsupportedNode); + } +} diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index 393de051649..f4325fa7f64 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/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; } diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java index 1936a0f5178..4db11e86e24 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java +++ b/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 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 results) { + this.schema = schema; + this.results = results; + this.rawResponse = ""; + } } @Data diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java index b476b015577..068260ea331 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java +++ b/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(""); + } } diff --git a/core/src/test/java/org/opensearch/sql/analysis/JsonSupportVisitorTest.java b/core/src/test/java/org/opensearch/sql/analysis/JsonSupportVisitorTest.java new file mode 100644 index 00000000000..3409900450a --- /dev/null +++ b/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())); + } +} diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java index 0a93c96bbb2..5efe8390b42 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java +++ b/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()); + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/DateFormatIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/DateFormatIT.java index a0b4b198985..770318a577c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/DateFormatIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/DateFormatIT.java @@ -28,6 +28,7 @@ import org.json.JSONObject; import org.junit.Ignore; import org.junit.Test; +import org.opensearch.client.ResponseException; import org.opensearch.sql.legacy.exception.SqlParseException; public class DateFormatIT extends SQLIntegTestCase { @@ -50,7 +51,9 @@ protected void init() throws Exception { * this ends up excluding some of the expected values causing the assertion to fail. LIMIT overrides this. */ - @Test + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 + @Test(expected = SqlParseException.class) public void equalTo() throws SqlParseException { assertThat( dateQuery( @@ -59,7 +62,9 @@ public void equalTo() throws SqlParseException { ); } - @Test + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 + @Test(expected = SqlParseException.class) public void lessThan() throws SqlParseException { assertThat( dateQuery( @@ -68,6 +73,8 @@ public void lessThan() throws SqlParseException { ); } + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 @Test public void lessThanOrEqualTo() throws SqlParseException { assertThat( @@ -79,7 +86,9 @@ public void lessThanOrEqualTo() throws SqlParseException { ); } - @Test + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 + @Test(expected = SqlParseException.class) public void greaterThan() throws SqlParseException { assertThat( dateQuery( @@ -88,6 +97,8 @@ public void greaterThan() throws SqlParseException { ); } + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 @Test public void greaterThanOrEqualTo() throws SqlParseException { assertThat( @@ -99,6 +110,8 @@ public void greaterThanOrEqualTo() throws SqlParseException { ); } + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 @Test public void and() throws SqlParseException { assertThat( @@ -122,6 +135,8 @@ public void andWithDefaultTimeZone() throws SqlParseException { ); } + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 @Test public void or() throws SqlParseException { assertThat( @@ -134,7 +149,8 @@ public void or() throws SqlParseException { ); } - + // DATE_FORMAT with timezone argument is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1436 @Test public void sortByDateFormat() throws IOException { // Sort by expression in descending order, but sort inside in ascending order, so we increase our confidence diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/QueryFunctionsIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/QueryFunctionsIT.java index c59974a6225..57a0a722787 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/QueryFunctionsIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/QueryFunctionsIT.java @@ -28,6 +28,7 @@ import org.json.JSONObject; import org.junit.Test; import org.opensearch.action.search.SearchResponse; +import org.opensearch.client.ResponseException; import org.opensearch.common.xcontent.LoggingDeprecationHandler; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentType; @@ -133,7 +134,9 @@ public void scoreQueryWithNestedField() throws IOException { ); } - @Test + // Specifying city.keyword is not supported in V2 + // Rework on data types is in progress and tracked with https://github.com/opensearch-project/sql/pull/1314 + @Test(expected = ResponseException.class) public void wildcardQuery() throws IOException { assertThat( query( diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/QueryIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/QueryIT.java index eeff107f154..d70d3d1ff55 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/QueryIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/QueryIT.java @@ -102,6 +102,8 @@ public void multipleFromTest() throws IOException { Assert.assertEquals(14, getTotalHits(response)); } + // Duplicated column names like SELECT *, age is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/785 @Test public void selectAllWithFieldReturnsAll() throws IOException { JSONObject response = executeQuery(StringUtils.format( @@ -114,6 +116,8 @@ public void selectAllWithFieldReturnsAll() throws IOException { checkSelectAllAndFieldResponseSize(response); } + // Duplicated column names like SELECT *, age is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/785 @Test public void selectAllWithFieldReverseOrder() throws IOException { JSONObject response = executeQuery(StringUtils.format( @@ -126,6 +130,8 @@ public void selectAllWithFieldReverseOrder() throws IOException { checkSelectAllAndFieldResponseSize(response); } + // Duplicated column names like SELECT *, age is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/785 @Test public void selectAllWithMultipleFields() throws IOException { JSONObject response = executeQuery(StringUtils.format( @@ -138,6 +144,8 @@ public void selectAllWithMultipleFields() throws IOException { checkSelectAllAndFieldResponseSize(response); } + // Duplicated column names like SELECT *, age is not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/785 @Test public void selectAllWithFieldAndOrderBy() throws IOException { JSONObject response = executeQuery(StringUtils.format( @@ -577,7 +585,7 @@ public void notBetweenTest() throws IOException { public void inTest() throws IOException { JSONObject response = executeQuery( String.format(Locale.ROOT, "SELECT age FROM %s WHERE age IN (20, 22) LIMIT 1000", - TestsConstants.TEST_INDEX_PHRASE)); + TestsConstants.TEST_INDEX_PEOPLE)); JSONArray hits = getHits(response); for (int i = 0; i < hits.length(); i++) { @@ -725,7 +733,9 @@ public void notInTest() throws IOException { } } - @Test + // Unsupported date format in V2 + // Can be tracked through https://github.com/opensearch-project/sql/issues/856 + @Test(expected = ResponseException.class) public void dateSearch() throws IOException { DateTimeFormatter formatter = DateTimeFormat.forPattern(TestsConstants.DATE_FORMAT); DateTime dateToCompare = new DateTime(2014, 8, 18, 0, 0, 0); @@ -770,6 +780,8 @@ public void dateSearchBraces() throws IOException { } } + // Unsupported date format in V2 + // Can be tracked through https://github.com/opensearch-project/sql/issues/856 @Test public void dateBetweenSearch() throws IOException { DateTimeFormatter formatter = DateTimeFormat.forPattern(TestsConstants.DATE_FORMAT); @@ -1270,6 +1282,8 @@ public void queryWithDotAtStartOfIndexName() throws Exception { Assert.assertTrue(response.contains("PhD")); } + // Objects in IS NOT NULL not yet supported in V2 + // Can be tracked with https://github.com/opensearch-project/sql/issues/1425 @Test public void notLikeTests() throws IOException { JSONObject response = executeQuery( diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java index bc97f71b476..818c1b5e705 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java @@ -126,7 +126,7 @@ public void onResponse(T response) { @Override public void onFailure(Exception e) { - if (e instanceof SyntaxCheckException) { + if (e instanceof SyntaxCheckException || e instanceof UnsupportedOperationException) { fallBackHandler.accept(channel, e); } else { next.onFailure(e); @@ -165,10 +165,22 @@ private ResponseListener createQueryResponseListener( formatter = new CsvResponseFormatter(request.sanitize()); } else if (format.equals(Format.RAW)) { formatter = new RawResponseFormatter(); + } else if (format.equals(Format.JSON)) { + return new ResponseListener<>() { + @Override + public void onResponse(QueryResponse response) { + sendResponse(channel, OK, response.getRawResponse()); + } + + @Override + public void onFailure(Exception e) { + errorHandler.accept(channel, e); + } + }; } else { formatter = new JdbcResponseFormatter(PRETTY); } - return new ResponseListener() { + return new ResponseListener<>() { @Override public void onResponse(QueryResponse response) { sendResponse(channel, OK, diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index 9a136a3bec9..7086f5c8fab 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -49,7 +49,8 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context, result.add(plan.next()); } - QueryResponse response = new QueryResponse(physicalPlan.schema(), result); + String rawResponse = physicalPlan.getRawResponse(); + QueryResponse response = new QueryResponse(physicalPlan.schema(), result, rawResponse); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java index aadd73efdde..0c7e707d22b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java @@ -12,6 +12,7 @@ import java.util.Map; import java.util.stream.Collectors; import lombok.EqualsAndHashCode; +import lombok.Getter; import lombok.ToString; import org.opensearch.action.search.SearchResponse; import org.opensearch.search.SearchHits; @@ -43,6 +44,9 @@ public class OpenSearchResponse implements Iterable { */ @EqualsAndHashCode.Exclude private final OpenSearchExprValueFactory exprValueFactory; + @EqualsAndHashCode.Exclude + @Getter + private String rawResponse; /** * Constructor of ElasticsearchResponse. @@ -52,6 +56,7 @@ public OpenSearchResponse(SearchResponse searchResponse, this.hits = searchResponse.getHits(); this.aggregations = searchResponse.getAggregations(); this.exprValueFactory = exprValueFactory; + this.rawResponse = searchResponse.toString(); } /** diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java index e9746e1fae1..591ddff6a88 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java @@ -52,6 +52,9 @@ public class OpenSearchIndexScan extends TableScanOperator { /** Search response for current batch. */ private Iterator iterator; + @Getter + private String rawResponse; + /** * Constructor. */ @@ -80,7 +83,7 @@ public void open() { request = requestBuilder.build(); iterator = Collections.emptyIterator(); queryCount = 0; - fetchNextBatch(); + rawResponse = fetchNextBatch().getRawResponse(); } @Override @@ -99,11 +102,12 @@ public ExprValue next() { return iterator.next(); } - private void fetchNextBatch() { + protected OpenSearchResponse fetchNextBatch() { OpenSearchResponse response = client.search(request); if (!response.isEmpty()) { iterator = response.iterator(); } + return response; } @Override diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index 4a0c6e24f1c..f42778c6d8f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -25,7 +25,6 @@ import java.util.Arrays; import java.util.Iterator; import java.util.List; -import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; import lombok.RequiredArgsConstructor; @@ -45,6 +44,7 @@ import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; @@ -61,6 +61,8 @@ class OpenSearchExecutionEngineTest { @Mock private Split split; + @Mock private ProjectOperator fakeProjectOperator; + @BeforeEach void setUp() { doAnswer( @@ -207,6 +209,29 @@ public void onFailure(Exception e) { assertTrue(plan.hasClosed); } + @Test + void executeRawResponseSuccessfully() { + PhysicalPlan plan = fakeProjectOperator; + when(protector.protect(plan)).thenReturn(plan); + when(plan.getRawResponse()).thenReturn("searchResponse"); + + AtomicReference result = new AtomicReference<>(); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + executor.execute(plan, new ResponseListener() { + @Override + public void onResponse(QueryResponse response) { + result.set(response); + } + + @Override + public void onFailure(Exception e) { + fail(e); + } + }); + + assertEquals(result.get().getRawResponse(), "searchResponse"); + } + @RequiredArgsConstructor private static class FakePhysicalPlan extends TableScanOperator { private final Iterator it; diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java index 4291c09df07..4d1a84f47b2 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java @@ -16,6 +16,7 @@ @RequiredArgsConstructor public enum Format { JDBC("jdbc"), + JSON("json"), CSV("csv"), RAW("raw"), VIZ("viz"); diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index 082a3e95816..93ab0b71268 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -7,8 +7,12 @@ package org.opensearch.sql.sql; import java.util.Optional; +import java.util.stream.Collectors; import lombok.RequiredArgsConstructor; import org.antlr.v4.runtime.tree.ParseTree; +import org.opensearch.sql.analysis.JsonSupportVisitor; +import org.opensearch.sql.analysis.JsonSupportVisitorContext; +import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; @@ -75,6 +79,24 @@ private AbstractPlan plan( .isExplain(request.isExplainRequest()) .build())); + // There is no full support for JSON format yet for in memory operations, aliases, literals, + // and casts. Aggregation has differences with legacy results. + if (request.format().getFormatName().equals("json") && statement instanceof Query) { + if (parser.parseHints(request.getQuery()).getChildCount() > 1) { + throw new UnsupportedOperationException("Hints are not yet supported in the new engine."); + } + + // Go through the tree and throw exceptions when unsupported + JsonSupportVisitorContext jsonSupportVisitorContext = new JsonSupportVisitorContext(); + if (!((Query) statement).getPlan().accept(new JsonSupportVisitor(), + jsonSupportVisitorContext)) { + throw new UnsupportedOperationException( + "The following features are not supported with JSON format: " + .concat(jsonSupportVisitorContext.getUnsupportedNodes().stream() + .collect(Collectors.joining(", ")))); + } + } + return queryExecutionFactory.create(statement, queryListener, explainListener); } } diff --git a/sql/src/main/java/org/opensearch/sql/sql/antlr/SQLSyntaxParser.java b/sql/src/main/java/org/opensearch/sql/sql/antlr/SQLSyntaxParser.java index ee1e991bd4a..e8cd9f6909b 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/antlr/SQLSyntaxParser.java +++ b/sql/src/main/java/org/opensearch/sql/sql/antlr/SQLSyntaxParser.java @@ -32,4 +32,15 @@ public ParseTree parse(String query) { return parser.root(); } + /** + * Checks if the query contains hints as it is not yet support in V2. + * @param query a SQL query + * @return boolean value if query contains hints + */ + public ParseTree parseHints(String query) { + OpenSearchSQLLexer lexer = new OpenSearchSQLLexer(new CaseInsensitiveCharStream(query)); + OpenSearchSQLParser hintsParser = new OpenSearchSQLParser( + new CommonTokenStream(lexer, OpenSearchSQLLexer.SQLCOMMENT)); + return hintsParser.root(); + } } diff --git a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java index 508f80cee41..7ef3901aa11 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java +++ b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java @@ -8,6 +8,7 @@ import com.google.common.base.Strings; import com.google.common.collect.ImmutableSet; +import java.util.Arrays; import java.util.Collections; import java.util.Locale; import java.util.Map; @@ -33,6 +34,7 @@ public class SQLQueryRequest { "query", "fetch_size", "parameters"); private static final String QUERY_PARAMS_FORMAT = "format"; private static final String QUERY_PARAMS_SANITIZE = "sanitize"; + private static final String[] SUPPORTED_FORMATS = {"jdbc", "csv", "raw", "json"}; /** * JSON payload in REST request. @@ -121,8 +123,8 @@ private boolean isFetchSizeZeroIfPresent() { } private boolean isSupportedFormat() { - return Strings.isNullOrEmpty(format) || "jdbc".equalsIgnoreCase(format) - || "csv".equalsIgnoreCase(format) || "raw".equalsIgnoreCase(format); + return Strings.isNullOrEmpty(format) + || Arrays.stream(SUPPORTED_FORMATS).anyMatch(format::equalsIgnoreCase); } private String getFormat(Map params) { diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java index a351c306099..cdb463f1415 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -60,16 +60,9 @@ public void cleanup() throws InterruptedException { queryManager.awaitTermination(1, TimeUnit.SECONDS); } - @Test - public void canExecuteSqlQuery() { - doAnswer(invocation -> { - ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); - return null; - }).when(queryService).execute(any(), any()); - + private void execute(String query, String format) { sqlService.execute( - new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "jdbc"), + new SQLQueryRequest(new JSONObject(), query, QUERY, format), new ResponseListener() { @Override public void onResponse(QueryResponse response) { @@ -83,19 +76,72 @@ public void onFailure(Exception e) { }); } + private void executeWithUnsupportedException(String query, String format) { + sqlService.execute( + new SQLQueryRequest(new JSONObject(), query, QUERY, format), + new ResponseListener() { + @Override + public void onResponse(QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + assert (e instanceof UnsupportedOperationException); + } + }); + } + + @Test + public void canExecuteSqlQuery() { + doAnswer(invocation -> { + ResponseListener listener = invocation.getArgument(1); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), "")); + return null; + }).when(queryService).execute(any(), any()); + + execute("SELECT 123", "jdbc"); + } + @Test public void canExecuteCsvFormatRequest() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), "")); return null; }).when(queryService).execute(any(), any()); - sqlService.execute( - new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "csv"), - new ResponseListener() { + execute("SELECT 123", "csv"); + } + + @Test + public void canExecuteJsonFormatRequest() { + doAnswer(invocation -> { + ResponseListener listener = invocation.getArgument(1); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), "")); + return null; + }).when(queryService).execute(any(), any()); + + execute("SELECT FIELD FROM TABLE", "json"); + } + + @Test + public void canThrowUnsupportedExceptionWithUnsupportedJsonQuery() { + executeWithUnsupportedException("SELECT CAST(FIELD AS DOUBLE)", "json"); + } + + @Test + public void canThrowUnsupportedExceptionForHintsQuery() { + executeWithUnsupportedException("SELECT /*! HINTS */ 123", "json"); + } + + @Test + public void canExplainHintsQueryWithJsonFormat() { + sqlService.explain( + new SQLQueryRequest(new JSONObject(), "SELECT /*! HINTS */ 123", EXPLAIN, "json"), + new ResponseListener() { @Override - public void onResponse(QueryResponse response) { + public void onResponse(ExplainResponse response) { assertNotNull(response); } 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 index 415a77e17c9..62398b95209 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java @@ -291,7 +291,7 @@ public void can_parse_dayofmonth_functions() { assertNotNull(parser.parse("SELECT dayofmonth('2022-11-18')")); assertNotNull(parser.parse("SELECT day_of_month('2022-11-18')")); } - + @Test public void can_parse_day_of_week_functions() { assertNotNull(parser.parse("SELECT dayofweek('2022-11-18')")); @@ -720,6 +720,12 @@ public void canParseMultiMatchAlternateSyntax() { assertNotNull(parser.parse("SELECT * FROM test WHERE Field = multimatch(\"query\")")); } + @Test + public void canParseHints() { + assertNotNull(parser.parseHints("SELECT /*! HINTS */ field FROM test")); + assertNotNull(parser.parseHints("SELECT field FROM test")); + } + private static Stream matchPhraseQueryComplexQueries() { return Stream.of( "SELECT * FROM t WHERE matchphrasequery(c, 3)", diff --git a/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java b/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java index 52a1f534e9e..c2064332805 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java @@ -142,6 +142,15 @@ public void shouldSupportRawFormat() { assertTrue(csvRequest.isSupported()); } + @Test + public void shouldSupportJsonFormat() { + SQLQueryRequest jsonRequest = + SQLQueryRequestBuilder.request("SELECT 1") + .format("json") + .build(); + assertTrue(jsonRequest.isSupported()); + } + /** * SQL query request build helper to improve test data setup readability. */