From 7e2dc80e2a7599bcb3faf52e52566af09d3faabe Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 27 Jan 2023 19:08:05 -0800 Subject: [PATCH 01/46] Fixing integration tests broken during POC Signed-off-by: MaxKsyunz --- .../common/antlr/SyntaxCheckException.java | 3 +- .../org/opensearch/sql/analysis/Analyzer.java | 8 ++ .../sql/ast/AbstractNodeVisitor.java | 5 + .../opensearch/sql/ast/statement/Query.java | 1 + .../org/opensearch/sql/ast/tree/Paginate.java | 45 +++++++++ .../sql/executor/CanPaginateVisitor.java | 67 +++++++++++++ .../sql/executor/ExecutionEngine.java | 3 + .../sql/executor/PaginatedPlanCache.java | 96 +++++++++++++++++++ .../org/opensearch/sql/executor/QueryId.java | 1 + .../opensearch/sql/executor/QueryService.java | 12 +++ .../execution/ContinuePaginatedPlan.java | 52 ++++++++++ .../sql/executor/execution/PaginatedPlan.java | 45 +++++++++ .../executor/execution/QueryPlanFactory.java | 28 +++++- .../UnsupportCursorRequestException.java | 9 ++ .../sql/opensearch/executor/Cursor.java | 30 ++++++ .../sql/planner/DefaultImplementor.java | 8 ++ .../sql/planner/PaginateOperator.java | 75 +++++++++++++++ .../sql/planner/logical/LogicalPaginate.java | 28 ++++++ .../logical/LogicalPlanNodeVisitor.java | 4 + .../sql/planner/physical/PhysicalPlan.java | 5 +- .../physical/PhysicalPlanNodeVisitor.java | 5 + .../opensearch/sql/storage/StorageEngine.java | 3 +- .../sql/storage/read/TableScanBuilder.java | 3 + .../sql/executor/PaginatedPlanCacheTest.java | 45 +++++++++ .../sql/executor/QueryServiceTest.java | 6 +- .../execution/QueryPlanFactoryTest.java | 13 ++- .../MicroBatchStreamingExecutionTest.java | 4 +- .../planner/physical/ProjectOperatorTest.java | 33 ++++++- .../sql/executor/DefaultExecutionEngine.java | 4 +- integ-test/build.gradle | 15 ++- .../org/opensearch/sql/ppl/StandaloneIT.java | 16 +++- .../resources/datasource/datasources.json | 4 +- .../sql/legacy/plugin/RestSQLQueryAction.java | 2 +- .../sql/legacy/plugin/RestSqlAction.java | 21 +++- .../executor/OpenSearchExecutionEngine.java | 10 +- .../OpenSearchExecutionProtector.java | 7 ++ .../opensearch/request/OpenSearchRequest.java | 3 +- .../request/OpenSearchRequestBuilder.java | 18 ++-- .../opensearch/storage/OpenSearchIndex.java | 6 +- .../storage/OpenSearchIndexScan.java | 25 +---- .../scan/OpenSearchIndexScanBuilder.java | 5 + .../scan/OpenSearchIndexScanQueryBuilder.java | 4 + .../sql/opensearch/executor/CursorTest.java | 23 +++++ .../OpenSearchExecutionEngineTest.java | 21 ++-- .../OpenSearchExecutionProtectorTest.java | 17 ++-- .../storage/OpenSearchIndexScanTest.java | 24 +++-- .../storage/OpenSearchIndexTest.java | 17 ++-- plugin/build.gradle | 1 + .../plugin/config/OpenSearchPluginModule.java | 15 ++- .../transport/TransportPPLQueryAction.java | 3 +- .../sql/ppl/parser/AstStatementBuilder.java | 3 +- .../opensearch/sql/ppl/PPLServiceTest.java | 14 ++- .../ppl/parser/AstStatementBuilderTest.java | 4 +- .../sql/protocol/response/QueryResult.java | 8 ++ .../format/JdbcResponseFormatter.java | 6 ++ .../protocol/response/QueryResultTest.java | 12 ++- .../org/opensearch/sql/sql/SQLService.java | 26 +++-- .../sql/sql/domain/SQLQueryRequest.java | 36 ++++--- .../sql/sql/parser/AstStatementBuilder.java | 3 +- .../opensearch/sql/sql/SQLServiceTest.java | 13 ++- .../sql/sql/domain/SQLQueryRequestTest.java | 6 +- 61 files changed, 898 insertions(+), 131 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/ast/tree/Paginate.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java create mode 100644 core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java create mode 100644 core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java create mode 100644 core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java create mode 100644 core/src/main/java/org/opensearch/sql/planner/logical/LogicalPaginate.java create mode 100644 core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java diff --git a/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java b/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java index 806cb7208bb..684e66029e5 100644 --- a/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java +++ b/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java @@ -6,7 +6,8 @@ package org.opensearch.sql.common.antlr; -public class SyntaxCheckException extends RuntimeException { +public class +SyntaxCheckException extends RuntimeException { public SyntaxCheckException(String message) { super(message); } diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index 228b54ba0ca..b96f2472650 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -50,6 +50,7 @@ import org.opensearch.sql.ast.tree.Kmeans; import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.ML; +import org.opensearch.sql.ast.tree.Paginate; import org.opensearch.sql.ast.tree.Parse; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.RareTopN; @@ -85,6 +86,7 @@ import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalML; import org.opensearch.sql.planner.logical.LogicalMLCommons; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalRareTopN; @@ -529,6 +531,12 @@ public LogicalPlan visitML(ML node, AnalysisContext context) { return new LogicalML(child, node.getArguments()); } + @Override + public LogicalPlan visitPaginate(Paginate paginate, AnalysisContext context) { + LogicalPlan child = paginate.getChild().get(0).accept(this, context); + return new LogicalPaginate(paginate.getPageSize(), List.of(child)); + } + /** * The first argument is always "asc", others are optional. * Given nullFirst argument, use its value. Otherwise just use DEFAULT_ASC/DESC. 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..adcde61d426 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -47,6 +47,7 @@ import org.opensearch.sql.ast.tree.Kmeans; import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.ML; +import org.opensearch.sql.ast.tree.Paginate; import org.opensearch.sql.ast.tree.Parse; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.RareTopN; @@ -289,4 +290,8 @@ public T visitQuery(Query node, C context) { public T visitExplain(Explain node, C context) { return visitStatement(node, context); } + + public T visitPaginate(Paginate paginate, C context) { + return visitChildren(paginate, context); + } } diff --git a/core/src/main/java/org/opensearch/sql/ast/statement/Query.java b/core/src/main/java/org/opensearch/sql/ast/statement/Query.java index 17682cd47b9..82efdde4ddc 100644 --- a/core/src/main/java/org/opensearch/sql/ast/statement/Query.java +++ b/core/src/main/java/org/opensearch/sql/ast/statement/Query.java @@ -27,6 +27,7 @@ public class Query extends Statement { protected final UnresolvedPlan plan; + protected final int fetchSize; @Override public R accept(AbstractNodeVisitor visitor, C context) { diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Paginate.java b/core/src/main/java/org/opensearch/sql/ast/tree/Paginate.java new file mode 100644 index 00000000000..05f2c54db43 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Paginate.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.tree; + +import java.util.List; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.ToString; +import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.Node; + +@RequiredArgsConstructor +@EqualsAndHashCode(callSuper = false) +@ToString +public class Paginate extends UnresolvedPlan { + @Getter + private final int pageSize; + private UnresolvedPlan child; + + public Paginate(int pageSize, UnresolvedPlan child) { + this.pageSize = pageSize; + this.child = child; + } + + @Override + public List getChild() { + return List.of(child); + } + + @Override + public T accept(AbstractNodeVisitor nodeVisitor, C context) { + return nodeVisitor.visitPaginate(this, context); + } + + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + assert this.child == null; + this.child = child; + return this; + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java new file mode 100644 index 00000000000..a28d04c3bc4 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.Node; +import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.tree.Project; +import org.opensearch.sql.ast.tree.Relation; +import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; +import org.opensearch.sql.planner.physical.ProjectOperator; +import org.opensearch.sql.storage.TableScanOperator; + +/** + * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache. + * If plan.accept(new CanpaginateVisitor(...)) returns true, + * then PaginatedPlanCache.convertToCursor will succeed. + * Otherwise, it will fail. + * Currently, the conditions are: + * - only projection of a relation is supported. + * - projection only has * (a.k.a. allFields). + * - Relation only scans one table + * - The table is an open search index. + * See PaginatedPlanCache.canConvertToCursor for usage. + */ +public class CanPaginateVisitor extends AbstractNodeVisitor { + + @Override + public Boolean visitRelation(Relation node, Object context) { + if (!node.getChild().isEmpty()) { + // Relation instance should never have a child, but check just in case. + return Boolean.FALSE; + } + + // TODO use storageEngine from the calling PaginatedPlanCache to determine if + // node.getTableQualifiedName is provided by the storage engine. Return false if it's + // not the case. + return Boolean.TRUE; + } + + @Override + public Boolean visit(Node node, Object context) { + return Boolean.FALSE; + } + + @Override + public Boolean visitProject(Project node, Object context) { + var projections = node.getProjectList(); + if (projections.size() != 1) { + return Boolean.FALSE; + } + + if (!(projections.get(0) instanceof AllFields)) { + return Boolean.FALSE; + } + + var children = node.getChild(); + if (children.size() != 1) { + return Boolean.FALSE; + } + + return children.get(0).accept(this, context); + } +} 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..1c78df704cc 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java @@ -14,6 +14,7 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.physical.PhysicalPlan; /** @@ -53,6 +54,8 @@ void execute(PhysicalPlan plan, ExecutionContext context, class QueryResponse { private final Schema schema; private final List results; + + private final Cursor cursor; } @Data diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java new file mode 100644 index 00000000000..d6524649ac0 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -0,0 +1,96 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import lombok.Data; +import lombok.RequiredArgsConstructor; +import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.expression.NamedExpression; +import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.opensearch.executor.Cursor; +import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; +import org.opensearch.sql.planner.physical.ProjectOperator; +import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.read.TableScanBuilder; + +@RequiredArgsConstructor +public class PaginatedPlanCache { + private final StorageEngine storageEngine; + public static final PaginatedPlanCache None = new PaginatedPlanCache(null); + + public boolean canConvertToCursor(UnresolvedPlan plan) { + return plan.accept(new CanPaginateVisitor(), null); + } + + @RequiredArgsConstructor + @Data + static class SeriazationContext { + private final PaginatedPlanCache cache; + } + + public static class SerializationVisitor + extends PhysicalPlanNodeVisitor { + private static final byte[] NO_CURSOR = new byte[] {}; + + @Override + public byte[] visitPaginate(PaginateOperator node, SeriazationContext context) { + // Save cursor to read the next page. + // Could process node.getChild() here with another visitor -- one that saves the + // parameters for other physical operators -- ProjectOperator, etc. + return String.format("You got it!%d", node.getPageIndex() + 1).getBytes(); + } + + // Cursor is returned only if physical plan node is PaginateOerator. + @Override + protected byte[] visitNode(PhysicalPlan node, SeriazationContext context) { + return NO_CURSOR; + } + } + + /** + * Converts a physical plan tree to a cursor. May cache plan related data somewhere. + */ + public Cursor convertToCursor(PhysicalPlan plan) { + var serializer = new SerializationVisitor(); + var raw = plan.accept(serializer, new SeriazationContext(this)); + return new Cursor(raw); + } + + /** + * Convers a cursor to a physical plann tree. + */ + public PhysicalPlan convertToPlan(String cursor) { + // TODO HACKY_HACK -- create a plan + if (cursor.startsWith("You got it!")) { + int pageIndex = Integer.parseInt(cursor.substring("You got it!".length())); + + Table table = storageEngine.getTable(null, "phrases"); + TableScanBuilder scanBuilder = table.createScanBuilder(); + scanBuilder.pushDownOffset(5 * pageIndex); + PhysicalPlan scan = scanBuilder.build(); + var fields = table.getFieldTypes(); + List references = + Stream.of("phrase", "test field", "insert_time2") + .map(c -> + new NamedExpression(c, new ReferenceExpression(c, List.of(c), fields.get(c)))) + .collect(Collectors.toList()); + + return new PaginateOperator(new ProjectOperator(scan, references, List.of()), 5, pageIndex); + + } else { + throw new RuntimeException("Unsupported cursor"); + } + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryId.java b/core/src/main/java/org/opensearch/sql/executor/QueryId.java index 933cb5d82dc..43d6aed85eb 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryId.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryId.java @@ -16,6 +16,7 @@ * Query id of {@link AbstractPlan}. */ public class QueryId { + public static final QueryId None = new QueryId(""); /** * Query id. */ diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 94e70819204..f2825a7149c 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -70,6 +70,18 @@ public void executePlan(LogicalPlan plan, } } + /** + * Execute a physical plan without analyzing or planning anything. + */ + public void executePlan(PhysicalPlan plan, + ResponseListener listener) { + try { + executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); + } catch (Exception e) { + listener.onFailure(e); + } + } + /** * Explain the query in {@link UnresolvedPlan} using {@link ResponseListener} to * get and format explain response. diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java new file mode 100644 index 00000000000..15d4d973513 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor.execution; + +import org.apache.commons.lang3.NotImplementedException; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.executor.QueryId; +import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.planner.physical.PhysicalPlan; + +public class ContinuePaginatedPlan extends AbstractPlan { + + public static final ContinuePaginatedPlan None + = new ContinuePaginatedPlan(QueryId.None, "", null, + null, null); + private final String cursor; + private final QueryService queryService; + private final PaginatedPlanCache paginatedPlanCache; + + private final ResponseListener queryResponseListener; + + + /** + * Create an abstract plan that can continue paginating a given cursor. + */ + public ContinuePaginatedPlan(QueryId queryId, String cursor, QueryService queryService, + PaginatedPlanCache ppc, + ResponseListener + queryResponseListener) { + super(queryId); + this.cursor = cursor; + this.paginatedPlanCache = ppc; + this.queryService = queryService; + this.queryResponseListener = queryResponseListener; + } + + @Override + public void execute() { + PhysicalPlan plan = paginatedPlanCache.convertToPlan(cursor); + queryService.executePlan(plan, queryResponseListener); + } + + @Override + public void explain(ResponseListener listener) { + throw new NotImplementedException("Explain of query continuation is not supported"); + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java new file mode 100644 index 00000000000..a0d4a1eaea7 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor.execution; + +import org.opensearch.sql.ast.tree.Paginate; +import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.QueryId; +import org.opensearch.sql.executor.QueryService; + +public class PaginatedPlan extends AbstractPlan { + final UnresolvedPlan plan; + final int fetchSize; + final QueryService queryService; + final ResponseListener + queryResponseResponseListener; + + /** + * Create an abstract plan that can start paging a query. + */ + public PaginatedPlan(QueryId queryId, UnresolvedPlan plan, int fetchSize, + QueryService queryService, + ResponseListener + queryResponseResponseListener) { + super(queryId); + this.plan = plan; + this.fetchSize = fetchSize; + this.queryService = queryService; + this.queryResponseResponseListener = queryResponseResponseListener; + } + + @Override + public void execute() { + queryService.execute(new Paginate(fetchSize, plan), queryResponseResponseListener); + } + + @Override + public void explain(ResponseListener listener) { + + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 851381cc7a4..25ccec08d52 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -12,6 +12,7 @@ import com.google.common.base.Preconditions; import java.util.Optional; import lombok.RequiredArgsConstructor; +import org.apache.commons.lang3.NotImplementedException; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.statement.Explain; @@ -19,8 +20,10 @@ import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.legacy.plugin.UnsupportCursorRequestException; /** * QueryExecution Factory. @@ -37,6 +40,7 @@ public class QueryPlanFactory * Query Service. */ private final QueryService queryService; + private final PaginatedPlanCache paginatedPlanCache; /** * NO_CONSUMER_RESPONSE_LISTENER should never been called. It is only used as constructor @@ -69,6 +73,16 @@ public AbstractPlan create( return statement.accept(this, Pair.of(queryListener, explainListener)); } + /** + * Creates a ContinuePaginatedPlan from a cursor. + */ + public AbstractPlan create(String cursor, ResponseListener + queryResponseListener) { + QueryId queryId = QueryId.queryId(); + return new ContinuePaginatedPlan(queryId, cursor, queryService, paginatedPlanCache, + queryResponseListener); + } + @Override public AbstractPlan visitQuery( Query node, @@ -79,7 +93,18 @@ public AbstractPlan visitQuery( Preconditions.checkArgument( context.getLeft().isPresent(), "[BUG] query listener must be not null"); - return new QueryPlan(QueryId.queryId(), node.getPlan(), queryService, context.getLeft().get()); + if (node.getFetchSize() > 0) { + if (paginatedPlanCache.canConvertToCursor(node.getPlan())) { + return new PaginatedPlan(QueryId.queryId(), node.getPlan(), node.getFetchSize(), + queryService, + context.getLeft().get()); + } else { + throw new UnsupportCursorRequestException(); + } + } else { + return new QueryPlan(QueryId.queryId(), node.getPlan(), queryService, + context.getLeft().get()); + } } @Override @@ -97,4 +122,5 @@ public AbstractPlan visitExplain( create(node.getStatement(), Optional.of(NO_CONSUMER_RESPONSE_LISTENER), Optional.empty()), context.getRight().get()); } + } diff --git a/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java b/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java new file mode 100644 index 00000000000..7fcfceeebaf --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java @@ -0,0 +1,9 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.legacy.plugin; + +public class UnsupportCursorRequestException extends RuntimeException { +} diff --git a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java new file mode 100644 index 00000000000..02d06c6aec6 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import lombok.EqualsAndHashCode; +import lombok.Getter; + +@EqualsAndHashCode +public class Cursor { + public static final Cursor None = new Cursor(); + + @Getter + private final byte[] raw; + + private Cursor() { + raw = new byte[] {}; + } + + public Cursor(byte[] raw) { + this.raw = raw; + } + + @Override + public String toString() { + return new String(raw); + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index 4a6d4d82222..8bb1770d713 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -11,6 +11,7 @@ import org.opensearch.sql.planner.logical.LogicalEval; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; import org.opensearch.sql.planner.logical.LogicalProject; @@ -125,6 +126,12 @@ public PhysicalPlan visitLimit(LogicalLimit node, C context) { return new LimitOperator(visitChild(node, context), node.getLimit(), node.getOffset()); } + + @Override + public PhysicalPlan visitPaginate(LogicalPaginate plan, C context) { + return new PaginateOperator(visitChild(plan, context), plan.getPageSize()); + } + @Override public PhysicalPlan visitTableScanBuilder(TableScanBuilder plan, C context) { return plan.build(); @@ -141,6 +148,7 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { + "implementing and optimizing logical plan with relation involved"); } + protected PhysicalPlan visitChild(LogicalPlan node, C context) { // Logical operators visited here must have a single child return node.getChild().get(0).accept(this, context); diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java new file mode 100644 index 00000000000..f1add36e79e --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java @@ -0,0 +1,75 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner; + +import java.util.List; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; +import org.opensearch.sql.planner.physical.ProjectOperator; + +@RequiredArgsConstructor +public class PaginateOperator extends PhysicalPlan { + @Getter + private final PhysicalPlan input; + + @Getter + private final int pageSize; + + /** + * Which page is this? + * May not be necessary in the end. Currently used to increment the "cursor counter" -- + * See usage. + */ + @Getter + private final int pageIndex; + + int numReturned = 0; + + /** + * Page given physical plan, with pageSize elements per page, starting with the first page. + */ + public PaginateOperator(PhysicalPlan input, int pageSize) { + this.pageSize = pageSize; + this.input = input; + this.pageIndex = 0; + } + + @Override + public R accept(PhysicalPlanNodeVisitor visitor, C context) { + return visitor.visitPaginate(this, context); + } + + @Override + public boolean hasNext() { + return numReturned < pageSize && input.hasNext(); + } + + @Override + public void open() { + super.open(); + numReturned = 0; + } + + @Override + public ExprValue next() { + numReturned += 1; + return input.next(); + } + + public List getChild() { + return List.of(input); + } + + @Override + public ExecutionEngine.Schema schema() { + assert input instanceof ProjectOperator; + return input.schema(); + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPaginate.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPaginate.java new file mode 100644 index 00000000000..ab6eb341062 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPaginate.java @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner.logical; + +import java.util.List; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.ToString; + +@Getter +@ToString +@EqualsAndHashCode(callSuper = false) +public class LogicalPaginate extends LogicalPlan { + private final int pageSize; + + public LogicalPaginate(int pageSize, List childPlans) { + super(childPlans); + this.pageSize = pageSize; + } + + @Override + public R accept(LogicalPlanNodeVisitor visitor, C context) { + return visitor.visitPaginate(this, context); + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java index 9a41072fe7b..28cf6bcd792 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java @@ -100,4 +100,8 @@ public R visitML(LogicalML plan, C context) { public R visitAD(LogicalAD plan, C context) { return visitNode(plan, context); } + + public R visitPaginate(LogicalPaginate plan, C context) { + return visitNode(plan, context); + } } 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..92e032f8d3a 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 @@ -6,6 +6,7 @@ package org.opensearch.sql.planner.physical; +import java.io.Serializable; import java.util.Iterator; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.executor.ExecutionEngine; @@ -17,7 +18,8 @@ */ public abstract class PhysicalPlan implements PlanNode, Iterator, - AutoCloseable { + AutoCloseable, + Serializable { /** * Accept the {@link PhysicalPlanNodeVisitor}. * @@ -45,4 +47,5 @@ public ExecutionEngine.Schema schema() { throw new IllegalStateException(String.format("[BUG] schema can been only applied to " + "ProjectOperator, instead of %s", toString())); } + } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java index d4bc4a1ea9f..e52e5979ed3 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java @@ -6,6 +6,7 @@ package org.opensearch.sql.planner.physical; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.write.TableWriteOperator; @@ -88,4 +89,8 @@ public R visitAD(PhysicalPlan node, C context) { public R visitML(PhysicalPlan node, C context) { return visitNode(node, context); } + + public R visitPaginate(PaginateOperator node, C context) { + return visitNode(node, context); + } } diff --git a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java index 246a50ea093..73a96a8b7c6 100644 --- a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java +++ b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java @@ -8,7 +8,9 @@ import java.util.Collection; import java.util.Collections; +import java.util.List; import org.opensearch.sql.DataSourceSchemaName; +import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.function.FunctionResolver; /** @@ -29,5 +31,4 @@ public interface StorageEngine { default Collection getFunctions() { return Collections.emptyList(); } - } diff --git a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java index c0fdf36e709..e05cfad94ee 100644 --- a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java +++ b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java @@ -108,4 +108,7 @@ public boolean pushDownHighlight(LogicalHighlight highlight) { public R accept(LogicalPlanNodeVisitor visitor, C context) { return visitor.visitTableScanBuilder(this, context); } + + public void pushDownOffset(int i) { + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java new file mode 100644 index 00000000000..21b99c6a36d --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.opensearch.sql.ast.dsl.AstDSL; +import org.opensearch.sql.storage.StorageEngine; + +class PaginatedPlanCacheTest { + + @Mock + StorageEngine storageEngine; + + PaginatedPlanCache planCache; + @BeforeEach + void setUp() { + planCache = new PaginatedPlanCache(storageEngine); + } + + @Test + void canConvertToCursor_relation() { + Assertions.assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); + } + + @Test + void canConvertToCursor_project_allFields_relation() { + var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.allFields()); + Assertions.assertTrue(planCache.canConvertToCursor(unresolvedPlan)); + } + + @Test + void canConvertToCursor_project_some_fields_relation() { + var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.field("rando")); + Assertions.assertFalse(planCache.canConvertToCursor(unresolvedPlan)); + } +} diff --git a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java b/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java index 4df38027f4f..a0be4f8f2ec 100644 --- a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java @@ -15,11 +15,9 @@ import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.when; import java.util.Collections; import java.util.Optional; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -27,6 +25,7 @@ import org.opensearch.sql.analysis.Analyzer; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.PlanContext; import org.opensearch.sql.planner.Planner; import org.opensearch.sql.planner.logical.LogicalPlan; @@ -134,7 +133,8 @@ Helper executeSuccess(Split split) { invocation -> { ResponseListener listener = invocation.getArgument(2); listener.onResponse( - new ExecutionEngine.QueryResponse(schema, Collections.emptyList())); + new ExecutionEngine.QueryResponse(schema, Collections.emptyList(), + Cursor.None)); return null; }) .when(executionEngine) diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index cc4bf070fbe..b1940022aaa 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -25,6 +25,7 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryService; @ExtendWith(MockitoExtension.class) @@ -45,16 +46,18 @@ class QueryPlanFactoryTest { @Mock private ExecutionEngine.QueryResponse queryResponse; + @Mock + private PaginatedPlanCache paginatedPlanCache; private QueryPlanFactory factory; @BeforeEach void init() { - factory = new QueryPlanFactory(queryService); + factory = new QueryPlanFactory(queryService, paginatedPlanCache); } @Test public void createFromQueryShouldSuccess() { - Statement query = new Query(plan); + Statement query = new Query(plan, 0); AbstractPlan queryExecution = factory.create(query, Optional.of(queryListener), Optional.empty()); assertTrue(queryExecution instanceof QueryPlan); @@ -62,7 +65,7 @@ public void createFromQueryShouldSuccess() { @Test public void createFromExplainShouldSuccess() { - Statement query = new Explain(new Query(plan)); + Statement query = new Explain(new Query(plan, 0)); AbstractPlan queryExecution = factory.create(query, Optional.empty(), Optional.of(explainListener)); assertTrue(queryExecution instanceof ExplainPlan); @@ -70,7 +73,7 @@ public void createFromExplainShouldSuccess() { @Test public void createFromQueryWithoutQueryListenerShouldThrowException() { - Statement query = new Query(plan); + Statement query = new Query(plan, 0); IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> factory.create(query, @@ -80,7 +83,7 @@ public void createFromQueryWithoutQueryListenerShouldThrowException() { @Test public void createFromExplainWithoutExplainListenerShouldThrowException() { - Statement query = new Explain(new Query(plan)); + Statement query = new Explain(new Query(plan, 0)); IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> factory.create(query, diff --git a/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java b/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java index 1a2b6e3f2a4..75a62385308 100644 --- a/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java @@ -26,6 +26,7 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.PlanContext; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.storage.split.Split; @@ -169,7 +170,8 @@ Helper executeSuccess(Long... offsets) { ResponseListener listener = invocation.getArgument(2); listener.onResponse( - new ExecutionEngine.QueryResponse(null, Collections.emptyList())); + new ExecutionEngine.QueryResponse(null, Collections.emptyList(), + Cursor.None)); PlanContext planContext = invocation.getArgument(1); assertTrue(planContext.getSplit().isPresent()); diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index 24be5eb2b8d..6afa9db3bd8 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -10,6 +10,7 @@ import static org.hamcrest.Matchers.allOf; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.hasItems; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.iterableWithSize; import static org.mockito.Mockito.when; import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_MISSING; @@ -20,7 +21,13 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; import java.util.List; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -34,7 +41,7 @@ @ExtendWith(MockitoExtension.class) class ProjectOperatorTest extends PhysicalPlanTestBase { - @Mock + @Mock(serializable = true) private PhysicalPlan inputPlan; @Test @@ -206,4 +213,28 @@ public void project_parse_missing_will_fallback() { ExprValueUtils.tupleValue(ImmutableMap.of("action", "GET", "response", "200")), ExprValueUtils.tupleValue(ImmutableMap.of("action", "POST"))))); } + + @Test + public void project_serialize() throws IOException, ClassNotFoundException { + PhysicalPlan plan = project(inputPlan, DSL.named("action", DSL.ref("action", STRING))); + + var os = new ByteArrayOutputStream(); + var objstream = new ObjectOutputStream(os); + objstream.writeObject(plan); + objstream.close(); + os.close(); + String result = os.toString(); + + + var is = new ByteArrayInputStream(os.toByteArray()); + var outstream = new ObjectInputStream(is); + var newObj = outstream.readObject(); + + assertThat(newObj, instanceOf(ProjectOperator.class)); + + var newOp = (ProjectOperator) newObj; + + Assertions.assertEquals(1, newOp.getProjectList().size()); + Assertions.assertEquals("action", newOp.getProjectList().get(0).getName()); + } } diff --git a/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java b/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java index e4f9a185a30..18058302711 100644 --- a/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java +++ b/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java @@ -9,6 +9,7 @@ import java.util.List; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.physical.PhysicalPlan; /** @@ -32,7 +33,8 @@ public void execute( while (plan.hasNext()) { result.add(plan.next()); } - QueryResponse response = new QueryResponse(new Schema(new ArrayList<>()), new ArrayList<>()); + QueryResponse response = new QueryResponse(new Schema(new ArrayList<>()), new ArrayList<>(), + Cursor.None); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); diff --git a/integ-test/build.gradle b/integ-test/build.gradle index 911ee0b2539..ad8a2128cd8 100644 --- a/integ-test/build.gradle +++ b/integ-test/build.gradle @@ -109,6 +109,11 @@ compileTestJava { testClusters.all { testDistribution = 'archive' + + // debug with command, ./gradlew opensearch-sql:run -DdebugJVM. --debug-jvm does not work with keystore. + if (System.getProperty("debugJVM") != null) { + jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005' + } } testClusters.integTest { @@ -166,10 +171,16 @@ integTest { // Tell the test JVM if the cluster JVM is running under a debugger so that tests can use longer timeouts for // requests. The 'doFirst' delays reading the debug setting on the cluster till execution time. - doFirst { systemProperty 'cluster.debug', getDebug() } + doFirst { + if (System.getProperty("debug-jvm") != null) { + setDebug(true); + } + systemProperty 'cluster.debug', getDebug() + } + if (System.getProperty("test.debug") != null) { - jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:5005' + jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:5006' } if (System.getProperty("tests.rest.bwcsuite") == null) { diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java index 71988a8e31c..28c4250f8b6 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java @@ -34,6 +34,7 @@ import org.opensearch.sql.datasource.DataSourceServiceImpl; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; @@ -194,8 +195,9 @@ public StorageEngine storageEngine(OpenSearchClient client) { } @Provides - public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector) { - return new OpenSearchExecutionEngine(client, protector); + public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector, + PaginatedPlanCache paginatedPlanCache) { + return new OpenSearchExecutionEngine(client, protector, paginatedPlanCache); } @Provides @@ -225,12 +227,18 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl } @Provides - public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine) { + public PaginatedPlanCache paginatedPlanCache(StorageEngine storageEngine) { + return new PaginatedPlanCache(storageEngine); + } + @Provides + public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, + PaginatedPlanCache paginatedPlanCache) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); Planner planner = new Planner(LogicalPlanOptimizer.create()); - return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner)); + return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner), + paginatedPlanCache); } } } diff --git a/integ-test/src/test/resources/datasource/datasources.json b/integ-test/src/test/resources/datasource/datasources.json index 5f195747ae0..fbf18cfc462 100644 --- a/integ-test/src/test/resources/datasource/datasources.json +++ b/integ-test/src/test/resources/datasource/datasources.json @@ -3,7 +3,7 @@ "name" : "my_prometheus", "connector": "prometheus", "properties" : { - "prometheus.uri" : "http://localhost:9090" + "prometheus.uri" : "http://localhost:9091" } } -] \ No newline at end of file +] 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..9594d5a3b51 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 UnsupportCursorRequestException) { fallBackHandler.accept(channel, e); } else { next.onFailure(e); diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 88ed42010bd..6fdd855fd0e 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -42,6 +42,7 @@ import org.opensearch.sql.legacy.antlr.SqlAnalysisConfig; import org.opensearch.sql.legacy.antlr.SqlAnalysisException; import org.opensearch.sql.legacy.antlr.semantic.types.Type; +import org.opensearch.sql.legacy.cursor.CursorType; import org.opensearch.sql.legacy.domain.ColumnTypeProvider; import org.opensearch.sql.legacy.domain.QueryActionRequest; import org.opensearch.sql.legacy.esdomain.LocalClusterState; @@ -132,7 +133,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli } final SqlRequest sqlRequest = SqlRequestFactory.getSqlRequest(request); - if (sqlRequest.cursor() != null) { + if (isLegacyCursor(sqlRequest)) { if (isExplainRequest(request)) { throw new IllegalArgumentException("Invalid request. Cannot explain cursor"); } else { @@ -141,14 +142,15 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli } } - LOG.info("[{}] Incoming request {}: {}", QueryContext.getRequestId(), request.uri(), - QueryDataAnonymizer.anonymizeData(sqlRequest.getSql())); + +// LOG.info("[{}] Incoming request {}: {}", QueryContext.getRequestId(), request.uri(), +// QueryDataAnonymizer.anonymizeData(sqlRequest.getSql())); Format format = SqlRequestParam.getFormat(request.params()); // Route request to new query engine if it's supported already SQLQueryRequest newSqlRequest = new SQLQueryRequest(sqlRequest.getJsonContent(), - sqlRequest.getSql(), request.path(), request.params()); + sqlRequest.getSql(), request.path(), request.params(), sqlRequest.cursor()); return newSqlQueryHandler.prepareRequest(newSqlRequest, (restChannel, exception) -> { try{ @@ -175,6 +177,17 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli } } + + /** + * @param sqlRequest client request + * @return true if this cursor was generated by the legacy engine, false otherwise. + */ + private static boolean isLegacyCursor(SqlRequest sqlRequest) { + String cursor = sqlRequest.cursor(); + return cursor != null + && CursorType.getById(cursor.substring(0, 1)) != CursorType.NULL; + } + @Override protected Set responseParams() { Set responseParams = new HashSet<>(super.responseParams()); 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..1f4fb578b12 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 @@ -15,6 +15,7 @@ import org.opensearch.sql.executor.ExecutionContext; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.Explain; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -27,6 +28,7 @@ public class OpenSearchExecutionEngine implements ExecutionEngine { private final OpenSearchClient client; private final ExecutionProtector executionProtector; + private final PaginatedPlanCache paginatedPlanCache; @Override public void execute(PhysicalPlan physicalPlan, ResponseListener listener) { @@ -49,7 +51,13 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context, result.add(plan.next()); } - QueryResponse response = new QueryResponse(physicalPlan.schema(), result); + + // + // getContinuation expects hasNext to return false before it is called. + + Cursor qc = paginatedPlanCache.convertToCursor(plan); + + QueryResponse response = new QueryResponse(physicalPlan.schema(), result, qc); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java index f06ecb85768..572b5cbadec 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java @@ -11,6 +11,7 @@ import org.opensearch.sql.opensearch.planner.physical.ADOperator; import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator; import org.opensearch.sql.opensearch.planner.physical.MLOperator; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.AggregationOperator; import org.opensearch.sql.planner.physical.DedupeOperator; import org.opensearch.sql.planner.physical.EvalOperator; @@ -63,6 +64,12 @@ public PhysicalPlan visitRename(RenameOperator node, Object context) { return new RenameOperator(visitInput(node.getInput(), context), node.getMapping()); } + @Override + public PhysicalPlan visitPaginate(PaginateOperator node, Object context) { + return new PaginateOperator(visitInput(node.getInput(), context), node.getPageSize(), + node.getPageIndex()); + } + /** * Decorate with {@link ResourceMonitorPlan}. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index ce990780c1c..0de204c285b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -6,6 +6,7 @@ package org.opensearch.sql.opensearch.request; +import java.io.Serializable; import java.util.function.Consumer; import java.util.function.Function; import lombok.EqualsAndHashCode; @@ -19,7 +20,7 @@ /** * OpenSearch search request. */ -public interface OpenSearchRequest { +public interface OpenSearchRequest extends Serializable { /** * Apply the search action or scroll action on request based on context. * diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 439a970a4f4..ee97dc3ae8b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -9,7 +9,7 @@ import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; import static org.opensearch.search.sort.SortOrder.ASC; -import com.google.common.collect.Lists; +import java.io.Serializable; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -26,7 +26,6 @@ import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder; -import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.SortBuilder; import org.opensearch.search.sort.SortBuilders; import org.opensearch.sql.ast.expression.Literal; @@ -44,7 +43,7 @@ @EqualsAndHashCode @Getter @ToString -public class OpenSearchRequestBuilder { +public class OpenSearchRequestBuilder implements Serializable { /** * Default query timeout in minutes. @@ -74,15 +73,16 @@ public class OpenSearchRequestBuilder { private final OpenSearchExprValueFactory exprValueFactory; /** - * Query size of the request. + * Query size of the request -- how many rows will be returned. */ - private Integer querySize; + private int querySize; public OpenSearchRequestBuilder(String indexName, Integer maxResultWindow, Settings settings, OpenSearchExprValueFactory exprValueFactory) { - this(new OpenSearchRequest.IndexName(indexName), maxResultWindow, settings, exprValueFactory); + this(new OpenSearchRequest.IndexName(indexName), maxResultWindow, settings, + exprValueFactory); } /** @@ -111,11 +111,11 @@ public OpenSearchRequest build() { Integer from = sourceBuilder.from(); Integer size = sourceBuilder.size(); - if (from + size <= maxResultWindow) { - return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); - } else { + if (from + size > maxResultWindow) { sourceBuilder.size(maxResultWindow - from); return new OpenSearchScrollRequest(indexName, sourceBuilder, exprValueFactory); + } else { + return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index c694769b895..d7afac9060a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -19,6 +19,7 @@ import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator; import org.opensearch.sql.opensearch.planner.physical.MLOperator; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.request.system.OpenSearchDescribeIndexRequest; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder; import org.opensearch.sql.planner.DefaultImplementor; @@ -120,8 +121,9 @@ public LogicalPlan optimize(LogicalPlan plan) { @Override public TableScanBuilder createScanBuilder() { - OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, indexName, - getMaxResultWindow(), new OpenSearchExprValueFactory(getFieldTypes())); + var requestBuilder = new OpenSearchRequestBuilder(indexName, getMaxResultWindow(), + settings, new OpenSearchExprValueFactory(getFieldTypes())); + OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, requestBuilder); return new OpenSearchIndexScanBuilder(indexScan); } 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..9a1ddcba08e 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 @@ -28,7 +28,7 @@ public class OpenSearchIndexScan extends TableScanOperator { /** OpenSearch client. */ - private final OpenSearchClient client; + private final transient OpenSearchClient client; /** Search request builder. */ @EqualsAndHashCode.Include @@ -49,28 +49,13 @@ public class OpenSearchIndexScan extends TableScanOperator { /** Number of rows returned. */ private Integer queryCount; + /** Search response for current batch. */ - private Iterator iterator; - - /** - * Constructor. - */ - public OpenSearchIndexScan(OpenSearchClient client, Settings settings, - String indexName, Integer maxResultWindow, - OpenSearchExprValueFactory exprValueFactory) { - this(client, settings, - new OpenSearchRequest.IndexName(indexName),maxResultWindow, exprValueFactory); - } + private transient Iterator iterator; - /** - * Constructor. - */ - public OpenSearchIndexScan(OpenSearchClient client, Settings settings, - OpenSearchRequest.IndexName indexName, Integer maxResultWindow, - OpenSearchExprValueFactory exprValueFactory) { + public OpenSearchIndexScan(OpenSearchClient client, OpenSearchRequestBuilder builder) { this.client = client; - this.requestBuilder = new OpenSearchRequestBuilder( - indexName, maxResultWindow, settings,exprValueFactory); + this.requestBuilder = builder; } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index d7483cfcf06..72f21049490 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -92,6 +92,11 @@ public boolean pushDownProject(LogicalProject project) { return delegate.pushDownProject(project); } + @Override + public void pushDownOffset(int i) { + delegate.pushDownOffset(i); + } + @Override public boolean pushDownHighlight(LogicalHighlight highlight) { return delegate.pushDownHighlight(highlight); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index 7190d580002..370036cdeee 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -38,6 +38,10 @@ */ @VisibleForTesting class OpenSearchIndexScanQueryBuilder extends TableScanBuilder { + @Override + public void pushDownOffset(int i) { + indexScan.getRequestBuilder().getSourceBuilder().from(i); + } /** OpenSearch index scan to be optimized. */ @EqualsAndHashCode.Include diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java new file mode 100644 index 00000000000..6b2a1a7c574 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java @@ -0,0 +1,23 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class CursorTest { + + @Test + void EmptyArrayIsNone() { + Assertions.assertEquals(Cursor.None, new Cursor(new byte[]{})); + } + + @Test + void ToStringIsArrayValue() { + String cursorTxt = "This is a test"; + Assertions.assertEquals(cursorTxt, new Cursor(cursorTxt.getBytes()).toString()); + } +} 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..ff529a018f7 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; @@ -37,12 +36,14 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.ExecutionContext; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; @@ -82,7 +83,8 @@ void executeSuccessfully() { FakePhysicalPlan plan = new FakePhysicalPlan(expected.iterator()); when(protector.protect(plan)).thenReturn(plan); - OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); List actual = new ArrayList<>(); executor.execute( plan, @@ -110,7 +112,8 @@ void executeWithFailure() { when(plan.hasNext()).thenThrow(expected); when(protector.protect(plan)).thenReturn(plan); - OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); AtomicReference actual = new AtomicReference<>(); executor.execute( plan, @@ -131,11 +134,13 @@ public void onFailure(Exception e) { @Test void explainSuccessfully() { - OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); Settings settings = mock(Settings.class); when(settings.getSettingValue(QUERY_SIZE_LIMIT)).thenReturn(100); PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), - settings, "test", 10000, mock(OpenSearchExprValueFactory.class)); + new OpenSearchRequestBuilder("test", 10000, settings, + mock(OpenSearchExprValueFactory.class))); AtomicReference result = new AtomicReference<>(); executor.explain(plan, new ResponseListener() { @@ -155,7 +160,8 @@ public void onFailure(Exception e) { @Test void explainWithFailure() { - OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); PhysicalPlan plan = mock(PhysicalPlan.class); when(plan.accept(any(), any())).thenThrow(IllegalStateException.class); @@ -184,7 +190,8 @@ void callAddSplitAndOpenInOrder() { when(protector.protect(plan)).thenReturn(plan); when(executionContext.getSplit()).thenReturn(Optional.of(split)); - OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector); + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); List actual = new ArrayList<>(); executor.execute( plan, diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index 857ff601e14..c64d9e4ad96 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -57,6 +57,7 @@ import org.opensearch.sql.opensearch.planner.physical.ADOperator; import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator; import org.opensearch.sql.opensearch.planner.physical.MLOperator; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -124,9 +125,11 @@ public void testProtectIndexScan() { PhysicalPlanDSL.agg( filter( resourceMonitor( - new OpenSearchIndexScan( - client, settings, indexName, - maxResultWindow, exprValueFactory)), + new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder(indexName, + maxResultWindow, + settings, + exprValueFactory))), filterExpr), aggregators, groupByExprs), @@ -152,9 +155,11 @@ public void testProtectIndexScan() { PhysicalPlanDSL.rename( PhysicalPlanDSL.agg( filter( - new OpenSearchIndexScan( - client, settings, indexName, - maxResultWindow, exprValueFactory), + new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder(indexName, + maxResultWindow, + settings, + exprValueFactory)), filterExpr), aggregators, groupByExprs), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java index 9a606750a3b..4cb86015fa6 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java @@ -43,6 +43,7 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) @@ -66,7 +67,8 @@ void setup() { void queryEmptyResult() { mockResponse(); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "test", 3, exprValueFactory)) { + new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("test", 3, settings, + exprValueFactory))) { indexScan.open(); assertFalse(indexScan.hasNext()); } @@ -80,8 +82,10 @@ void queryAllResultsWithQuery() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); + OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder("employees", 10, settings, + exprValueFactory); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "employees", 10, exprValueFactory)) { + new OpenSearchIndexScan(client, builder)) { indexScan.open(); assertTrue(indexScan.hasNext()); @@ -105,7 +109,8 @@ void queryAllResultsWithScroll() { new ExprValue[]{employee(3, "Allen", "IT")}); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "employees", 2, exprValueFactory)) { + new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("employees", 2, settings, + exprValueFactory))) { indexScan.open(); assertTrue(indexScan.hasNext()); @@ -131,7 +136,8 @@ void querySomeResultsWithQuery() { employee(4, "Bob", "HR")}); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "employees", 10, exprValueFactory)) { + new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("employees", 10, settings, + exprValueFactory))) { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); @@ -156,7 +162,8 @@ void querySomeResultsWithScroll() { new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "employees", 2, exprValueFactory)) { + new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("employees", 2, settings, + exprValueFactory))) { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); @@ -225,7 +232,8 @@ void pushDownHighlightWithRepeatingFields() { new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, settings, "test", 2, exprValueFactory)) { + new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("test", 2, settings, + exprValueFactory))) { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); Map args = new HashMap<>(); @@ -251,7 +259,9 @@ public PushDownAssertion(OpenSearchClient client, OpenSearchExprValueFactory valueFactory, Settings settings) { this.client = client; - this.indexScan = new OpenSearchIndexScan(client, settings, "test", 10000, valueFactory); + this.indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("test", 10000, + settings, valueFactory)); this.response = mock(OpenSearchResponse.class); this.factory = valueFactory; when(response.isEmpty()).thenReturn(true); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 74c18f7c3d1..890951e4d3a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -54,6 +54,7 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.mapping.IndexMapping; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; @@ -159,9 +160,9 @@ void implementRelationOperatorOnly() { LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - assertEquals( - new OpenSearchIndexScan(client, settings, indexName, maxResultWindow, exprValueFactory), - index.implement(plan)); + OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, + settings, exprValueFactory); + assertEquals(new OpenSearchIndexScan(client, builder), index.implement(plan)); } @Test @@ -171,8 +172,10 @@ void implementRelationOperatorWithOptimization() { LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); + OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, + settings, exprValueFactory); assertEquals( - new OpenSearchIndexScan(client, settings, indexName, maxResultWindow, exprValueFactory), + new OpenSearchIndexScan(client, builder), index.implement(index.optimize(plan))); } @@ -220,8 +223,10 @@ void implementOtherLogicalOperators() { PhysicalPlanDSL.eval( PhysicalPlanDSL.remove( PhysicalPlanDSL.rename( - new OpenSearchIndexScan(client, settings, indexName, - maxResultWindow, exprValueFactory), + new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder( + indexName, maxResultWindow, + settings, exprValueFactory)), mappings), exclude), newEvalField), diff --git a/plugin/build.gradle b/plugin/build.gradle index f0bad12c2dc..c51bd8ce7d8 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -238,6 +238,7 @@ afterEvaluate { testClusters.integTest { plugin(project.tasks.bundlePlugin.archiveFile) + testDistribution = "ARCHIVE" // debug with command, ./gradlew opensearch-sql:run -DdebugJVM. --debug-jvm does not work with keystore. if (System.getProperty("debugJVM") != null) { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 5ab4bbaecd0..8fea281586c 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -15,6 +15,7 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; @@ -58,8 +59,9 @@ public StorageEngine storageEngine(OpenSearchClient client, Settings settings) { } @Provides - public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector) { - return new OpenSearchExecutionEngine(client, protector); + public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector, + PaginatedPlanCache paginatedPlanCache) { + return new OpenSearchExecutionEngine(client, protector, paginatedPlanCache); } @Provides @@ -72,6 +74,11 @@ public ExecutionProtector protector(ResourceMonitor resourceMonitor) { return new OpenSearchExecutionProtector(resourceMonitor); } + @Provides + public PaginatedPlanCache paginatedPlanCache(StorageEngine storageEngine) { + return new PaginatedPlanCache(storageEngine); + } + @Provides @Singleton public QueryManager queryManager(NodeClient nodeClient) { @@ -93,11 +100,11 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl */ @Provides public QueryPlanFactory queryPlanFactory( - DataSourceService dataSourceService, ExecutionEngine executionEngine) { + DataSourceService dataSourceService, ExecutionEngine executionEngine, PaginatedPlanCache p) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); Planner planner = new Planner(LogicalPlanOptimizer.create()); - return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner)); + return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner), p); } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index 6825b2ac923..d9dad9d5353 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -139,7 +139,8 @@ private ResponseListener createListener( @Override public void onResponse(ExecutionEngine.QueryResponse response) { String responseContent = - formatter.format(new QueryResult(response.getSchema(), response.getResults())); + formatter.format(new QueryResult(response.getSchema(), response.getResults(), + response.getCursor())); listener.onResponse(new TransportPPLQueryResponse(responseContent)); } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java index e4f40e9a115..3b7e5a78dde 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java @@ -33,7 +33,7 @@ public class AstStatementBuilder extends OpenSearchPPLParserBaseVisitor { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); return null; }).when(queryService).execute(any(), any()); @@ -87,7 +93,7 @@ public void onFailure(Exception e) { public void testExecuteCsvFormatShouldPass() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); return null; }).when(queryService).execute(any(), any()); @@ -161,7 +167,7 @@ public void onFailure(Exception e) { public void testPrometheusQuery() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); return null; }).when(queryService).execute(any(), any()); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java index 47600246920..cdb0e37ee50 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java @@ -39,7 +39,7 @@ public void buildQueryStatement() { "search source=t a=1", new Query( project( - filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFields.of()))); + filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFields.of()), 0)); } @Test @@ -50,7 +50,7 @@ public void buildExplainStatement() { new Query( project( filter(relation("t"), compare("=", field("a"), intLiteral(1))), - AllFields.of())))); + AllFields.of()), 0))); } private void assertEqual(String query, Statement expectedStatement) { diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java index 915a61f3611..3ea5846b87c 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java @@ -16,6 +16,7 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; +import org.opensearch.sql.opensearch.executor.Cursor; /** * Query response that encapsulates query results and isolate {@link ExprValue} @@ -32,6 +33,13 @@ public class QueryResult implements Iterable { */ private final Collection exprValues; + @Getter + private final Cursor cursor; + + + public QueryResult(ExecutionEngine.Schema schema, Collection exprValues) { + this(schema, exprValues, Cursor.None); + } /** * size of results. diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java index 943287cb62b..f52ee222467 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java @@ -15,6 +15,7 @@ import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.exception.QueryEngineException; import org.opensearch.sql.executor.ExecutionEngine.Schema; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.opensearch.response.error.ErrorMessage; import org.opensearch.sql.opensearch.response.error.ErrorMessageFactory; import org.opensearch.sql.protocol.response.QueryResult; @@ -42,6 +43,9 @@ protected Object buildJsonObject(QueryResult response) { json.total(response.size()) .size(response.size()) .status(200); + if (!response.getCursor().equals(Cursor.None)) { + json.cursor(response.getCursor().toString()); + } return json.build(); } @@ -95,6 +99,8 @@ public static class JdbcResponse { private final long total; private final long size; private final int status; + + private final String cursor; } @RequiredArgsConstructor diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java index 319965e2d0e..3db405339bc 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java @@ -19,6 +19,7 @@ import java.util.Collections; import org.junit.jupiter.api.Test; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.opensearch.executor.Cursor; class QueryResultTest { @@ -35,7 +36,7 @@ void size() { tupleValue(ImmutableMap.of("name", "John", "age", 20)), tupleValue(ImmutableMap.of("name", "Allen", "age", 30)), tupleValue(ImmutableMap.of("name", "Smith", "age", 40)) - )); + ), Cursor.None); assertEquals(3, response.size()); } @@ -45,7 +46,7 @@ void columnNameTypes() { schema, Collections.singletonList( tupleValue(ImmutableMap.of("name", "John", "age", 20)) - )); + ), Cursor.None); assertEquals( ImmutableMap.of("name", "string", "age", "integer"), @@ -59,7 +60,8 @@ void columnNameTypesWithAlias() { new ExecutionEngine.Schema.Column("name", "n", STRING))); QueryResult response = new QueryResult( schema, - Collections.singletonList(tupleValue(ImmutableMap.of("n", "John")))); + Collections.singletonList(tupleValue(ImmutableMap.of("n", "John"))), + Cursor.None); assertEquals( ImmutableMap.of("n", "string"), @@ -71,7 +73,7 @@ void columnNameTypesWithAlias() { void columnNameTypesFromEmptyExprValues() { QueryResult response = new QueryResult( schema, - Collections.emptyList()); + Collections.emptyList(), Cursor.None); assertEquals( ImmutableMap.of("name", "string", "age", "integer"), response.columnNameTypes() @@ -100,7 +102,7 @@ void iterate() { Arrays.asList( tupleValue(ImmutableMap.of("name", "John", "age", 20)), tupleValue(ImmutableMap.of("name", "Allen", "age", 30)) - )); + ), Cursor.None); int i = 0; for (Object[] objects : response) { 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..0912bf34826 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -65,16 +65,22 @@ private AbstractPlan plan( SQLQueryRequest request, Optional> queryListener, Optional> explainListener) { - // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) - ParseTree cst = parser.parse(request.getQuery()); - Statement statement = - cst.accept( - new AstStatementBuilder( - new AstBuilder(request.getQuery()), - AstStatementBuilder.StatementBuilderContext.builder() - .isExplain(request.isExplainRequest()) - .build())); + if (request.getCursor().isPresent()) { + // Handle v2 cursor here -- legacy cursor was handled earlier. + return queryExecutionFactory.create(request.getCursor().get(), queryListener.get()); + } else { + // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) + ParseTree cst = parser.parse(request.getQuery()); + Statement statement = + cst.accept( + new AstStatementBuilder( + new AstBuilder(request.getQuery()), + AstStatementBuilder.StatementBuilderContext.builder() + .isExplain(request.isExplainRequest()) + .fetchSize(request.getFetchSize()) + .build())); - return queryExecutionFactory.create(statement, queryListener, explainListener); + return queryExecutionFactory.create(statement, queryListener, explainListener); + } } } 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..bd0db9e9a00 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 @@ -28,9 +28,9 @@ @EqualsAndHashCode @RequiredArgsConstructor public class SQLQueryRequest { - - private static final Set SUPPORTED_FIELDS = ImmutableSet.of( - "query", "fetch_size", "parameters"); + private static final String QUERY_FIELD_CURSOR = "cursor"; + private static final Set SUPPORTED_FIELDS = Set.of( + "query", "fetch_size", "parameters", QUERY_FIELD_CURSOR); private static final String QUERY_PARAMS_FORMAT = "format"; private static final String QUERY_PARAMS_SANITIZE = "sanitize"; @@ -64,33 +64,38 @@ public class SQLQueryRequest { @Accessors(fluent = true) private boolean sanitize = true; + private String cursor = ""; /** * Constructor of SQLQueryRequest that passes request params. */ public SQLQueryRequest( - JSONObject jsonContent, String query, String path, Map params) { + JSONObject jsonContent, String query, String path, Map params, String cursor) { this.jsonContent = jsonContent; this.query = query; this.path = path; this.params = params; this.format = getFormat(params); this.sanitize = shouldSanitize(params); + // TODO hack + this.cursor = cursor == null? "" : cursor; } /** * Pre-check if the request can be supported by meeting ALL the following criteria: * 1.Only supported fields present in request body, ex. "filter" and "cursor" are not supported - * 2.No fetch_size or "fetch_size=0". In other word, it's not a cursor request - * 3.Response format is default or can be supported. + * 2.Response format is default or can be supported. * * @return true if supported. */ public boolean isSupported() { - return isOnlySupportedFieldInPayload() - && isFetchSizeZeroIfPresent() + return (isCursor() || isOnlySupportedFieldInPayload()) && isSupportedFormat(); } + private boolean isCursor() { + return cursor != null && cursor.isEmpty() == false; + } + /** * Check if request is to explain rather than execute the query. * @return true if it is a explain request @@ -113,11 +118,20 @@ public Format format() { } private boolean isOnlySupportedFieldInPayload() { - return SUPPORTED_FIELDS.containsAll(jsonContent.keySet()); + return jsonContent == null || SUPPORTED_FIELDS.containsAll(jsonContent.keySet()); + } + + + public Optional getCursor() { + return cursor != "" ? Optional.of(cursor) : Optional.empty(); + } + + public boolean mayReturnCursor() { + return cursor != "" || getFetchSize() > 0; } - private boolean isFetchSizeZeroIfPresent() { - return (jsonContent.optInt("fetch_size") == 0); + public int getFetchSize() { + return jsonContent.optInt("fetch_size"); } private boolean isSupportedFormat() { diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java index 40d549764a6..593e7b51ff4 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java +++ b/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java @@ -26,7 +26,7 @@ public class AstStatementBuilder extends OpenSearchSQLParserBaseVisitor { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); return null; }).when(queryService).execute(any(), any()); sqlService.execute( new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "jdbc"), - new ResponseListener() { + new ResponseListener<>() { @Override public void onResponse(QueryResponse response) { assertNotNull(response); @@ -87,7 +92,7 @@ public void onFailure(Exception e) { public void canExecuteCsvFormatRequest() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList())); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); return null; }).when(queryService).execute(any(), any()); 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..cf14d40e015 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 @@ -14,6 +14,8 @@ import com.google.common.collect.ImmutableMap; import java.util.Map; import org.json.JSONObject; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.opensearch.sql.protocol.response.format.Format; @@ -81,6 +83,7 @@ public void shouldSupportExplain() { } @Test + @Disabled("SQLQueryRequest does support cursor requests") public void shouldNotSupportCursorRequest() { SQLQueryRequest fetchSizeRequest = SQLQueryRequestBuilder.request("SELECT 1") @@ -183,7 +186,8 @@ SQLQueryRequest build() { jsonContent = "{\"query\": \"" + query + "\"}"; } if (params != null) { - return new SQLQueryRequest(new JSONObject(jsonContent), query, path, params); + return new SQLQueryRequest(new JSONObject(jsonContent), query, path, params, + ""); } return new SQLQueryRequest(new JSONObject(jsonContent), query, path, format); } From d359751c4a631b87547bd65cb40c05e92ac56287 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 30 Jan 2023 19:18:33 -0800 Subject: [PATCH 02/46] Comment to clarify an exception. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/executor/execution/QueryPlanFactory.java | 1 + 1 file changed, 1 insertion(+) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 25ccec08d52..0a449cc1b2a 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -99,6 +99,7 @@ public AbstractPlan visitQuery( queryService, context.getLeft().get()); } else { + // This should be picked up by the legacy engine. throw new UnsupportCursorRequestException(); } } else { From 9c3f7febc4b11f216c2b3e229431ba373c466e57 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Sun, 5 Feb 2023 23:58:08 -0800 Subject: [PATCH 03/46] Add support for paginated scroll request, first page. Implement PaginatedPlanCache.convertToPlan for second page to work. Signed-off-by: MaxKsyunz --- .../common/antlr/SyntaxCheckException.java | 3 +- .../sql/executor/CanPaginateVisitor.java | 5 +- .../sql/executor/PaginatedPlanCache.java | 59 +++----- .../opensearch/sql/executor/QueryService.java | 11 -- .../execution/ContinuePaginatedPlan.java | 5 +- .../sql/executor/execution/PaginatedPlan.java | 7 +- .../execution/PaginatedQueryService.java | 63 ++++++++ .../executor/execution/QueryPlanFactory.java | 6 +- .../sql/opensearch/executor/Cursor.java | 3 +- .../sql/planner/PaginateOperator.java | 13 ++ .../CreatePagingTableScanBuilder.java | 45 ++++++ .../optimizer/LogicalPlanOptimizer.java | 23 +++ .../sql/planner/physical/PhysicalPlan.java | 31 +++- .../sql/planner/physical/ProjectOperator.java | 9 ++ .../org/opensearch/sql/storage/Table.java | 5 + .../org/opensearch/sql/ppl/StandaloneIT.java | 114 +-------------- .../sql/sql/StandalonePaginationIT.java | 135 ++++++++++++++++++ .../sql/util/InternalRestHighLevelClient.java | 19 +++ .../opensearch/sql/util/StandaloneModule.java | 130 +++++++++++++++++ .../executor/OpenSearchExecutionEngine.java | 4 - .../protector/ResourceMonitorPlan.java | 5 + .../opensearch/request/OpenSearchRequest.java | 3 +- .../request/OpenSearchRequestBuilder.java | 4 +- .../request/OpenSearchScrollRequest.java | 13 ++ .../opensearch/storage/OpenSearchIndex.java | 8 ++ .../storage/OpenSearchPagedIndexScan.java | 69 +++++++++ .../OpenSearchPagedRequestBuilder.java | 92 ++++++++++++ .../storage/OpenSearchPagedScanBuilder.java | 30 ++++ .../storage/OpenSearchIndexScanTest.java | 3 +- .../storage/OpenSearchIndexTest.java | 7 +- .../plugin/config/OpenSearchPluginModule.java | 12 +- 31 files changed, 740 insertions(+), 196 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java create mode 100644 core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/util/InternalRestHighLevelClient.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java diff --git a/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java b/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java index 684e66029e5..806cb7208bb 100644 --- a/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java +++ b/common/src/main/java/org/opensearch/sql/common/antlr/SyntaxCheckException.java @@ -6,8 +6,7 @@ package org.opensearch.sql.common.antlr; -public class -SyntaxCheckException extends RuntimeException { +public class SyntaxCheckException extends RuntimeException { public SyntaxCheckException(String message) { super(message); } diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index a28d04c3bc4..e89a04216ec 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -10,9 +10,6 @@ import org.opensearch.sql.ast.expression.AllFields; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.Relation; -import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; -import org.opensearch.sql.planner.physical.ProjectOperator; -import org.opensearch.sql.storage.TableScanOperator; /** * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache. @@ -42,7 +39,7 @@ public Boolean visitRelation(Relation node, Object context) { } @Override - public Boolean visit(Node node, Object context) { + public Boolean visitChildren(Node node, Object context) { return Boolean.FALSE; } diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index d6524649ac0..0d3c4757080 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -5,7 +5,6 @@ package org.opensearch.sql.executor; -import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.ObjectInputStream; import java.util.List; @@ -23,10 +22,12 @@ import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; @RequiredArgsConstructor public class PaginatedPlanCache { + public static final String CURSOR_PREFIX = "n:"; private final StorageEngine storageEngine; public static final PaginatedPlanCache None = new PaginatedPlanCache(null); @@ -40,57 +41,37 @@ static class SeriazationContext { private final PaginatedPlanCache cache; } - public static class SerializationVisitor - extends PhysicalPlanNodeVisitor { - private static final byte[] NO_CURSOR = new byte[] {}; - - @Override - public byte[] visitPaginate(PaginateOperator node, SeriazationContext context) { - // Save cursor to read the next page. - // Could process node.getChild() here with another visitor -- one that saves the - // parameters for other physical operators -- ProjectOperator, etc. - return String.format("You got it!%d", node.getPageIndex() + 1).getBytes(); - } - - // Cursor is returned only if physical plan node is PaginateOerator. - @Override - protected byte[] visitNode(PhysicalPlan node, SeriazationContext context) { - return NO_CURSOR; - } - } - /** * Converts a physical plan tree to a cursor. May cache plan related data somewhere. */ public Cursor convertToCursor(PhysicalPlan plan) { - var serializer = new SerializationVisitor(); - var raw = plan.accept(serializer, new SeriazationContext(this)); - return new Cursor(raw); + if (plan instanceof PaginateOperator) { + var raw = CURSOR_PREFIX + plan.toCursor(); + return new Cursor(raw.getBytes()); + } else { + return Cursor.None; + } } /** - * Convers a cursor to a physical plann tree. + * Converts a cursor to a physical plan tree. */ public PhysicalPlan convertToPlan(String cursor) { - // TODO HACKY_HACK -- create a plan - if (cursor.startsWith("You got it!")) { - int pageIndex = Integer.parseInt(cursor.substring("You got it!".length())); + if (cursor.startsWith(CURSOR_PREFIX)) { + String sExpression = cursor.substring(CURSOR_PREFIX.length()); - Table table = storageEngine.getTable(null, "phrases"); - TableScanBuilder scanBuilder = table.createScanBuilder(); - scanBuilder.pushDownOffset(5 * pageIndex); - PhysicalPlan scan = scanBuilder.build(); - var fields = table.getFieldTypes(); - List references = - Stream.of("phrase", "test field", "insert_time2") - .map(c -> - new NamedExpression(c, new ReferenceExpression(c, List.of(c), fields.get(c)))) - .collect(Collectors.toList()); + // TODO Parse sExpression and initialize variables below. + // storageEngine needs to create the TableScanOperator. + int pageSize = -1; + int currentPageIndex = -1; + List projectList = List.of(); + TableScanOperator scan = null; - return new PaginateOperator(new ProjectOperator(scan, references, List.of()), 5, pageIndex); + return new PaginateOperator(new ProjectOperator(scan, projectList, List.of()), + pageSize, currentPageIndex); } else { - throw new RuntimeException("Unsupported cursor"); + throw new UnsupportedOperationException("Unsupported cursor"); } } } diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index f2825a7149c..ed251e2b33c 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -70,17 +70,6 @@ public void executePlan(LogicalPlan plan, } } - /** - * Execute a physical plan without analyzing or planning anything. - */ - public void executePlan(PhysicalPlan plan, - ResponseListener listener) { - try { - executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); - } catch (Exception e) { - listener.onFailure(e); - } - } /** * Explain the query in {@link UnresolvedPlan} using {@link ResponseListener} to diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java index 15d4d973513..0115e6de731 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java @@ -10,7 +10,6 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryId; -import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.planner.physical.PhysicalPlan; public class ContinuePaginatedPlan extends AbstractPlan { @@ -19,7 +18,7 @@ public class ContinuePaginatedPlan extends AbstractPlan { = new ContinuePaginatedPlan(QueryId.None, "", null, null, null); private final String cursor; - private final QueryService queryService; + private final PaginatedQueryService queryService; private final PaginatedPlanCache paginatedPlanCache; private final ResponseListener queryResponseListener; @@ -28,7 +27,7 @@ public class ContinuePaginatedPlan extends AbstractPlan { /** * Create an abstract plan that can continue paginating a given cursor. */ - public ContinuePaginatedPlan(QueryId queryId, String cursor, QueryService queryService, + public ContinuePaginatedPlan(QueryId queryId, String cursor, PaginatedQueryService queryService, PaginatedPlanCache ppc, ResponseListener queryResponseListener) { diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java index a0d4a1eaea7..62c913ee672 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java @@ -10,12 +10,11 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryId; -import org.opensearch.sql.executor.QueryService; public class PaginatedPlan extends AbstractPlan { final UnresolvedPlan plan; final int fetchSize; - final QueryService queryService; + final PaginatedQueryService queryService; final ResponseListener queryResponseResponseListener; @@ -23,7 +22,7 @@ public class PaginatedPlan extends AbstractPlan { * Create an abstract plan that can start paging a query. */ public PaginatedPlan(QueryId queryId, UnresolvedPlan plan, int fetchSize, - QueryService queryService, + PaginatedQueryService queryService, ResponseListener queryResponseResponseListener) { super(queryId); @@ -40,6 +39,6 @@ public void execute() { @Override public void explain(ResponseListener listener) { - + throw new UnsupportedOperationException("implement PaginatedPlan.explain"); } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java new file mode 100644 index 00000000000..e08c6dc59da --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java @@ -0,0 +1,63 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor.execution; + +import lombok.RequiredArgsConstructor; +import org.opensearch.sql.analysis.AnalysisContext; +import org.opensearch.sql.analysis.Analyzer; +import org.opensearch.sql.ast.tree.Paginate; +import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.ExecutionContext; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.planner.Planner; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.physical.PhysicalPlan; + +@RequiredArgsConstructor +public class PaginatedQueryService { + private final Analyzer analyzer; + + private final ExecutionEngine executionEngine; + + private final Planner planner; + + /** + * Execute a pagination request. Passes the exception the listener. + */ + public void execute(Paginate plan, ResponseListener listener) { + try { + executePlan(analyze(plan), listener); + } catch (Exception e) { + listener.onFailure(e); + } + } + + public void executePlan(LogicalPlan plan, + ResponseListener listener) { + executionEngine.execute(plan(plan), ExecutionContext.emptyExecutionContext(), listener); + } + + /** + * Execute a physical plan without analyzing or planning anything. + */ + public void executePlan(PhysicalPlan plan, + ResponseListener listener) { + try { + executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); + } catch (Exception e) { + listener.onFailure(e); + } + } + + public LogicalPlan analyze(UnresolvedPlan plan) { + return analyzer.analyze(plan, new AnalysisContext()); + } + + public PhysicalPlan plan(LogicalPlan plan) { + return planner.plan(plan); + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 0a449cc1b2a..881046ec0b9 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -12,7 +12,6 @@ import com.google.common.base.Preconditions; import java.util.Optional; import lombok.RequiredArgsConstructor; -import org.apache.commons.lang3.NotImplementedException; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.statement.Explain; @@ -40,6 +39,7 @@ public class QueryPlanFactory * Query Service. */ private final QueryService queryService; + private final PaginatedQueryService paginatedQueryService; private final PaginatedPlanCache paginatedPlanCache; /** @@ -79,7 +79,7 @@ public AbstractPlan create( public AbstractPlan create(String cursor, ResponseListener queryResponseListener) { QueryId queryId = QueryId.queryId(); - return new ContinuePaginatedPlan(queryId, cursor, queryService, paginatedPlanCache, + return new ContinuePaginatedPlan(queryId, cursor, paginatedQueryService, paginatedPlanCache, queryResponseListener); } @@ -96,7 +96,7 @@ public AbstractPlan visitQuery( if (node.getFetchSize() > 0) { if (paginatedPlanCache.canConvertToCursor(node.getPlan())) { return new PaginatedPlan(QueryId.queryId(), node.getPlan(), node.getFetchSize(), - queryService, + paginatedQueryService, context.getLeft().get()); } else { // This should be picked up by the legacy engine. diff --git a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java index 02d06c6aec6..6751ca107c6 100644 --- a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java +++ b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java @@ -23,8 +23,7 @@ public Cursor(byte[] raw) { this.raw = raw; } - @Override - public String toString() { + public String asString() { return new String(raw); } } diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java index f1add36e79e..227d4b4602c 100644 --- a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java @@ -72,4 +72,17 @@ public ExecutionEngine.Schema schema() { assert input instanceof ProjectOperator; return input.schema(); } + + @Override + public String toCursor() { + // Save cursor to read the next page. + // Could process node.getChild() here with another visitor -- one that saves the + // parameters for other physical operators -- ProjectOperator, etc. + // cursor format: n:|" + String child = getChild().get(0).toCursor(); + + var nextPage = getPageIndex() + 1; + return createSection("Paginate", Integer.toString(nextPage), + Integer.toString(getPageSize()), child); + } } diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java new file mode 100644 index 00000000000..bc97e373c22 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner.optimizer; + +import static org.opensearch.sql.planner.optimizer.pattern.Patterns.table; + +import com.facebook.presto.matching.Capture; +import com.facebook.presto.matching.Captures; +import com.facebook.presto.matching.Pattern; +import lombok.Getter; +import lombok.experimental.Accessors; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalRelation; +import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.read.TableScanBuilder; + +public class CreatePagingTableScanBuilder + implements Rule { + /** Capture the table inside matched logical relation operator. */ + private final Capture capture; + + /** Pattern that matches logical relation operator. */ + @Accessors(fluent = true) + @Getter + private final Pattern pattern; + + /** + * Constructor. + */ + public CreatePagingTableScanBuilder() { + this.capture = Capture.newCapture(); + this.pattern = Pattern.typeOf(LogicalRelation.class) + .with(table().capturedAs(capture)); + } + + @Override + public LogicalPlan apply(LogicalRelation plan, Captures captures) { + TableScanBuilder scanBuilder = captures.get(capture).createPagedScanBuilder(); + // TODO: Remove this after Prometheus refactored to new table scan builder too + return (scanBuilder == null) ? plan : scanBuilder; + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java index 70847b869b5..58a96c3efc6 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java @@ -60,6 +60,29 @@ public static LogicalPlanOptimizer create() { new CreateTableWriteBuilder())); } + /** + * Create {@link LogicalPlanOptimizer} with pre-defined rules. + */ + public static LogicalPlanOptimizer paginationCreate() { + return new LogicalPlanOptimizer(Arrays.asList( + /* + * Phase 1: Transformations that rely on relational algebra equivalence + */ + new MergeFilterAndFilter(), + new PushFilterUnderSort(), + /* + * Phase 2: Transformations that rely on data source push down capability + */ + new CreatePagingTableScanBuilder(), + TableScanPushDown.PUSH_DOWN_FILTER, + TableScanPushDown.PUSH_DOWN_AGGREGATION, + TableScanPushDown.PUSH_DOWN_SORT, + TableScanPushDown.PUSH_DOWN_LIMIT, + TableScanPushDown.PUSH_DOWN_HIGHLIGHT, + TableScanPushDown.PUSH_DOWN_PROJECT, + new CreateTableWriteBuilder())); + } + /** * Optimize {@link LogicalPlan}. */ 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 92e032f8d3a..8b7f1cad576 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 @@ -8,6 +8,7 @@ import java.io.Serializable; import java.util.Iterator; +import java.util.List; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.planner.PlanNode; @@ -18,8 +19,10 @@ */ public abstract class PhysicalPlan implements PlanNode, Iterator, - AutoCloseable, - Serializable { + AutoCloseable { + + public static final List FORBIDDEN_CHARS = List.of("(", ")", ","); + /** * Accept the {@link PhysicalPlanNodeVisitor}. * @@ -48,4 +51,28 @@ public ExecutionEngine.Schema schema() { + "ProjectOperator, instead of %s", toString())); } + public String toCursor() { + throw new IllegalStateException(String.format("%s needs to implement ToCursor", + this.getClass())); + } + + /** + * Creates an S-expression that represents a plan node. + * @param plan Label for the plan. + * @param params List of serialized parameters. Including the child plans. + * @return A string that represents the plan called with those parameters. + */ + protected String createSection(String plan, String... params) { + if (FORBIDDEN_CHARS.stream().anyMatch(plan::contains)) { + var error = String.format("plan key '%s' contains forbidden character", + plan); + throw new RuntimeException(error); + } + + // TODO: check that each param is either a valid s-expression or + // does not contain forbidden characters. + return "(" + plan + "," + + String.join(",", params) + + ")"; + } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 496e4e6ddb1..91b07d2d849 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -94,4 +94,13 @@ public ExecutionEngine.Schema schema() { .map(expr -> new ExecutionEngine.Schema.Column(expr.getName(), expr.getAlias(), expr.type())).collect(Collectors.toList())); } + + @Override + public String toCursor() { + String child = getChild().get(0).toCursor(); + String namedExpressions = "TODO"; + String parseExpressions = "TODO"; + // TODO serialize named expressions and parse expressions + return createSection("Project", namedExpressions, parseExpressions, child); + } } diff --git a/core/src/main/java/org/opensearch/sql/storage/Table.java b/core/src/main/java/org/opensearch/sql/storage/Table.java index 496281fa8d7..8117e2cc307 100644 --- a/core/src/main/java/org/opensearch/sql/storage/Table.java +++ b/core/src/main/java/org/opensearch/sql/storage/Table.java @@ -92,4 +92,9 @@ default TableWriteBuilder createWriteBuilder(LogicalWrite plan) { default StreamingSource asStreamingSource() { throw new UnsupportedOperationException(); } + + default TableScanBuilder createPagedScanBuilder() { + var error = String.format("'%s' does not support pagination", getClass().toString()); + throw new UnsupportedOperationException(error); + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java index 28c4250f8b6..5d8afb0e516 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java @@ -12,54 +12,29 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import java.io.IOException; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; -import lombok.RequiredArgsConstructor; import org.junit.jupiter.api.Test; import org.opensearch.client.Request; -import org.opensearch.client.RestClient; import org.opensearch.client.RestHighLevelClient; -import org.opensearch.common.inject.AbstractModule; import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; -import org.opensearch.common.inject.Provides; -import org.opensearch.common.inject.Singleton; -import org.opensearch.sql.analysis.Analyzer; -import org.opensearch.sql.analysis.ExpressionAnalyzer; +import org.opensearch.sql.util.InternalRestHighLevelClient; +import org.opensearch.sql.util.StandaloneModule; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasource.DataSourceServiceImpl; -import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; -import org.opensearch.sql.executor.PaginatedPlanCache; -import org.opensearch.sql.executor.QueryManager; -import org.opensearch.sql.executor.QueryService; -import org.opensearch.sql.executor.execution.QueryPlanFactory; -import org.opensearch.sql.expression.function.BuiltinFunctionRepository; -import org.opensearch.sql.monitor.AlwaysHealthyMonitor; -import org.opensearch.sql.monitor.ResourceMonitor; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.client.OpenSearchRestClient; -import org.opensearch.sql.opensearch.executor.OpenSearchExecutionEngine; -import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; -import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.security.SecurityAccess; import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory; -import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; -import org.opensearch.sql.planner.Planner; -import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; -import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; import org.opensearch.sql.ppl.domain.PPLQueryRequest; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.SimpleJsonResponseFormatter; -import org.opensearch.sql.sql.SQLService; -import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.storage.DataSourceFactory; -import org.opensearch.sql.storage.StorageEngine; -import org.opensearch.sql.util.ExecuteOnCallerThreadQueryManager; /** * Run PPL with query engine outside OpenSearch cluster. This IT doesn't require our plugin @@ -68,13 +43,11 @@ */ public class StandaloneIT extends PPLIntegTestCase { - private RestHighLevelClient restClient; - private PPLService pplService; @Override public void init() { - restClient = new InternalRestHighLevelClient(client()); + RestHighLevelClient restClient = new InternalRestHighLevelClient(client()); OpenSearchClient client = new OpenSearchRestClient(restClient); DataSourceService dataSourceService = new DataSourceServiceImpl( new ImmutableSet.Builder() @@ -160,85 +133,4 @@ public List getSettings() { }; } - /** - * Internal RestHighLevelClient only for testing purpose. - */ - static class InternalRestHighLevelClient extends RestHighLevelClient { - public InternalRestHighLevelClient(RestClient restClient) { - super(restClient, RestClient::close, Collections.emptyList()); - } - } - - @RequiredArgsConstructor - public class StandaloneModule extends AbstractModule { - - private final RestHighLevelClient client; - - private final Settings settings; - - private final DataSourceService dataSourceService; - - private final BuiltinFunctionRepository functionRepository = - BuiltinFunctionRepository.getInstance(); - - @Override - protected void configure() {} - - @Provides - public OpenSearchClient openSearchClient() { - return new OpenSearchRestClient(client); - } - - @Provides - public StorageEngine storageEngine(OpenSearchClient client) { - return new OpenSearchStorageEngine(client, settings); - } - - @Provides - public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector, - PaginatedPlanCache paginatedPlanCache) { - return new OpenSearchExecutionEngine(client, protector, paginatedPlanCache); - } - - @Provides - public ResourceMonitor resourceMonitor() { - return new AlwaysHealthyMonitor(); - } - - @Provides - public ExecutionProtector protector(ResourceMonitor resourceMonitor) { - return new OpenSearchExecutionProtector(resourceMonitor); - } - - @Provides - @Singleton - public QueryManager queryManager() { - return new ExecuteOnCallerThreadQueryManager(); - } - - @Provides - public PPLService pplService(QueryManager queryManager, QueryPlanFactory queryPlanFactory) { - return new PPLService(new PPLSyntaxParser(), queryManager, queryPlanFactory); - } - - @Provides - public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPlanFactory) { - return new SQLService(new SQLSyntaxParser(), queryManager, queryPlanFactory); - } - - @Provides - public PaginatedPlanCache paginatedPlanCache(StorageEngine storageEngine) { - return new PaginatedPlanCache(storageEngine); - } - @Provides - public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, - PaginatedPlanCache paginatedPlanCache) { - Analyzer analyzer = - new Analyzer( - new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); - Planner planner = new Planner(LogicalPlanOptimizer.create()); - return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner), - paginatedPlanCache); - } - } } diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java new file mode 100644 index 00000000000..046bb5679e1 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -0,0 +1,135 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.sql; + +import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import java.io.IOException; +import java.util.List; +import java.util.Map; +import lombok.Getter; +import org.junit.Test; +import org.opensearch.client.Request; +import org.opensearch.client.RestHighLevelClient; +import org.opensearch.common.inject.Injector; +import org.opensearch.common.inject.ModulesBuilder; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.data.model.ExprBooleanValue; +import org.opensearch.sql.datasource.DataSourceService; +import org.opensearch.sql.datasource.DataSourceServiceImpl; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.executor.execution.PaginatedQueryService; +import org.opensearch.sql.expression.LiteralExpression; +import org.opensearch.sql.expression.NamedExpression; +import org.opensearch.sql.legacy.SQLIntegTestCase; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.client.OpenSearchRestClient; +import org.opensearch.sql.opensearch.executor.Cursor; +import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalProject; +import org.opensearch.sql.planner.logical.LogicalRelation; +import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.storage.DataSourceFactory; +import org.opensearch.sql.util.InternalRestHighLevelClient; +import org.opensearch.sql.util.StandaloneModule; + +public class StandalonePaginationIT extends SQLIntegTestCase { + + private PaginatedQueryService paginatedQueryService; + + private PaginatedPlanCache paginatedPlanCache; + + private OpenSearchClient client; + + @Override + public void init() { + RestHighLevelClient restClient = new InternalRestHighLevelClient(client()); + client = new OpenSearchRestClient(restClient); + DataSourceService dataSourceService = new DataSourceServiceImpl( + new ImmutableSet.Builder() + .add(new OpenSearchDataSourceFactory(client, defaultSettings())) + .build()); + dataSourceService.addDataSource(defaultOpenSearchDataSourceMetadata()); + + ModulesBuilder modules = new ModulesBuilder(); + modules.add(new StandaloneModule(new InternalRestHighLevelClient(client()), defaultSettings(), dataSourceService)); + Injector injector = modules.createInjector(); + + paginatedQueryService = injector.getInstance(PaginatedQueryService.class); + paginatedPlanCache = injector.getInstance(PaginatedPlanCache.class); + } + + @Test + public void testPagination() throws IOException { + class TestResponder + implements ResponseListener { + @Getter + Cursor cursor = Cursor.None; + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + cursor = response.getCursor(); + assertTrue(true); + } + + @Override + public void onFailure(Exception e) { + + assertFalse(true); + } + }; + + // arrange + { + Request request1 = new Request("PUT", "/test/_doc/1?refresh=true"); + request1.setJsonEntity("{\"name\": \"hello\", \"age\": 20}"); + client().performRequest(request1); + Request request2 = new Request("PUT", "/test/_doc/2?refresh=true"); + request2.setJsonEntity("{\"name\": \"world\", \"age\": 30}"); + client().performRequest(request2); + } + + // act 1, asserts in firstResponder + var t = new OpenSearchIndex(client, defaultSettings(), "test"); + LogicalPlan p = new LogicalPaginate(1, List.of(new LogicalProject( + new LogicalRelation("test", t), + List.of(new NamedExpression("count()", new LiteralExpression(ExprBooleanValue.of(true)))), + List.of() + ))); + var firstResponder = new TestResponder(); + paginatedQueryService.executePlan(p, firstResponder); + + // act 2, asserts in secondResponder + + PhysicalPlan plan = paginatedPlanCache.convertToPlan(firstResponder.getCursor().asString()); + var secondResponder = new TestResponder(); + paginatedQueryService.executePlan(plan, secondResponder); + } + + private Settings defaultSettings() { + return new Settings() { + private final Map defaultSettings = new ImmutableMap.Builder() + .put(Key.QUERY_SIZE_LIMIT, 200) + .build(); + + @Override + public T getSettingValue(Key key) { + return (T) defaultSettings.get(key); + } + + @Override + public List getSettings() { + return (List) defaultSettings; + } + }; + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/util/InternalRestHighLevelClient.java b/integ-test/src/test/java/org/opensearch/sql/util/InternalRestHighLevelClient.java new file mode 100644 index 00000000000..57726089ae7 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/util/InternalRestHighLevelClient.java @@ -0,0 +1,19 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.util; + +import java.util.Collections; +import org.opensearch.client.RestClient; +import org.opensearch.client.RestHighLevelClient; + +/** + * Internal RestHighLevelClient only for testing purpose. + */ +public class InternalRestHighLevelClient extends RestHighLevelClient { + public InternalRestHighLevelClient(RestClient restClient) { + super(restClient, RestClient::close, Collections.emptyList()); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java new file mode 100644 index 00000000000..3551c22a1d6 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java @@ -0,0 +1,130 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.util; + +import lombok.RequiredArgsConstructor; +import org.opensearch.client.RestHighLevelClient; +import org.opensearch.common.inject.AbstractModule; +import org.opensearch.common.inject.Provides; +import org.opensearch.common.inject.Singleton; +import org.opensearch.sql.analysis.Analyzer; +import org.opensearch.sql.analysis.ExpressionAnalyzer; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.datasource.DataSourceService; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.executor.QueryManager; +import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.execution.PaginatedQueryService; +import org.opensearch.sql.executor.execution.QueryPlanFactory; +import org.opensearch.sql.expression.function.BuiltinFunctionRepository; +import org.opensearch.sql.monitor.AlwaysHealthyMonitor; +import org.opensearch.sql.monitor.ResourceMonitor; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.client.OpenSearchRestClient; +import org.opensearch.sql.opensearch.executor.OpenSearchExecutionEngine; +import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; +import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; +import org.opensearch.sql.planner.Planner; +import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; +import org.opensearch.sql.ppl.PPLService; +import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; +import org.opensearch.sql.sql.SQLService; +import org.opensearch.sql.sql.antlr.SQLSyntaxParser; +import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.sql.util.ExecuteOnCallerThreadQueryManager; + +@RequiredArgsConstructor +public class StandaloneModule extends AbstractModule { + + private final RestHighLevelClient client; + + private final Settings settings; + + private final DataSourceService dataSourceService; + + private final BuiltinFunctionRepository functionRepository = + BuiltinFunctionRepository.getInstance(); + + @Override + protected void configure() { + } + + @Provides + public OpenSearchClient openSearchClient() { + return new OpenSearchRestClient(client); + } + + @Provides + public StorageEngine storageEngine(OpenSearchClient client) { + return new OpenSearchStorageEngine(client, settings); + } + + @Provides + public ExecutionEngine executionEngine(OpenSearchClient client, ExecutionProtector protector, + PaginatedPlanCache paginatedPlanCache) { + return new OpenSearchExecutionEngine(client, protector, paginatedPlanCache); + } + + @Provides + public ResourceMonitor resourceMonitor() { + return new AlwaysHealthyMonitor(); + } + + @Provides + public ExecutionProtector protector(ResourceMonitor resourceMonitor) { + return new OpenSearchExecutionProtector(resourceMonitor); + } + + @Provides + @Singleton + public QueryManager queryManager() { + return new ExecuteOnCallerThreadQueryManager(); + } + + @Provides + public PPLService pplService(QueryManager queryManager, QueryPlanFactory queryPlanFactory) { + return new PPLService(new PPLSyntaxParser(), queryManager, queryPlanFactory); + } + + @Provides + public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPlanFactory) { + return new SQLService(new SQLSyntaxParser(), queryManager, queryPlanFactory); + } + + @Provides + public PaginatedPlanCache paginatedPlanCache(StorageEngine storageEngine) { + return new PaginatedPlanCache(storageEngine); + } + + @Provides + public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, + PaginatedPlanCache paginatedPlanCache, + QueryService qs, + PaginatedQueryService pqs) { + + return new QueryPlanFactory(qs, pqs, paginatedPlanCache); + } + + @Provides + public QueryService querySerivce(ExecutionEngine executionEngine) { + Analyzer analyzer = + new Analyzer( + new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); + Planner planner = new Planner(LogicalPlanOptimizer.create()); + return new QueryService(analyzer, executionEngine, planner); + } + + @Provides + public PaginatedQueryService paginatedQueryService(ExecutionEngine executionEngine) { + Analyzer analyzer = + new Analyzer( + new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); + Planner planner = new Planner(LogicalPlanOptimizer.paginationCreate()); + return new PaginatedQueryService(analyzer, executionEngine, planner); + } +} 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 1f4fb578b12..cec2864c11f 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 @@ -51,10 +51,6 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context, result.add(plan.next()); } - - // - // getContinuation expects hasNext to return false before it is called. - Cursor qc = paginatedPlanCache.convertToCursor(plan); QueryResponse response = new QueryResponse(physicalPlan.schema(), result, qc); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index 9c59e4acaf8..307b40dce79 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -82,4 +82,9 @@ public ExprValue next() { } return delegate.next(); } + + @Override + public String toCursor() { + return delegate.toCursor(); + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index 0de204c285b..ce990780c1c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -6,7 +6,6 @@ package org.opensearch.sql.opensearch.request; -import java.io.Serializable; import java.util.function.Consumer; import java.util.function.Function; import lombok.EqualsAndHashCode; @@ -20,7 +19,7 @@ /** * OpenSearch search request. */ -public interface OpenSearchRequest extends Serializable { +public interface OpenSearchRequest { /** * Apply the search action or scroll action on request based on context. * diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index ee97dc3ae8b..a432bf1ca8f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -9,7 +9,6 @@ import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; import static org.opensearch.search.sort.SortOrder.ASC; -import java.io.Serializable; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -43,7 +42,7 @@ @EqualsAndHashCode @Getter @ToString -public class OpenSearchRequestBuilder implements Serializable { +public class OpenSearchRequestBuilder { /** * Default query timeout in minutes. @@ -105,6 +104,7 @@ public OpenSearchRequestBuilder(OpenSearchRequest.IndexName indexName, /** * Build DSL request. * + * @return query request or scroll request */ public OpenSearchRequest build() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 4509e443c0a..5f4e8a77cf6 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -51,6 +51,7 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { * multi-thread so this state has to be maintained here. */ @Setter + @Getter private String scrollId; /** Search request source builder. */ @@ -140,4 +141,16 @@ public SearchScrollRequest scrollRequest() { public void reset() { scrollId = null; } + + /** + * Convert a scroll request to string that can be included in a cursor. + * @return a string representing the scroll request. + */ + public String toCursor() { + if (isScrollStarted()) { + return scrollId; + } else { + return "TESTING ONLY"; + } + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index d7afac9060a..e293c8b0e78 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -127,6 +127,14 @@ public TableScanBuilder createScanBuilder() { return new OpenSearchIndexScanBuilder(indexScan); } + @Override + public TableScanBuilder createPagedScanBuilder() { + var requestBuilder = new OpenSearchPagedRequestBuilder(indexName, + settings, new OpenSearchExprValueFactory(getFieldTypes())); + var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); + return new OpenSearchPagedScanBuilder(indexScan); + } + @VisibleForTesting @RequiredArgsConstructor public static class OpenSearchDefaultImplementor diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java new file mode 100644 index 00000000000..9ff5cfa1da4 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java @@ -0,0 +1,69 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import java.util.Iterator; +import lombok.EqualsAndHashCode; +import lombok.ToString; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.storage.TableScanOperator; + +@EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) +@ToString(onlyExplicitlyIncluded = true) +public class OpenSearchPagedIndexScan extends TableScanOperator { + private final OpenSearchClient client; + private final OpenSearchPagedRequestBuilder requestBuilder; + @EqualsAndHashCode.Include + @ToString.Include + private OpenSearchScrollRequest request; + private Iterator iterator; + + public OpenSearchPagedIndexScan(OpenSearchClient client, + OpenSearchPagedRequestBuilder requestBuilder) { + this.client = client; + this.requestBuilder = requestBuilder; + } + + @Override + public String explain() { + throw new RuntimeException("Implement OpenSearchPagedIndexScan.explain"); + } + + @Override + public boolean hasNext() { + return iterator.hasNext(); + } + + @Override + public ExprValue next() { + return iterator.next(); + } + + @Override + public void open() { + super.open(); + request = requestBuilder.build(); + OpenSearchResponse response = client.search(request); + if (!response.isEmpty()) { + iterator = response.iterator(); + } + } + + @Override + public void close() { + super.close(); + + client.cleanup(request); + } + + @Override + public String toCursor() { + return createSection("OpenSearchPagedIndexScan", request.toCursor()); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java new file mode 100644 index 00000000000..9e15fd12f67 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java @@ -0,0 +1,92 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.commons.lang3.tuple.Pair; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.search.aggregations.AggregationBuilder; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.SortBuilder; +import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; +import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; + +public class OpenSearchPagedRequestBuilder { + + + private final OpenSearchRequest.IndexName indexName; + private final SearchSourceBuilder sourceBuilder; + private final OpenSearchExprValueFactory exprValueFactory; + private final int querySize; + + /** + * Constructor. + * @param indexName index being scanned + * @param settings other settings + * @param exprValueFactory value factory + */ + public OpenSearchPagedRequestBuilder(OpenSearchRequest.IndexName indexName, Settings settings, + OpenSearchExprValueFactory exprValueFactory) { + this.indexName = indexName; + this.sourceBuilder = new SearchSourceBuilder(); + this.exprValueFactory = exprValueFactory; + this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + sourceBuilder.from(0); + sourceBuilder.size(querySize); + sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); + } + + public OpenSearchScrollRequest build() { + return new OpenSearchScrollRequest(indexName, sourceBuilder, exprValueFactory); + } + + public void pushDown(QueryBuilder query) { + throw new RuntimeException(); + } + + public void pushDownAggregation( + Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { + + throw new UnsupportedOperationException("pagination of aggregation requests is not supported"); + } + + public void pushDownSort(List> sortBuilders) { + throw new UnsupportedOperationException("sorting of paged requests is not supported"); + + } + + public void pushDownLimit(Integer limit, Integer offset) { + throw new UnsupportedOperationException("limit of paged requests is not supported"); + } + + public void pushDownHighlight(String field, Map arguments) { + throw new UnsupportedOperationException("highlight of paged requests is not supported"); + } + + /** + * Push down project expression to OpenSearch. + */ + public void pushDownProjects(Set projects) { + final Set projectsSet = + projects.stream().map(ReferenceExpression::getAttr).collect(Collectors.toSet()); + sourceBuilder.fetchSource(projectsSet.toArray(new String[0]), new String[0]); + } + + public void pushTypeMapping(Map typeMapping) { + exprValueFactory.setTypeMapping(typeMapping); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java new file mode 100644 index 00000000000..09232cb4de7 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import lombok.EqualsAndHashCode; +import org.opensearch.sql.storage.TableScanOperator; +import org.opensearch.sql.storage.read.TableScanBuilder; + +/** + * Builder for a paged opensearch request. + * Override pushDown* methods from TableScaneBuilder as more features + * support pagination. + */ +public class OpenSearchPagedScanBuilder extends TableScanBuilder { + @EqualsAndHashCode.Include + OpenSearchPagedIndexScan indexScan; + + public OpenSearchPagedScanBuilder(OpenSearchPagedIndexScan indexScan) { + this.indexScan = indexScan; + } + + + @Override + public TableScanOperator build() { + return indexScan; + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java index 4cb86015fa6..bd0d0089fb3 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java @@ -82,7 +82,8 @@ void queryAllResultsWithQuery() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder("employees", 10, settings, + OpenSearchRequestBuilder + builder = new OpenSearchRequestBuilder("employees", 10, settings, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, builder)) { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 890951e4d3a..4adacb7dd28 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -23,7 +23,6 @@ import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.OPENSEARCH_TEXT_KEYWORD; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.eval; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.remove; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.rename; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; @@ -160,7 +159,8 @@ void implementRelationOperatorOnly() { LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, + OpenSearchRequestBuilder + builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, settings, exprValueFactory); assertEquals(new OpenSearchIndexScan(client, builder), index.implement(plan)); } @@ -172,7 +172,8 @@ void implementRelationOperatorWithOptimization() { LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, + OpenSearchRequestBuilder + builder = new OpenSearchRequestBuilder(indexName, maxResultWindow, settings, exprValueFactory); assertEquals( new OpenSearchIndexScan(client, builder), diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 8fea281586c..09aa88a3f50 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -18,6 +18,7 @@ import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.execution.PaginatedQueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.expression.function.BuiltinFunctionRepository; import org.opensearch.sql.monitor.ResourceMonitor; @@ -99,12 +100,17 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl * {@link QueryPlanFactory}. */ @Provides - public QueryPlanFactory queryPlanFactory( - DataSourceService dataSourceService, ExecutionEngine executionEngine, PaginatedPlanCache p) { + public QueryPlanFactory queryPlanFactory(DataSourceService dataSourceService, + ExecutionEngine executionEngine, + PaginatedPlanCache paginatedPlanCache) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); Planner planner = new Planner(LogicalPlanOptimizer.create()); - return new QueryPlanFactory(new QueryService(analyzer, executionEngine, planner), p); + Planner paginationPlanner = new Planner(LogicalPlanOptimizer.paginationCreate()); + QueryService queryService = new QueryService(analyzer, executionEngine, planner); + PaginatedQueryService paginatedQueryService + = new PaginatedQueryService(analyzer, executionEngine, paginationPlanner); + return new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache); } } From 1ee718b20a4e4734ff86f9572c08eca15cb820d2 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 6 Feb 2023 01:16:55 -0800 Subject: [PATCH 04/46] Progress on paginated scroll request, subsequent page. Signed-off-by: MaxKsyunz --- .../sql/executor/PaginatedPlanCache.java | 7 ++- .../sql/planner/physical/ProjectOperator.java | 6 +- .../opensearch/sql/storage/StorageEngine.java | 5 ++ .../sql/executor/PaginatedPlanCacheTest.java | 4 +- .../execution/QueryPlanFactoryTest.java | 5 +- .../planner/physical/ProjectOperatorTest.java | 24 -------- .../sql/sql/StandalonePaginationIT.java | 15 +++-- .../opensearch/request/OpenSearchRequest.java | 4 ++ .../request/OpenSearchScrollRequest.java | 5 +- .../storage/ContinueScrollRequest.java | 59 +++++++++++++++++++ ...er.java => InitialPageRequestBuilder.java} | 12 ++-- .../opensearch/storage/OpenSearchIndex.java | 2 +- .../storage/OpenSearchPagedIndexScan.java | 11 ++-- .../storage/OpenSearchStorageEngine.java | 16 +++++ .../storage/PagedRequestBuilder.java | 14 +++++ .../storage/SubsequentPageRequestBuilder.java | 32 ++++++++++ 16 files changed, 171 insertions(+), 50 deletions(-) create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java rename opensearch/src/main/java/org/opensearch/sql/opensearch/storage/{OpenSearchPagedRequestBuilder.java => InitialPageRequestBuilder.java} (89%) create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 0d3c4757080..25c94080604 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -58,14 +58,15 @@ public Cursor convertToCursor(PhysicalPlan plan) { */ public PhysicalPlan convertToPlan(String cursor) { if (cursor.startsWith(CURSOR_PREFIX)) { - String sExpression = cursor.substring(CURSOR_PREFIX.length()); + String expression = cursor.substring(CURSOR_PREFIX.length()); - // TODO Parse sExpression and initialize variables below. + // TODO Parse expression and initialize variables below. // storageEngine needs to create the TableScanOperator. int pageSize = -1; int currentPageIndex = -1; List projectList = List.of(); - TableScanOperator scan = null; + String scanAsString = ""; + TableScanOperator scan = storageEngine.getTableScan(scanAsString); return new PaginateOperator(new ProjectOperator(scan, projectList, List.of()), pageSize, currentPageIndex); diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 91b07d2d849..9a3542ad1fc 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -99,8 +99,8 @@ public ExecutionEngine.Schema schema() { public String toCursor() { String child = getChild().get(0).toCursor(); String namedExpressions = "TODO"; - String parseExpressions = "TODO"; - // TODO serialize named expressions and parse expressions - return createSection("Project", namedExpressions, parseExpressions, child); + // TODO serialize named expressions. + // Skipping parsedExpressions for now. + return createSection("Project", namedExpressions, child); } } diff --git a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java index 73a96a8b7c6..e854477b72b 100644 --- a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java +++ b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java @@ -31,4 +31,9 @@ public interface StorageEngine { default Collection getFunctions() { return Collections.emptyList(); } + + default TableScanOperator getTableScan(String scanAsString) { + String error = String.format("%s.getTableScan needs to be implemented", getClass()); + throw new UnsupportedOperationException(error); + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java index 21b99c6a36d..0b879a8c2de 100644 --- a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java @@ -5,10 +5,7 @@ package org.opensearch.sql.executor; -import static org.junit.jupiter.api.Assertions.*; - import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; @@ -21,6 +18,7 @@ class PaginatedPlanCacheTest { StorageEngine storageEngine; PaginatedPlanCache planCache; + @BeforeEach void setUp() { planCache = new PaginatedPlanCache(storageEngine); diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index b1940022aaa..f6abaa0d541 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -46,13 +46,16 @@ class QueryPlanFactoryTest { @Mock private ExecutionEngine.QueryResponse queryResponse; + @Mock + private PaginatedQueryService paginatedQueryService; + @Mock private PaginatedPlanCache paginatedPlanCache; private QueryPlanFactory factory; @BeforeEach void init() { - factory = new QueryPlanFactory(queryService, paginatedPlanCache); + factory = new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache); } @Test diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index 6afa9db3bd8..4ab8253c040 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -213,28 +213,4 @@ public void project_parse_missing_will_fallback() { ExprValueUtils.tupleValue(ImmutableMap.of("action", "GET", "response", "200")), ExprValueUtils.tupleValue(ImmutableMap.of("action", "POST"))))); } - - @Test - public void project_serialize() throws IOException, ClassNotFoundException { - PhysicalPlan plan = project(inputPlan, DSL.named("action", DSL.ref("action", STRING))); - - var os = new ByteArrayOutputStream(); - var objstream = new ObjectOutputStream(os); - objstream.writeObject(plan); - objstream.close(); - os.close(); - String result = os.toString(); - - - var is = new ByteArrayInputStream(os.toByteArray()); - var outstream = new ObjectInputStream(is); - var newObj = outstream.readObject(); - - assertThat(newObj, instanceOf(ProjectOperator.class)); - - var newOp = (ProjectOperator) newObj; - - Assertions.assertEquals(1, newOp.getProjectList().size()); - Assertions.assertEquals("action", newOp.getProjectList().get(0).getName()); - } } diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 046bb5679e1..6bb03528bbe 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -21,13 +21,16 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprBooleanValue; +import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasource.DataSourceServiceImpl; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.execution.PaginatedQueryService; +import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.NamedExpression; +import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.legacy.SQLIntegTestCase; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.client.OpenSearchRestClient; @@ -100,10 +103,12 @@ public void onFailure(Exception e) { // act 1, asserts in firstResponder var t = new OpenSearchIndex(client, defaultSettings(), "test"); - LogicalPlan p = new LogicalPaginate(1, List.of(new LogicalProject( - new LogicalRelation("test", t), - List.of(new NamedExpression("count()", new LiteralExpression(ExprBooleanValue.of(true)))), - List.of() + LogicalPlan p = new LogicalPaginate(1, List.of( + new LogicalProject( + new LogicalRelation("test", t), List.of( + DSL.named("name", DSL.ref("name", ExprCoreType.STRING)), + DSL.named("age", DSL.ref("age", ExprCoreType.LONG))), + List.of() ))); var firstResponder = new TestResponder(); paginatedQueryService.executePlan(p, firstResponder); @@ -113,6 +118,8 @@ public void onFailure(Exception e) { PhysicalPlan plan = paginatedPlanCache.convertToPlan(firstResponder.getCursor().asString()); var secondResponder = new TestResponder(); paginatedQueryService.executePlan(plan, secondResponder); + + // act 3: confirm that there's no cursor. } private Settings defaultSettings() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index ce990780c1c..4027845be61 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -50,6 +50,10 @@ OpenSearchResponse search(Function searchAction, */ OpenSearchExprValueFactory getExprValueFactory(); + default String toCursor() { + return ""; + } + /** * OpenSearch Index Name. * Indices are seperated by ",". diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 5f4e8a77cf6..456739cef0e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -11,7 +11,6 @@ import java.util.function.Function; import lombok.EqualsAndHashCode; import lombok.Getter; -import lombok.RequiredArgsConstructor; import lombok.Setter; import lombok.ToString; import org.opensearch.action.search.SearchRequest; @@ -146,11 +145,13 @@ public void reset() { * Convert a scroll request to string that can be included in a cursor. * @return a string representing the scroll request. */ + @Override public String toCursor() { if (isScrollStarted()) { + // TODO: probably should serialize exprValueFactory here as well. return scrollId; } else { - return "TESTING ONLY"; + return ""; } } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java new file mode 100644 index 00000000000..e07404bd18f --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java @@ -0,0 +1,59 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import java.util.function.Consumer; +import java.util.function.Function; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.ToString; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; + +public class ContinueScrollRequest implements OpenSearchRequest { + final String initialScrollId; + + // ScrollId that OpenSearch returns after search. + String responseScrollId; + + @EqualsAndHashCode.Exclude + @ToString.Exclude + @Getter + private final OpenSearchExprValueFactory exprValueFactory; + + public ContinueScrollRequest(String scrollId, OpenSearchExprValueFactory exprValueFactory) { + this.initialScrollId = scrollId; + this.exprValueFactory = exprValueFactory; + } + + @Override + public OpenSearchResponse search(Function searchAction, + Function scrollAction) { + SearchResponse openSearchResponse; + + openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)); + responseScrollId = openSearchResponse.getScrollId(); + + return new OpenSearchResponse(openSearchResponse, exprValueFactory); + } + + @Override + public void clean(Consumer cleanAction) { + cleanAction.accept(responseScrollId); + } + + @Override + public SearchSourceBuilder getSourceBuilder() { + throw new UnsupportedOperationException( + "SearchSourceBuilder is unavailable for ContinueScrollRequest"); + } + +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java similarity index 89% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java index 9e15fd12f67..6a453a95265 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java @@ -11,6 +11,7 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.index.query.QueryBuilder; import org.opensearch.search.aggregations.AggregationBuilder; @@ -25,9 +26,9 @@ import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -public class OpenSearchPagedRequestBuilder { - +public class InitialPageRequestBuilder implements PagedRequestBuilder { + @Getter private final OpenSearchRequest.IndexName indexName; private final SearchSourceBuilder sourceBuilder; private final OpenSearchExprValueFactory exprValueFactory; @@ -39,8 +40,8 @@ public class OpenSearchPagedRequestBuilder { * @param settings other settings * @param exprValueFactory value factory */ - public OpenSearchPagedRequestBuilder(OpenSearchRequest.IndexName indexName, Settings settings, - OpenSearchExprValueFactory exprValueFactory) { + public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, Settings settings, + OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; this.sourceBuilder = new SearchSourceBuilder(); this.exprValueFactory = exprValueFactory; @@ -50,12 +51,13 @@ public OpenSearchPagedRequestBuilder(OpenSearchRequest.IndexName indexName, Sett sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); } + @Override public OpenSearchScrollRequest build() { return new OpenSearchScrollRequest(indexName, sourceBuilder, exprValueFactory); } public void pushDown(QueryBuilder query) { - throw new RuntimeException(); + throw new UnsupportedOperationException("pushdown of a query is not supported"); } public void pushDownAggregation( diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index e293c8b0e78..b109fbd8360 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -129,7 +129,7 @@ public TableScanBuilder createScanBuilder() { @Override public TableScanBuilder createPagedScanBuilder() { - var requestBuilder = new OpenSearchPagedRequestBuilder(indexName, + var requestBuilder = new InitialPageRequestBuilder(indexName, settings, new OpenSearchExprValueFactory(getFieldTypes())); var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); return new OpenSearchPagedScanBuilder(indexScan); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java index 9ff5cfa1da4..111a7bec944 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java @@ -10,6 +10,7 @@ import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.storage.TableScanOperator; @@ -18,14 +19,14 @@ @ToString(onlyExplicitlyIncluded = true) public class OpenSearchPagedIndexScan extends TableScanOperator { private final OpenSearchClient client; - private final OpenSearchPagedRequestBuilder requestBuilder; + private final PagedRequestBuilder requestBuilder; @EqualsAndHashCode.Include @ToString.Include - private OpenSearchScrollRequest request; + private OpenSearchRequest request; private Iterator iterator; public OpenSearchPagedIndexScan(OpenSearchClient client, - OpenSearchPagedRequestBuilder requestBuilder) { + PagedRequestBuilder requestBuilder) { this.client = client; this.requestBuilder = requestBuilder; } @@ -64,6 +65,8 @@ public void close() { @Override public String toCursor() { - return createSection("OpenSearchPagedIndexScan", request.toCursor()); + // TODO this assumes exactly one index is scanned. + var indexName = requestBuilder.getIndexName().getIndexNames()[0]; + return createSection("OpenSearchPagedIndexScan", indexName, request.toCursor()); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java index 4a3393abc94..f34ed79c546 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java @@ -12,9 +12,12 @@ import org.opensearch.sql.DataSourceSchemaName; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.storage.system.OpenSearchSystemIndex; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.TableScanOperator; /** OpenSearch storage engine implementation. */ @RequiredArgsConstructor @@ -33,4 +36,17 @@ public Table getTable(DataSourceSchemaName dataSourceSchemaName, String name) { return new OpenSearchIndex(client, settings, name); } } + + @Override + public TableScanOperator getTableScan(String scanAsString) { + // TODO extract indexName and scrollId from scanAsString + String indexName =""; + String scrollId = ""; + var index = new OpenSearchIndex(client, settings, indexName); + var requestBuilder = new SubsequentPageRequestBuilder( + new OpenSearchRequest.IndexName(indexName), + scrollId, + new OpenSearchExprValueFactory(index.getFieldTypes())); + return new OpenSearchPagedIndexScan(client, requestBuilder); + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java new file mode 100644 index 00000000000..c138e327550 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java @@ -0,0 +1,14 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; + +public interface PagedRequestBuilder { + OpenSearchRequest build(); + OpenSearchRequest.IndexName getIndexName(); +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java new file mode 100644 index 00000000000..dca63e8e99f --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage; + +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; + +public class SubsequentPageRequestBuilder implements PagedRequestBuilder { + private OpenSearchRequest.IndexName indexName; + final String scrollId; + private OpenSearchExprValueFactory exprValueFactory; + + public SubsequentPageRequestBuilder(OpenSearchRequest.IndexName indexName, String scanAsString, + OpenSearchExprValueFactory exprValueFactory) { + this.indexName = indexName; + scrollId = scanAsString; + this.exprValueFactory = exprValueFactory; + } + + @Override + public OpenSearchRequest build() { + return new ContinueScrollRequest(scrollId, exprValueFactory); + } + + @Override + public OpenSearchRequest.IndexName getIndexName() { + return indexName; + } +} From f3cade65a269fed4d95b7355bc9ea68e91e6926b Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 8 Feb 2023 12:00:25 -0800 Subject: [PATCH 05/46] Move `ExpressionSerializer` from `opensearch` to `core`. Signed-off-by: Yury-Fridlyand --- .../expression}/serialization/DefaultExpressionSerializer.java | 2 +- .../sql/expression}/serialization/ExpressionSerializer.java | 2 +- .../storage/scan/OpenSearchIndexScanAggregationBuilder.java | 2 +- .../storage/scan/OpenSearchIndexScanQueryBuilder.java | 2 +- .../sql/opensearch/storage/script/ExpressionScriptEngine.java | 2 +- .../storage/script/aggregation/AggregationQueryBuilder.java | 2 +- .../script/aggregation/dsl/AggregationBuilderHelper.java | 2 +- .../script/aggregation/dsl/BucketAggregationBuilder.java | 2 +- .../script/aggregation/dsl/MetricAggregationBuilder.java | 2 +- .../opensearch/storage/script/filter/FilterQueryBuilder.java | 2 +- .../opensearch/storage/script/ExpressionScriptEngineTest.java | 2 +- .../storage/script/aggregation/AggregationQueryBuilderTest.java | 2 +- .../script/aggregation/dsl/BucketAggregationBuilderTest.java | 2 +- .../script/aggregation/dsl/MetricAggregationBuilderTest.java | 2 +- .../storage/script/filter/FilterQueryBuilderTest.java | 2 +- .../storage/serialization/DefaultExpressionSerializerTest.java | 2 ++ plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java | 2 +- 17 files changed, 18 insertions(+), 16 deletions(-) rename {opensearch/src/main/java/org/opensearch/sql/opensearch/storage => core/src/main/java/org/opensearch/sql/expression}/serialization/DefaultExpressionSerializer.java (95%) rename {opensearch/src/main/java/org/opensearch/sql/opensearch/storage => core/src/main/java/org/opensearch/sql/expression}/serialization/ExpressionSerializer.java (90%) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializer.java b/core/src/main/java/org/opensearch/sql/expression/serialization/DefaultExpressionSerializer.java similarity index 95% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializer.java rename to core/src/main/java/org/opensearch/sql/expression/serialization/DefaultExpressionSerializer.java index dc67da9de5d..33c22b2ea5d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializer.java +++ b/core/src/main/java/org/opensearch/sql/expression/serialization/DefaultExpressionSerializer.java @@ -4,7 +4,7 @@ */ -package org.opensearch.sql.opensearch.storage.serialization; +package org.opensearch.sql.expression.serialization; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/ExpressionSerializer.java b/core/src/main/java/org/opensearch/sql/expression/serialization/ExpressionSerializer.java similarity index 90% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/ExpressionSerializer.java rename to core/src/main/java/org/opensearch/sql/expression/serialization/ExpressionSerializer.java index b7caeb30f81..f96921e29c3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serialization/ExpressionSerializer.java +++ b/core/src/main/java/org/opensearch/sql/expression/serialization/ExpressionSerializer.java @@ -4,7 +4,7 @@ */ -package org.opensearch.sql.opensearch.storage.serialization; +package org.opensearch.sql.expression.serialization; import org.opensearch.sql.expression.Expression; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index e52fc566cd6..719ef52b79c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -18,7 +18,7 @@ import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; -import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalSort; import org.opensearch.sql.storage.TableScanOperator; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index 370036cdeee..a67ceeae9ae 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -23,7 +23,7 @@ import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; -import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java index 855aae645d2..a48da591806 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java @@ -18,7 +18,7 @@ import org.opensearch.sql.expression.Expression; import org.opensearch.sql.opensearch.storage.script.aggregation.ExpressionAggregationScriptFactory; import org.opensearch.sql.opensearch.storage.script.filter.ExpressionFilterScriptFactory; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Custom expression script engine that supports using core engine expression code in DSL diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java index ae3239eea04..f8ede598c1f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java @@ -36,7 +36,7 @@ import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.storage.script.aggregation.dsl.BucketAggregationBuilder; import org.opensearch.sql.opensearch.storage.script.aggregation.dsl.MetricAggregationBuilder; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Build the AggregationBuilder from the list of {@link NamedAggregator} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java index 001d7af9700..9cba3622a0d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java @@ -18,7 +18,7 @@ import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.opensearch.storage.script.ScriptUtils; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Abstract Aggregation Builder. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java index 1a6a82be966..5e910e73c3a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java @@ -24,7 +24,7 @@ import org.opensearch.sql.ast.expression.SpanUnit; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.span.SpanExpression; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Bucket Aggregation Builder. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java index 5e7d34abce0..d00a4eaf89a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java @@ -31,7 +31,7 @@ import org.opensearch.sql.opensearch.response.agg.StatsParser; import org.opensearch.sql.opensearch.response.agg.TopHitsParser; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Build the Metric Aggregation and List of {@link MetricParser} from {@link NamedAggregator}. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java index 5f36954d4a7..ff206a6a1e2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java @@ -38,7 +38,7 @@ import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.QueryStringQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.SimpleQueryStringQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.WildcardQuery; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @RequiredArgsConstructor public class FilterQueryBuilder extends ExpressionNodeVisitor { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java index 3d497c2f5b7..b106f396fb0 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java @@ -28,7 +28,7 @@ import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; import org.opensearch.sql.opensearch.storage.script.filter.ExpressionFilterScriptFactory; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilderTest.java index b62d5452068..00b6046f961 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilderTest.java @@ -52,7 +52,7 @@ import org.opensearch.sql.expression.aggregation.AvgAggregator; import org.opensearch.sql.expression.aggregation.CountAggregator; import org.opensearch.sql.expression.aggregation.NamedAggregator; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilderTest.java index 11dd0f849bd..45787ecb567 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilderTest.java @@ -46,7 +46,7 @@ import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java index 94f152f9132..d8e81026b68 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java @@ -43,7 +43,7 @@ import org.opensearch.sql.expression.aggregation.SumAggregator; import org.opensearch.sql.expression.aggregation.TakeAggregator; import org.opensearch.sql.expression.function.FunctionName; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilderTest.java index 2ad1f59d392..0771efd8283 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilderTest.java @@ -54,7 +54,7 @@ import org.opensearch.sql.expression.FunctionExpression; import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java index 72a319dbfe6..53a89d5421a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java @@ -21,6 +21,8 @@ import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.ExpressionNodeVisitor; import org.opensearch.sql.expression.env.Environment; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class DefaultExpressionSerializerTest { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 2c401af352f..08ad3b963a1 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -68,7 +68,7 @@ import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory; import org.opensearch.sql.opensearch.storage.script.ExpressionScriptEngine; -import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.plugin.config.OpenSearchPluginModule; import org.opensearch.sql.plugin.datasource.DataSourceSettings; import org.opensearch.sql.plugin.rest.RestPPLQueryAction; From db342b829d64cddbd57066809969fae0c9d52e63 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 8 Feb 2023 12:02:15 -0800 Subject: [PATCH 06/46] Rename `Cursor` `asString` to `toString`. Signed-off-by: Yury-Fridlyand --- .../java/org/opensearch/sql/opensearch/executor/Cursor.java | 2 +- .../java/org/opensearch/sql/sql/StandalonePaginationIT.java | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java index 6751ca107c6..73289d9066c 100644 --- a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java +++ b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java @@ -23,7 +23,7 @@ public Cursor(byte[] raw) { this.raw = raw; } - public String asString() { + public String toString() { return new String(raw); } } diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 6bb03528bbe..4f01955fe93 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -20,7 +20,6 @@ import org.opensearch.common.inject.ModulesBuilder; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasource.DataSourceServiceImpl; @@ -28,9 +27,6 @@ import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.execution.PaginatedQueryService; import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.expression.LiteralExpression; -import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.legacy.SQLIntegTestCase; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.client.OpenSearchRestClient; @@ -115,7 +111,7 @@ public void onFailure(Exception e) { // act 2, asserts in secondResponder - PhysicalPlan plan = paginatedPlanCache.convertToPlan(firstResponder.getCursor().asString()); + PhysicalPlan plan = paginatedPlanCache.convertToPlan(firstResponder.getCursor().toString()); var secondResponder = new TestResponder(); paginatedQueryService.executePlan(plan, secondResponder); From c8f0935fdfb91ed25baec741943f971a786f84ee Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 8 Feb 2023 12:03:05 -0800 Subject: [PATCH 07/46] Disable scroll cleaning. Signed-off-by: Yury-Fridlyand --- .../sql/opensearch/client/OpenSearchNodeClient.java | 2 +- .../sql/opensearch/client/OpenSearchRestClient.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java index 8818c394a17..30584581e31 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java @@ -172,7 +172,7 @@ public Map meta() { @Override public void cleanup(OpenSearchRequest request) { - request.clean(scrollId -> client.prepareClearScroll().addScrollId(scrollId).get()); + request.clean(scrollId -> {}/* client.prepareClearScroll().addScrollId(scrollId).get() */); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java index d9f9dbbe5d5..096899516f9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java @@ -174,7 +174,7 @@ public Map meta() { @Override public void cleanup(OpenSearchRequest request) { - request.clean(scrollId -> { + request.clean(scrollId -> {/* try { ClearScrollRequest clearRequest = new ClearScrollRequest(); clearRequest.addScrollId(scrollId); @@ -182,7 +182,7 @@ public void cleanup(OpenSearchRequest request) { } catch (IOException e) { throw new IllegalStateException( "Failed to clean up resources for search request " + request, e); - } + }*/ }); } From fffc36de8144a36a3e1a83bb3df1e827241e85cb Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 8 Feb 2023 12:04:34 -0800 Subject: [PATCH 08/46] Add full cursor serialization and deserialization. Signed-off-by: Yury-Fridlyand --- .../sql/executor/PaginatedPlanCache.java | 72 ++++++++++++++++--- .../sql/planner/physical/ProjectOperator.java | 17 +++-- .../opensearch/sql/storage/StorageEngine.java | 2 +- .../storage/OpenSearchStorageEngine.java | 5 +- 4 files changed, 77 insertions(+), 19 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 25c94080604..dc81d03aa4d 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -7,6 +7,7 @@ import java.io.IOException; import java.io.ObjectInputStream; +import java.util.ArrayList; import java.util.List; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -15,6 +16,7 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -58,19 +60,69 @@ public Cursor convertToCursor(PhysicalPlan plan) { */ public PhysicalPlan convertToPlan(String cursor) { if (cursor.startsWith(CURSOR_PREFIX)) { - String expression = cursor.substring(CURSOR_PREFIX.length()); + try { + String expression = cursor.substring(CURSOR_PREFIX.length()); - // TODO Parse expression and initialize variables below. - // storageEngine needs to create the TableScanOperator. - int pageSize = -1; - int currentPageIndex = -1; - List projectList = List.of(); - String scanAsString = ""; - TableScanOperator scan = storageEngine.getTableScan(scanAsString); + // TODO Parse expression and initialize variables below. + // storageEngine needs to create the TableScanOperator. - return new PaginateOperator(new ProjectOperator(scan, projectList, List.of()), - pageSize, currentPageIndex); + // TODO Parse with ANTLR or serialize as JSON/XML + if (!expression.startsWith("(Paginate,")) { + throw new UnsupportedOperationException("Unsupported cursor"); + } + expression = expression.substring(expression.indexOf(',') + 1); + int currentPageIndex = Integer.parseInt(expression, 0, expression.indexOf(','), 10); + expression = expression.substring(expression.indexOf(',') + 1); + int pageSize = Integer.parseInt(expression, 0, expression.indexOf(','), 10); + + expression = expression.substring(expression.indexOf(',') + 1); + if (!expression.startsWith("(Project,")) { + throw new UnsupportedOperationException("Unsupported cursor"); + } + expression = expression.substring(expression.indexOf(',') + 1); + if (!expression.startsWith("(namedParseExpressions,")) { + throw new UnsupportedOperationException("Unsupported cursor"); + } + expression = expression.substring(expression.indexOf(',') + 1); + var serializer = new DefaultExpressionSerializer(); + // TODO parse npe + List namedParseExpressions = List.of(); + + expression = expression.substring(expression.indexOf(',') + 1); + List projectList = new ArrayList<>(); + if (!expression.startsWith("(projectList,")) { + throw new UnsupportedOperationException("Unsupported cursor"); + } + expression = expression.substring(expression.indexOf(',') + 1); + while (expression.startsWith("(named,")) { + expression = expression.substring(expression.indexOf(',') + 1); + var name = expression.substring(0, expression.indexOf(',')); + expression = expression.substring(expression.indexOf(',') + 1); + var alias = expression.substring(0, expression.indexOf(',')); + if (alias.isEmpty()) { + alias = null; + } + expression = expression.substring(expression.indexOf(',') + 1); + projectList.add(new NamedExpression(name, + serializer.deserialize(expression.substring(0, expression.indexOf(')'))), alias)); + expression = expression.substring(expression.indexOf(',') + 1); + } + + if (!expression.startsWith("(OpenSearchPagedIndexScan,")) { + throw new UnsupportedOperationException("Unsupported cursor"); + } + expression = expression.substring(expression.indexOf(',') + 1); + var indexName = expression.substring(0, expression.indexOf(',')); + expression = expression.substring(expression.indexOf(',') + 1); + var scrollId = expression.substring(0, expression.indexOf(')')); + TableScanOperator scan = storageEngine.getTableScan(indexName, scrollId); + + return new PaginateOperator(new ProjectOperator(scan, projectList, namedParseExpressions), + pageSize, currentPageIndex); + } catch (Exception e) { + throw new UnsupportedOperationException("Unsupported cursor", e); + } } else { throw new UnsupportedOperationException("Unsupported cursor"); } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 9a3542ad1fc..f07b43768fd 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -22,6 +22,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; /** * Project the fields specified in {@link ProjectOperator#projectList} from input. @@ -98,9 +99,17 @@ public ExecutionEngine.Schema schema() { @Override public String toCursor() { String child = getChild().get(0).toCursor(); - String namedExpressions = "TODO"; - // TODO serialize named expressions. - // Skipping parsedExpressions for now. - return createSection("Project", namedExpressions, child); + var serializer = new DefaultExpressionSerializer(); + String projects = createSection("projectList", + projectList.stream().map(ne -> createSection("named", + ne.getName(), ne.getAlias() == null ? "" : ne.getAlias(), serializer.serialize(ne.getDelegated()) + )) + .toArray(String[]::new)); + String namedExpressions = createSection("namedParseExpressions", + namedParseExpressions.stream().map(ne -> createSection("named", + ne.getName(), ne.getAlias() == null ? "" : ne.getAlias(), serializer.serialize(ne.getDelegated()) + )) + .toArray(String[]::new)); + return createSection("Project", namedExpressions, projects, child); } } diff --git a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java index e854477b72b..18e9e92886c 100644 --- a/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java +++ b/core/src/main/java/org/opensearch/sql/storage/StorageEngine.java @@ -32,7 +32,7 @@ default Collection getFunctions() { return Collections.emptyList(); } - default TableScanOperator getTableScan(String scanAsString) { + default TableScanOperator getTableScan(String indexName, String scrollId) { String error = String.format("%s.getTableScan needs to be implemented", getClass()); throw new UnsupportedOperationException(error); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java index f34ed79c546..0b0e231760d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java @@ -38,10 +38,7 @@ public Table getTable(DataSourceSchemaName dataSourceSchemaName, String name) { } @Override - public TableScanOperator getTableScan(String scanAsString) { - // TODO extract indexName and scrollId from scanAsString - String indexName =""; - String scrollId = ""; + public TableScanOperator getTableScan(String indexName, String scrollId) { var index = new OpenSearchIndex(client, settings, indexName); var requestBuilder = new SubsequentPageRequestBuilder( new OpenSearchRequest.IndexName(indexName), From d84497749b7917e29e4a8251945b173519cbe3d7 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 8 Feb 2023 12:05:20 -0800 Subject: [PATCH 09/46] Misc fixes. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java | 4 ++-- .../sql/opensearch/request/OpenSearchScrollRequest.java | 3 ++- .../sql/opensearch/response/OpenSearchResponse.java | 1 + .../sql/opensearch/storage/ContinueScrollRequest.java | 2 +- .../sql/opensearch/storage/InitialPageRequestBuilder.java | 2 +- .../sql/opensearch/storage/OpenSearchPagedIndexScan.java | 2 +- 6 files changed, 8 insertions(+), 6 deletions(-) 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 9594d5a3b51..bd9387a68e0 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 @@ -119,7 +119,7 @@ private ResponseListener fallBackListener( return new ResponseListener() { @Override public void onResponse(T response) { - LOG.error("[{}] Request is handled by new SQL query engine", + LOG.info("[{}] Request is handled by new SQL query engine", QueryContext.getRequestId()); next.onResponse(response); } @@ -172,7 +172,7 @@ private ResponseListener createQueryResponseListener( @Override public void onResponse(QueryResponse response) { sendResponse(channel, OK, - formatter.format(new QueryResult(response.getSchema(), response.getResults()))); + formatter.format(new QueryResult(response.getSchema(), response.getResults(), response.getCursor()))); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 456739cef0e..77553888426 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -33,7 +33,7 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { /** Default scroll context timeout in minutes. */ - public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1L); + public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(100L); /** * {@link OpenSearchRequest.IndexName}. @@ -86,6 +86,7 @@ public OpenSearchResponse search(Function searchA } else { openSearchResponse = searchAction.apply(searchRequest()); } + setScrollId(openSearchResponse.getScrollId()); return new OpenSearchResponse(openSearchResponse, exprValueFactory); 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..c2042d7c0ea 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 @@ -71,6 +71,7 @@ public OpenSearchResponse(SearchHits hits, OpenSearchExprValueFactory exprValueF */ public boolean isEmpty() { return (hits.getHits() == null) || (hits.getHits().length == 0) && aggregations == null; + // TODO TBD ^ ^ } public boolean isAggregationResponse() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java index e07404bd18f..3c431a15ff2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java @@ -39,7 +39,7 @@ public OpenSearchResponse search(Function searchA Function scrollAction) { SearchResponse openSearchResponse; - openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)); + openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)); responseScrollId = openSearchResponse.getScrollId(); return new OpenSearchResponse(openSearchResponse, exprValueFactory); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java index 6a453a95265..e36e0df529c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java @@ -45,7 +45,7 @@ public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, Settings this.indexName = indexName; this.sourceBuilder = new SearchSourceBuilder(); this.exprValueFactory = exprValueFactory; - this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT);//TODO fetch_size sourceBuilder.from(0); sourceBuilder.size(querySize); sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java index 111a7bec944..ea09502b0ef 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java @@ -53,7 +53,7 @@ public void open() { OpenSearchResponse response = client.search(request); if (!response.isEmpty()) { iterator = response.iterator(); - } + } // TODO else - last page is empty - } @Override From 333432ebc071483394e9755c295d252e67f82459 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 9 Feb 2023 20:41:38 -0800 Subject: [PATCH 10/46] Further work on pagination. * Added push down page size from `LogicalPaginate` to `LogicalRelation`. * Improved cursor encoding and decoding. * Added cursor compression. * Fixed issuing `SearchScrollRequest`. * Fixed returning last empty page. * Minor code grooming/commenting. Signed-off-by: Yury-Fridlyand --- .../sql/executor/PaginatedPlanCache.java | 124 +++++++++++------- .../sql/planner/DefaultImplementor.java | 3 - .../sql/planner/PaginateOperator.java | 5 +- .../sql/planner/logical/LogicalRelation.java | 6 + .../optimizer/LogicalPlanOptimizer.java | 3 + .../planner/optimizer/pattern/Patterns.java | 11 ++ .../CreatePagingTableScanBuilder.java | 6 +- .../planner/optimizer/rule/PushPageSize.java | 62 +++++++++ .../sql/planner/physical/PhysicalPlan.java | 2 +- .../sql/planner/physical/ProjectOperator.java | 13 +- .../org/opensearch/sql/storage/Table.java | 2 +- .../opensearch/request/OpenSearchRequest.java | 2 +- .../storage/ContinueScrollRequest.java | 17 ++- .../storage/InitialPageRequestBuilder.java | 12 +- .../opensearch/storage/OpenSearchIndex.java | 4 +- .../storage/OpenSearchPagedIndexScan.java | 11 +- 16 files changed, 205 insertions(+), 78 deletions(-) rename core/src/main/java/org/opensearch/sql/planner/optimizer/{ => rule}/CreatePagingTableScanBuilder.java (86%) create mode 100644 core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index dc81d03aa4d..e1172a71d62 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -5,27 +5,25 @@ package org.opensearch.sql.executor; -import java.io.IOException; -import java.io.ObjectInputStream; +import com.google.common.hash.HashCode; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.util.ArrayList; import java.util.List; -import java.util.stream.Collectors; -import java.util.stream.Stream; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; import lombok.Data; import lombok.RequiredArgsConstructor; +import lombok.SneakyThrows; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.storage.StorageEngine; -import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.TableScanOperator; -import org.opensearch.sql.storage.read.TableScanBuilder; @RequiredArgsConstructor public class PaginatedPlanCache { @@ -39,7 +37,7 @@ public boolean canConvertToCursor(UnresolvedPlan plan) { @RequiredArgsConstructor @Data - static class SeriazationContext { + static class SerializationContext { private final PaginatedPlanCache cache; } @@ -48,74 +46,104 @@ static class SeriazationContext { */ public Cursor convertToCursor(PhysicalPlan plan) { if (plan instanceof PaginateOperator) { - var raw = CURSOR_PREFIX + plan.toCursor(); + var cursor = plan.toCursor(); + if (cursor == null || cursor.isEmpty()) { + return Cursor.None; + } + var raw = CURSOR_PREFIX + compress(cursor); return new Cursor(raw.getBytes()); } else { return Cursor.None; } } + @SneakyThrows + public static String compress(String str) { + if (str == null || str.length() == 0) { + return null; + } + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + GZIPOutputStream gzip = new GZIPOutputStream(out); + gzip.write(str.getBytes()); + gzip.close(); + return HashCode.fromBytes(out.toByteArray()).toString(); + } + + @SneakyThrows + public static String decompress(String input) { + if (input == null || input.length() == 0) { + return null; + } + GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream( + HashCode.fromString(input).asBytes())); + return new String(gzip.readAllBytes()); + } + + /** + * Parse `NamedExpression`s from cursor. + * @param listToFill List to fill with data. + * @param cursor Cursor to parse. + * @return Remaining part of the cursor. + */ + private String parseNamedExpressions(List listToFill, String cursor) { + var serializer = new DefaultExpressionSerializer(); + while (!cursor.startsWith(")") && !cursor.startsWith("(")) { + listToFill.add((NamedExpression) + serializer.deserialize(cursor.substring(0, + Math.min(cursor.indexOf(','), cursor.indexOf(')'))))); + cursor = cursor.substring(cursor.indexOf(',') + 1); + } + return cursor; + } + /** * Converts a cursor to a physical plan tree. */ public PhysicalPlan convertToPlan(String cursor) { if (cursor.startsWith(CURSOR_PREFIX)) { try { - String expression = cursor.substring(CURSOR_PREFIX.length()); - - // TODO Parse expression and initialize variables below. - // storageEngine needs to create the TableScanOperator. + cursor = cursor.substring(CURSOR_PREFIX.length()); + cursor = decompress(cursor); // TODO Parse with ANTLR or serialize as JSON/XML - if (!expression.startsWith("(Paginate,")) { + if (!cursor.startsWith("(Paginate,")) { throw new UnsupportedOperationException("Unsupported cursor"); } - expression = expression.substring(expression.indexOf(',') + 1); - int currentPageIndex = Integer.parseInt(expression, 0, expression.indexOf(','), 10); + cursor = cursor.substring(cursor.indexOf(',') + 1); + int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); - expression = expression.substring(expression.indexOf(',') + 1); - int pageSize = Integer.parseInt(expression, 0, expression.indexOf(','), 10); + cursor = cursor.substring(cursor.indexOf(',') + 1); + int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); - expression = expression.substring(expression.indexOf(',') + 1); - if (!expression.startsWith("(Project,")) { + cursor = cursor.substring(cursor.indexOf(',') + 1); + if (!cursor.startsWith("(Project,")) { throw new UnsupportedOperationException("Unsupported cursor"); } - expression = expression.substring(expression.indexOf(',') + 1); - if (!expression.startsWith("(namedParseExpressions,")) { + cursor = cursor.substring(cursor.indexOf(',') + 1); + if (!cursor.startsWith("(namedParseExpressions,")) { throw new UnsupportedOperationException("Unsupported cursor"); } - expression = expression.substring(expression.indexOf(',') + 1); - var serializer = new DefaultExpressionSerializer(); - // TODO parse npe - List namedParseExpressions = List.of(); - expression = expression.substring(expression.indexOf(',') + 1); + cursor = cursor.substring(cursor.indexOf(',') + 1); + List namedParseExpressions = new ArrayList<>(); + cursor = parseNamedExpressions(namedParseExpressions, cursor); + + cursor = cursor.substring(cursor.indexOf(',') + 1); List projectList = new ArrayList<>(); - if (!expression.startsWith("(projectList,")) { + if (!cursor.startsWith("(projectList,")) { throw new UnsupportedOperationException("Unsupported cursor"); } - expression = expression.substring(expression.indexOf(',') + 1); - while (expression.startsWith("(named,")) { - expression = expression.substring(expression.indexOf(',') + 1); - var name = expression.substring(0, expression.indexOf(',')); - expression = expression.substring(expression.indexOf(',') + 1); - var alias = expression.substring(0, expression.indexOf(',')); - if (alias.isEmpty()) { - alias = null; - } - expression = expression.substring(expression.indexOf(',') + 1); - projectList.add(new NamedExpression(name, - serializer.deserialize(expression.substring(0, expression.indexOf(')'))), alias)); - expression = expression.substring(expression.indexOf(',') + 1); - } + cursor = cursor.substring(cursor.indexOf(',') + 1); + cursor = parseNamedExpressions(projectList, cursor); - if (!expression.startsWith("(OpenSearchPagedIndexScan,")) { + if (!cursor.startsWith("(OpenSearchPagedIndexScan,")) { throw new UnsupportedOperationException("Unsupported cursor"); } - expression = expression.substring(expression.indexOf(',') + 1); - var indexName = expression.substring(0, expression.indexOf(',')); - expression = expression.substring(expression.indexOf(',') + 1); - var scrollId = expression.substring(0, expression.indexOf(')')); + cursor = cursor.substring(cursor.indexOf(',') + 1); + var indexName = cursor.substring(0, cursor.indexOf(',')); + cursor = cursor.substring(cursor.indexOf(',') + 1); + var scrollId = cursor.substring(0, cursor.indexOf(')')); TableScanOperator scan = storageEngine.getTableScan(indexName, scrollId); return new PaginateOperator(new ProjectOperator(scan, projectList, namedParseExpressions), diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index 8bb1770d713..a02d908a0f1 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -126,7 +126,6 @@ public PhysicalPlan visitLimit(LogicalLimit node, C context) { return new LimitOperator(visitChild(node, context), node.getLimit(), node.getOffset()); } - @Override public PhysicalPlan visitPaginate(LogicalPaginate plan, C context) { return new PaginateOperator(visitChild(plan, context), plan.getPageSize()); @@ -148,10 +147,8 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { + "implementing and optimizing logical plan with relation involved"); } - protected PhysicalPlan visitChild(LogicalPlan node, C context) { // Logical operators visited here must have a single child return node.getChild().get(0).accept(this, context); } - } diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java index 227d4b4602c..a99beea8365 100644 --- a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java @@ -82,7 +82,8 @@ public String toCursor() { String child = getChild().get(0).toCursor(); var nextPage = getPageIndex() + 1; - return createSection("Paginate", Integer.toString(nextPage), - Integer.toString(getPageSize()), child); + return child == null || child.isEmpty() + ? null : createSection("Paginate", Integer.toString(nextPage), + Integer.toString(getPageSize()), child); } } diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java index a49c3d5cbe3..0ece74690e7 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java @@ -9,6 +9,7 @@ import com.google.common.collect.ImmutableList; import lombok.EqualsAndHashCode; import lombok.Getter; +import lombok.Setter; import lombok.ToString; import org.opensearch.sql.storage.Table; @@ -25,6 +26,10 @@ public class LogicalRelation extends LogicalPlan { @Getter private final Table table; + @Getter + @Setter + private Integer pageSize; + /** * Constructor of LogicalRelation. */ @@ -32,6 +37,7 @@ public LogicalRelation(String relationName, Table table) { super(ImmutableList.of()); this.relationName = relationName; this.table = table; + this.pageSize = null; } @Override diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java index 58a96c3efc6..13bcfabe74d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java @@ -13,8 +13,10 @@ import java.util.List; import java.util.stream.Collectors; import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder; import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter; import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort; +import org.opensearch.sql.planner.optimizer.rule.PushPageSize; import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown; import org.opensearch.sql.planner.optimizer.rule.write.CreateTableWriteBuilder; @@ -73,6 +75,7 @@ public static LogicalPlanOptimizer paginationCreate() { /* * Phase 2: Transformations that rely on data source push down capability */ + new PushPageSize(), new CreatePagingTableScanBuilder(), TableScanPushDown.PUSH_DOWN_FILTER, TableScanPushDown.PUSH_DOWN_AGGREGATION, diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java index 856d8df7ead..6e548975063 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java @@ -16,6 +16,7 @@ import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalRelation; @@ -112,6 +113,16 @@ public static Property table() { : Optional.empty()); } + /** + * Logical pagination with page size. + */ + public static Property pagination() { + return Property.optionalProperty("pagination", + plan -> plan instanceof LogicalPaginate + ? Optional.of(((LogicalPaginate) plan).getPageSize()) + : Optional.empty()); + } + /** * Logical write with table field. */ diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java similarity index 86% rename from core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java rename to core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java index bc97e373c22..6ef874b3766 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.planner.optimizer; +package org.opensearch.sql.planner.optimizer.rule; import static org.opensearch.sql.planner.optimizer.pattern.Patterns.table; @@ -14,6 +14,7 @@ import lombok.experimental.Accessors; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalRelation; +import org.opensearch.sql.planner.optimizer.Rule; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; @@ -38,7 +39,8 @@ public CreatePagingTableScanBuilder() { @Override public LogicalPlan apply(LogicalRelation plan, Captures captures) { - TableScanBuilder scanBuilder = captures.get(capture).createPagedScanBuilder(); + TableScanBuilder scanBuilder = captures.get(capture) + .createPagedScanBuilder(plan.getPageSize()); // TODO: Remove this after Prometheus refactored to new table scan builder too return (scanBuilder == null) ? plan : scanBuilder; } diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java new file mode 100644 index 00000000000..78f856535f3 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java @@ -0,0 +1,62 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner.optimizer.rule; + +import static org.opensearch.sql.planner.optimizer.pattern.Patterns.pagination; + +import com.facebook.presto.matching.Capture; +import com.facebook.presto.matching.Captures; +import com.facebook.presto.matching.Pattern; +import lombok.Getter; +import lombok.experimental.Accessors; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalRelation; +import org.opensearch.sql.planner.optimizer.Rule; + +import java.util.Objects; + +public class PushPageSize + implements Rule { + /** Capture the table inside matched logical paginate operator. */ + private final Capture capture; + + /** Pattern that matches logical paginate operator. */ + @Accessors(fluent = true) + @Getter + private final Pattern pattern; + + /** + * Constructor. + */ + public PushPageSize() { + this.capture = Capture.newCapture(); + this.pattern = Pattern.typeOf(LogicalPaginate.class) + .with(pagination().capturedAs(capture)); + } + + private LogicalRelation findLogicalRelation(LogicalPlan plan) { //TODO TBD multiple relations? + for (var subplan : plan.getChild()) { + if (subplan instanceof LogicalRelation) { + return (LogicalRelation) subplan; + } + var found = findLogicalRelation(subplan); + if (found != null) { + return found; + } + } + return null; + } + + @Override + public LogicalPlan apply(LogicalPaginate plan, Captures captures) { + var relation = findLogicalRelation(plan); + if (relation != null) { + relation.setPageSize(captures.get(capture)); + } + return plan; + } +} 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 8b7f1cad576..c52a477a51b 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 @@ -52,7 +52,7 @@ public ExecutionEngine.Schema schema() { } public String toCursor() { - throw new IllegalStateException(String.format("%s needs to implement ToCursor", + throw new IllegalStateException(String.format("%s needs to implement toCursor", this.getClass())); } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index f07b43768fd..c61b35e0cb6 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -99,17 +99,14 @@ public ExecutionEngine.Schema schema() { @Override public String toCursor() { String child = getChild().get(0).toCursor(); + if (child == null || child.isEmpty()) { + return null; + } var serializer = new DefaultExpressionSerializer(); String projects = createSection("projectList", - projectList.stream().map(ne -> createSection("named", - ne.getName(), ne.getAlias() == null ? "" : ne.getAlias(), serializer.serialize(ne.getDelegated()) - )) - .toArray(String[]::new)); + projectList.stream().map(serializer::serialize).toArray(String[]::new)); String namedExpressions = createSection("namedParseExpressions", - namedParseExpressions.stream().map(ne -> createSection("named", - ne.getName(), ne.getAlias() == null ? "" : ne.getAlias(), serializer.serialize(ne.getDelegated()) - )) - .toArray(String[]::new)); + namedParseExpressions.stream().map(serializer::serialize).toArray(String[]::new)); return createSection("Project", namedExpressions, projects, child); } } diff --git a/core/src/main/java/org/opensearch/sql/storage/Table.java b/core/src/main/java/org/opensearch/sql/storage/Table.java index 8117e2cc307..a7f2b606ca9 100644 --- a/core/src/main/java/org/opensearch/sql/storage/Table.java +++ b/core/src/main/java/org/opensearch/sql/storage/Table.java @@ -93,7 +93,7 @@ default StreamingSource asStreamingSource() { throw new UnsupportedOperationException(); } - default TableScanBuilder createPagedScanBuilder() { + default TableScanBuilder createPagedScanBuilder(int pageSize) { var error = String.format("'%s' does not support pagination", getClass().toString()); throw new UnsupportedOperationException(error); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index 4027845be61..c5b6d60af36 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -56,7 +56,7 @@ default String toCursor() { /** * OpenSearch Index Name. - * Indices are seperated by ",". + * Indices are separated by ",". */ @EqualsAndHashCode class IndexName { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java index 3c431a15ff2..ddc76f6ba22 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java @@ -18,6 +18,8 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.DEFAULT_SCROLL_TIMEOUT; + public class ContinueScrollRequest implements OpenSearchRequest { final String initialScrollId; @@ -39,10 +41,15 @@ public OpenSearchResponse search(Function searchA Function scrollAction) { SearchResponse openSearchResponse; - openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)); - responseScrollId = openSearchResponse.getScrollId(); + openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId) + .scroll(DEFAULT_SCROLL_TIMEOUT)); - return new OpenSearchResponse(openSearchResponse, exprValueFactory); + // TODO if terminated_early - something went wrong, e.g. no scroll returned. + var response = new OpenSearchResponse(openSearchResponse, exprValueFactory); + if (!response.isEmpty()) { + responseScrollId = openSearchResponse.getScrollId(); + } // else - last empty page, we should ignore the scroll even if it is returned + return response; } @Override @@ -56,4 +63,8 @@ public SearchSourceBuilder getSourceBuilder() { "SearchSourceBuilder is unavailable for ContinueScrollRequest"); } + @Override + public String toCursor() { + return responseScrollId; + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java index e36e0df529c..f2e4fd08b4a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java @@ -40,15 +40,17 @@ public class InitialPageRequestBuilder implements PagedRequestBuilder { * @param settings other settings * @param exprValueFactory value factory */ - public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, Settings settings, + public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, + int pageSize, + Settings settings, OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; this.sourceBuilder = new SearchSourceBuilder(); this.exprValueFactory = exprValueFactory; - this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT);//TODO fetch_size - sourceBuilder.from(0); - sourceBuilder.size(querySize); - sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); + this.querySize = pageSize; + sourceBuilder.from(0) + .size(querySize) + .timeout(DEFAULT_QUERY_TIMEOUT); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index b109fbd8360..1366dc374bc 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -128,8 +128,8 @@ public TableScanBuilder createScanBuilder() { } @Override - public TableScanBuilder createPagedScanBuilder() { - var requestBuilder = new InitialPageRequestBuilder(indexName, + public TableScanBuilder createPagedScanBuilder(int pageSize) { + var requestBuilder = new InitialPageRequestBuilder(indexName, pageSize, settings, new OpenSearchExprValueFactory(getFieldTypes())); var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); return new OpenSearchPagedScanBuilder(indexScan); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java index ea09502b0ef..646b414183d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java @@ -5,7 +5,10 @@ package org.opensearch.sql.opensearch.storage; +import java.util.Collections; import java.util.Iterator; +import java.util.List; + import lombok.EqualsAndHashCode; import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; @@ -53,7 +56,9 @@ public void open() { OpenSearchResponse response = client.search(request); if (!response.isEmpty()) { iterator = response.iterator(); - } // TODO else - last page is empty - + } else { + iterator = Collections.emptyIterator(); + } } @Override @@ -67,6 +72,8 @@ public void close() { public String toCursor() { // TODO this assumes exactly one index is scanned. var indexName = requestBuilder.getIndexName().getIndexNames()[0]; - return createSection("OpenSearchPagedIndexScan", indexName, request.toCursor()); + var cursor = request.toCursor(); + return cursor == null || cursor.isEmpty() + ? "" : createSection("OpenSearchPagedIndexScan", indexName, cursor); } } From 85c8825076c1c0fb260d58b4915c4e61f1a075a2 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Feb 2023 10:48:15 -0800 Subject: [PATCH 11/46] Pagination fix for empty indices. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/executor/PaginatedPlanCache.java | 6 ++++-- .../sql/opensearch/request/OpenSearchScrollRequest.java | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index e1172a71d62..73c058316c9 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -88,7 +88,10 @@ public static String decompress(String input) { */ private String parseNamedExpressions(List listToFill, String cursor) { var serializer = new DefaultExpressionSerializer(); - while (!cursor.startsWith(")") && !cursor.startsWith("(")) { + if (cursor.startsWith(")")) { //empty list + return cursor.substring(cursor.indexOf(',') + 1); + } + while (!cursor.startsWith("(")) { listToFill.add((NamedExpression) serializer.deserialize(cursor.substring(0, Math.min(cursor.indexOf(','), cursor.indexOf(')'))))); @@ -129,7 +132,6 @@ public PhysicalPlan convertToPlan(String cursor) { List namedParseExpressions = new ArrayList<>(); cursor = parseNamedExpressions(namedParseExpressions, cursor); - cursor = cursor.substring(cursor.indexOf(',') + 1); List projectList = new ArrayList<>(); if (!cursor.startsWith("(projectList,")) { throw new UnsupportedOperationException("Unsupported cursor"); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 77553888426..8eef94c8375 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -87,9 +87,11 @@ public OpenSearchResponse search(Function searchA openSearchResponse = searchAction.apply(searchRequest()); } - setScrollId(openSearchResponse.getScrollId()); - - return new OpenSearchResponse(openSearchResponse, exprValueFactory); + var response = new OpenSearchResponse(openSearchResponse, exprValueFactory); + if (!response.isEmpty()) { + setScrollId(openSearchResponse.getScrollId()); + } // else - last empty page, we should ignore the scroll even if it is returned + return response; } @Override From 484a8fea9c07ddee2184f809323dd6579d4f23d9 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Feb 2023 14:06:53 -0800 Subject: [PATCH 12/46] Fix error reporting on wrong cursor. Signed-off-by: Yury-Fridlyand --- .../sql/executor/execution/ContinuePaginatedPlan.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java index 0115e6de731..adefae778e1 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java @@ -40,8 +40,12 @@ public ContinuePaginatedPlan(QueryId queryId, String cursor, PaginatedQueryServi @Override public void execute() { - PhysicalPlan plan = paginatedPlanCache.convertToPlan(cursor); - queryService.executePlan(plan, queryResponseListener); + try { + PhysicalPlan plan = paginatedPlanCache.convertToPlan(cursor); + queryService.executePlan(plan, queryResponseListener); + } catch (Exception e) { + queryResponseListener.onFailure(e); + } } @Override From cccce53c2c5cc41214de20f45e3767e7a172f2fb Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 14 Feb 2023 12:23:27 -0800 Subject: [PATCH 13/46] Minor comments and error reporting improvement. Signed-off-by: Yury-Fridlyand --- .../sql/executor/CanPaginateVisitor.java | 47 ++++++++++++++++++- .../sql/planner/physical/PhysicalPlan.java | 7 ++- .../sql/legacy/plugin/RestSqlAction.java | 2 +- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index e89a04216ec..c85bbf0cc8c 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -8,12 +8,18 @@ import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.Relation; +import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.ast.tree.Values; + +import java.util.concurrent.atomic.AtomicBoolean; /** * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache. - * If plan.accept(new CanpaginateVisitor(...)) returns true, + * If plan.accept(new CanPaginateVisitor(...)) returns true, * then PaginatedPlanCache.convertToCursor will succeed. * Otherwise, it will fail. * Currently, the conditions are: @@ -38,6 +44,42 @@ public Boolean visitRelation(Relation node, Object context) { return Boolean.TRUE; } + private Boolean canPaginate(Node node, Object context) { + AtomicBoolean result = new AtomicBoolean(true); + node.getChild().forEach(n -> result.set(result.get() && n.accept(this, context))); + return result.get(); + } + + /* + For queries without `FROM` clause. + Required to overload `toCursor` function in `ValuesOperator` and modify cursor parsing. + @Override + public Boolean visitValues(Values node, Object context) { + return canPaginate(node, context); + } + + For queries with LIMIT clause: + Required to overload `toCursor` function in `LimitOperator` and modify cursor parsing. + @Override + public Boolean visitLimit(Limit node, Object context) { + return canPaginate(node, context); + } + + For queries with ORDER BY clause: + Required to overload `toCursor` function in `SortOperator` and modify cursor parsing. + @Override + public Boolean visitSort(Sort node, Object context) { + return canPaginate(node, context); + } + + For queries with WHERE clause: + Required to overload `toCursor` function in `FilterOperator` and modify cursor parsing. + @Override + public Boolean visitFilter(Filter node, Object context) { + return canPaginate(node, context); + } + */ + @Override public Boolean visitChildren(Node node, Object context) { return Boolean.FALSE; @@ -45,6 +87,9 @@ public Boolean visitChildren(Node node, Object context) { @Override public Boolean visitProject(Project node, Object context) { + // Allow queries with 'SELECT *' only. Those restriction could be removed, but consider + // in-memory aggregation performed by window function (see WindowOperator). + // SELECT max(age) OVER (PARTITION BY city) ... var projections = node.getProjectList(); if (projections.size() != 1) { return Boolean.FALSE; 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 c52a477a51b..5890b6f15f9 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 @@ -6,7 +6,6 @@ package org.opensearch.sql.planner.physical; -import java.io.Serializable; import java.util.Iterator; import java.util.List; import org.opensearch.sql.data.model.ExprValue; @@ -48,12 +47,12 @@ public void add(Split split) { public ExecutionEngine.Schema schema() { throw new IllegalStateException(String.format("[BUG] schema can been only applied to " - + "ProjectOperator, instead of %s", toString())); + + "ProjectOperator, instead of %s", this.getClass().getSimpleName())); } public String toCursor() { - throw new IllegalStateException(String.format("%s needs to implement toCursor", - this.getClass())); + throw new IllegalStateException(String.format("%s is not compatible with cursor feature", + this.getClass().getSimpleName())); } /** diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 6fdd855fd0e..67f5f3bb43e 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -157,7 +157,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli if (newSqlRequest.isExplainRequest()) { LOG.info("Request is falling back to old SQL engine due to: " + exception.getMessage()); } - LOG.debug("[{}] Request {} is not supported and falling back to old SQL engine", + LOG.info("[{}] Request {} is not supported and falling back to old SQL engine", QueryContext.getRequestId(), newSqlRequest); QueryAction queryAction = explainRequest(client, sqlRequest, format); executeSqlRequest(request, queryAction, client, restChannel); From 2895883b2f3a3658eee8723037d20b5dcf0c570f Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 14 Feb 2023 20:26:35 -0800 Subject: [PATCH 14/46] Add an end-to-end integration test. Signed-off-by: Yury-Fridlyand --- .../sql/sql/StandalonePaginationIT.java | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 4f01955fe93..dcd5d8c9913 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -6,6 +6,7 @@ package org.opensearch.sql.sql; import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata; +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -13,8 +14,14 @@ import java.util.List; import java.util.Map; import lombok.Getter; +import lombok.SneakyThrows; +import org.json.JSONArray; +import org.json.JSONObject; import org.junit.Test; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.opensearch.client.Request; +import org.opensearch.client.Response; import org.opensearch.client.RestHighLevelClient; import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; @@ -42,6 +49,7 @@ import org.opensearch.sql.util.InternalRestHighLevelClient; import org.opensearch.sql.util.StandaloneModule; +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class StandalonePaginationIT extends SQLIntegTestCase { private PaginatedQueryService paginatedQueryService; @@ -51,7 +59,14 @@ public class StandalonePaginationIT extends SQLIntegTestCase { private OpenSearchClient client; @Override + @SneakyThrows public void init() { + loadIndex(Index.ACCOUNT); + loadIndex(Index.ONLINE); + loadIndex(Index.BEER); + loadIndex(Index.BANK); + executeRequest(new Request("PUT", "/empty")); + RestHighLevelClient restClient = new InternalRestHighLevelClient(client()); client = new OpenSearchRestClient(restClient); DataSourceService dataSourceService = new DataSourceServiceImpl( @@ -69,7 +84,7 @@ public void init() { } @Test - public void testPagination() throws IOException { + public void test_pagination_whitebox() throws IOException { class TestResponder implements ResponseListener { @Getter @@ -118,6 +133,46 @@ public void onFailure(Exception e) { // act 3: confirm that there's no cursor. } + @Test + @SneakyThrows + public void test_pagination_blackbox() { + var indices = getResponseBody(client().performRequest(new Request("GET", "_cat/indices?h=i")), true).split("\n"); + for (var index : indices) { + var response = executeJdbcRequest(String.format("select * from %s", index)); + var indexSize = response.getInt("total"); + var rows = response.getJSONArray("datarows"); + var schema = response.getJSONArray("schema"); + for (var pageSize : List.of(1, 5, 10, 100, 1000)) { + var testReportPrefix = String.format("index: %s, page size: %d || ", index, pageSize); + var rowsPaged = new JSONArray(); + var rowsReturned = 0; + response = new JSONObject(executeFetchQuery( + String.format("select * from %s", index), pageSize, "jdbc")); + while (response.has("cursor")) { + var cursor = response.getString("cursor"); + assertTrue(testReportPrefix + "Cursor returned from legacy engine", + cursor.startsWith("n:")); + rowsReturned += response.getInt("total"); + var datarows = response.getJSONArray("datarows"); + for (int i = 0; i < datarows.length(); i++) { + rowsPaged.put(datarows.get(i)); + } + assertTrue("Paged response schema doesn't match to non-paged", + schema.similar(response.getJSONArray("schema"))); + response = executeCursorQuery(cursor); + } + assertEquals(testReportPrefix + "Last page is not empty", + 0, response.getInt("total")); + assertEquals(testReportPrefix + "Last page is not empty", + 0, response.getJSONArray("datarows").length()); + assertEquals(testReportPrefix + "Paged responses return another row count that non-paged", + indexSize, rowsReturned); + assertTrue(testReportPrefix + "Paged accumulated result has other rows than non-paged", + rows.similar(rowsPaged)); + } + } + } + private Settings defaultSettings() { return new Settings() { private final Map defaultSettings = new ImmutableMap.Builder() From dd6fcd67823d338e63e44a705c3efca9f39bd898 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Feb 2023 17:18:38 -0800 Subject: [PATCH 15/46] Add `explain` request handlers. Signed-off-by: Yury-Fridlyand --- .../opensearch/sql/executor/execution/PaginatedPlan.java | 3 ++- sql/src/main/java/org/opensearch/sql/sql/SQLService.java | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java index 62c913ee672..36b1b23e5e8 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java @@ -39,6 +39,7 @@ public void execute() { @Override public void explain(ResponseListener listener) { - throw new UnsupportedOperationException("implement PaginatedPlan.explain"); + listener.onFailure(new UnsupportedOperationException( + "`explain` feature for paginated requests is not implemented yet.")); } } 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 0912bf34826..908f03e4f26 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -67,6 +67,12 @@ private AbstractPlan plan( Optional> explainListener) { if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. + if (queryListener.isEmpty() && explainListener.isPresent()) { // explain request + explainListener.get().onFailure(new UnsupportedOperationException( + "`explain` request for cursor requests is not supported. " + + "Use `explain` for the initial query request.")); + } + // non-explain request return queryExecutionFactory.create(request.getCursor().get(), queryListener.get()); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) From 2a19e5647c73c39a66e1f72537c55aa30d4c934e Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Feb 2023 20:05:08 -0800 Subject: [PATCH 16/46] Add IT for explain. Signed-off-by: Yury-Fridlyand --- .../sql/sql/StandalonePaginationIT.java | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index dcd5d8c9913..46613539a4a 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -22,6 +22,7 @@ import org.junit.jupiter.api.DisplayNameGenerator; import org.opensearch.client.Request; import org.opensearch.client.Response; +import org.opensearch.client.ResponseException; import org.opensearch.client.RestHighLevelClient; import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; @@ -133,6 +134,7 @@ public void onFailure(Exception e) { // act 3: confirm that there's no cursor. } + // Test takes 3+ min due to a big amount of requests issued @Test @SneakyThrows public void test_pagination_blackbox() { @@ -173,6 +175,25 @@ public void test_pagination_blackbox() { } } + @Test + @SneakyThrows + public void test_explain_not_supported() { + var request = new Request("POST", "_plugins/_sql/_explain"); + // Request should be rejected before index names are resolved + request.setJsonEntity("{ \"query\": \"select * from something\", \"fetch_size\": 10 }"); + var exception = assertThrows(ResponseException.class, () -> client().performRequest(request)); + var response = new JSONObject(new String(exception.getResponse().getEntity().getContent().readAllBytes())); + assertEquals("`explain` feature for paginated requests is not implemented yet.", + response.getJSONObject("error").getString("details")); + + // Request should be rejected before cursor parsed + request.setJsonEntity("{ \"cursor\" : \"n:0000\" }"); + exception = assertThrows(ResponseException.class, () -> client().performRequest(request)); + response = new JSONObject(new String(exception.getResponse().getEntity().getContent().readAllBytes())); + assertEquals("`explain` request for cursor requests is not supported. Use `explain` for the initial query request.", + response.getJSONObject("error").getString("details")); + } + private Settings defaultSettings() { return new Settings() { private final Map defaultSettings = new ImmutableMap.Builder() From a3ef2bfb049849309eab2e1b5a9b9809b3cedb65 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 17 Feb 2023 01:01:54 -0800 Subject: [PATCH 17/46] Address issues flagged by checkstyle build step (#229) Signed-off-by: MaxKsyunz --- .../sql/executor/CanPaginateVisitor.java | 7 +------ .../sql/executor/PaginatedPlanCache.java | 18 ++++++++++++++---- .../planner/optimizer/rule/PushPageSize.java | 2 -- .../planner/physical/RemoveOperatorTest.java | 3 +-- .../storage/ContinueScrollRequest.java | 4 ++-- .../storage/OpenSearchPagedIndexScan.java | 3 --- .../storage/PagedRequestBuilder.java | 1 + .../storage/SubsequentPageRequestBuilder.java | 13 ++++--------- .../OpenSearchIndexScanAggregationBuilder.java | 2 +- .../scan/OpenSearchIndexScanQueryBuilder.java | 2 +- .../storage/script/ExpressionScriptEngine.java | 2 +- .../aggregation/AggregationQueryBuilder.java | 2 +- .../dsl/AggregationBuilderHelper.java | 2 +- .../dsl/BucketAggregationBuilder.java | 2 +- .../dsl/MetricAggregationBuilder.java | 2 +- .../script/filter/FilterQueryBuilder.java | 2 +- .../sql/opensearch/executor/CursorTest.java | 4 ++-- .../OpenSearchExecutionEngineTest.java | 2 +- .../script/ExpressionScriptEngineTest.java | 2 +- .../org/opensearch/sql/plugin/SQLPlugin.java | 2 +- .../org/opensearch/sql/ppl/PPLServiceTest.java | 6 +++++- .../ppl/parser/AstStatementBuilderTest.java | 3 ++- .../sql/sql/domain/SQLQueryRequest.java | 7 ++++--- .../org/opensearch/sql/sql/SQLServiceTest.java | 6 +++++- 24 files changed, 52 insertions(+), 47 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index c85bbf0cc8c..e3f7c7ad603 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -5,17 +5,12 @@ package org.opensearch.sql.executor; +import java.util.concurrent.atomic.AtomicBoolean; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.AllFields; -import org.opensearch.sql.ast.tree.Filter; -import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.Relation; -import org.opensearch.sql.ast.tree.Sort; -import org.opensearch.sql.ast.tree.Values; - -import java.util.concurrent.atomic.AtomicBoolean; /** * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache. diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 73c058316c9..49cd02dcf21 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -57,10 +57,15 @@ public Cursor convertToCursor(PhysicalPlan plan) { } } + /** + * Compress serialized query plan. + * @param str string representing a query plan + * @return str compressed with gzip. + */ @SneakyThrows public static String compress(String str) { if (str == null || str.length() == 0) { - return null; + return null; } ByteArrayOutputStream out = new ByteArrayOutputStream(); @@ -70,10 +75,15 @@ public static String compress(String str) { return HashCode.fromBytes(out.toByteArray()).toString(); } + /** + * Decompresses a query plan that was compress with {@link PaginatedPlanCache.compress}. + * @param input compressed query plan + * @return seria + */ @SneakyThrows public static String decompress(String input) { if (input == null || input.length() == 0) { - return null; + return null; } GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream( HashCode.fromString(input).asBytes())); @@ -114,10 +124,10 @@ public PhysicalPlan convertToPlan(String cursor) { throw new UnsupportedOperationException("Unsupported cursor"); } cursor = cursor.substring(cursor.indexOf(',') + 1); - int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); + final int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); cursor = cursor.substring(cursor.indexOf(',') + 1); - int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); + final int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); cursor = cursor.substring(cursor.indexOf(',') + 1); if (!cursor.startsWith("(Project,")) { diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java index 78f856535f3..95cd23d6ca0 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java @@ -17,8 +17,6 @@ import org.opensearch.sql.planner.logical.LogicalRelation; import org.opensearch.sql.planner.optimizer.Rule; -import java.util.Objects; - public class PushPageSize implements Rule { /** Capture the table inside matched logical paginate operator. */ diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java index bf046bf0a6b..1cc7d5532fb 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java @@ -117,8 +117,7 @@ public void invalid_to_retrieve_schema_from_remove() { IllegalStateException exception = assertThrows(IllegalStateException.class, () -> plan.schema()); assertEquals( - "[BUG] schema can been only applied to ProjectOperator, " - + "instead of RemoveOperator(input=inputPlan, removeList=[response, referer])", + "[BUG] schema can been only applied to ProjectOperator, instead of RemoveOperator", exception.getMessage()); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java index ddc76f6ba22..b33b7fd9a33 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java @@ -5,6 +5,8 @@ package org.opensearch.sql.opensearch.storage; +import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.DEFAULT_SCROLL_TIMEOUT; + import java.util.function.Consumer; import java.util.function.Function; import lombok.EqualsAndHashCode; @@ -18,8 +20,6 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; -import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.DEFAULT_SCROLL_TIMEOUT; - public class ContinueScrollRequest implements OpenSearchRequest { final String initialScrollId; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java index 646b414183d..fa4e4bf1065 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java @@ -7,14 +7,11 @@ import java.util.Collections; import java.util.Iterator; -import java.util.List; - import lombok.EqualsAndHashCode; import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.request.OpenSearchRequest; -import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.storage.TableScanOperator; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java index c138e327550..ae89a238a0e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java @@ -10,5 +10,6 @@ public interface PagedRequestBuilder { OpenSearchRequest build(); + OpenSearchRequest.IndexName getIndexName(); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java index dca63e8e99f..89f8f71933f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java @@ -5,20 +5,15 @@ package org.opensearch.sql.opensearch.storage; +import lombok.RequiredArgsConstructor; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +@RequiredArgsConstructor public class SubsequentPageRequestBuilder implements PagedRequestBuilder { - private OpenSearchRequest.IndexName indexName; + final OpenSearchRequest.IndexName indexName; final String scrollId; - private OpenSearchExprValueFactory exprValueFactory; - - public SubsequentPageRequestBuilder(OpenSearchRequest.IndexName indexName, String scanAsString, - OpenSearchExprValueFactory exprValueFactory) { - this.indexName = indexName; - scrollId = scanAsString; - this.exprValueFactory = exprValueFactory; - } + final OpenSearchExprValueFactory exprValueFactory; @Override public OpenSearchRequest build() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 719ef52b79c..4e1b20db6e3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -15,10 +15,10 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; -import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalSort; import org.opensearch.sql.storage.TableScanOperator; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index a67ceeae9ae..cb0940c410e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -20,10 +20,10 @@ import org.opensearch.sql.expression.ExpressionNodeVisitor; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; -import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java index a48da591806..9e8b47f6b05 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngine.java @@ -16,9 +16,9 @@ import org.opensearch.script.ScriptContext; import org.opensearch.script.ScriptEngine; import org.opensearch.sql.expression.Expression; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; import org.opensearch.sql.opensearch.storage.script.aggregation.ExpressionAggregationScriptFactory; import org.opensearch.sql.opensearch.storage.script.filter.ExpressionFilterScriptFactory; -import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Custom expression script engine that supports using core engine expression code in DSL diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java index f8ede598c1f..14f38bd5f37 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/AggregationQueryBuilder.java @@ -30,13 +30,13 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; import org.opensearch.sql.opensearch.response.agg.CompositeAggregationParser; import org.opensearch.sql.opensearch.response.agg.MetricParser; import org.opensearch.sql.opensearch.response.agg.NoBucketAggregationParser; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.storage.script.aggregation.dsl.BucketAggregationBuilder; import org.opensearch.sql.opensearch.storage.script.aggregation.dsl.MetricAggregationBuilder; -import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Build the AggregationBuilder from the list of {@link NamedAggregator} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java index 9cba3622a0d..d9360276a05 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/AggregationBuilderHelper.java @@ -17,8 +17,8 @@ import org.opensearch.sql.expression.FunctionExpression; import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.storage.script.ScriptUtils; import org.opensearch.sql.expression.serialization.ExpressionSerializer; +import org.opensearch.sql.opensearch.storage.script.ScriptUtils; /** * Abstract Aggregation Builder. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java index 5e910e73c3a..215be3b3565 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/BucketAggregationBuilder.java @@ -23,8 +23,8 @@ import org.opensearch.search.sort.SortOrder; import org.opensearch.sql.ast.expression.SpanUnit; import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.expression.span.SpanExpression; import org.opensearch.sql.expression.serialization.ExpressionSerializer; +import org.opensearch.sql.expression.span.SpanExpression; /** * Bucket Aggregation Builder. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java index d00a4eaf89a..db8d1fdf1eb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java @@ -25,13 +25,13 @@ import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; import org.opensearch.sql.opensearch.response.agg.FilterParser; import org.opensearch.sql.opensearch.response.agg.MetricParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; import org.opensearch.sql.opensearch.response.agg.StatsParser; import org.opensearch.sql.opensearch.response.agg.TopHitsParser; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; -import org.opensearch.sql.expression.serialization.ExpressionSerializer; /** * Build the Metric Aggregation and List of {@link MetricParser} from {@link NamedAggregator}. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java index ff206a6a1e2..a82869ec038 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/FilterQueryBuilder.java @@ -24,6 +24,7 @@ import org.opensearch.sql.expression.FunctionExpression; import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.expression.function.FunctionName; +import org.opensearch.sql.expression.serialization.ExpressionSerializer; import org.opensearch.sql.opensearch.storage.script.filter.lucene.LikeQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.LuceneQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.RangeQuery; @@ -38,7 +39,6 @@ import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.QueryStringQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.SimpleQueryStringQuery; import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.WildcardQuery; -import org.opensearch.sql.expression.serialization.ExpressionSerializer; @RequiredArgsConstructor public class FilterQueryBuilder extends ExpressionNodeVisitor { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java index 6b2a1a7c574..6833eb8e887 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/CursorTest.java @@ -11,12 +11,12 @@ class CursorTest { @Test - void EmptyArrayIsNone() { + void emptyArrayIsNone() { Assertions.assertEquals(Cursor.None, new Cursor(new byte[]{})); } @Test - void ToStringIsArrayValue() { + void toStringIsArrayValue() { String cursorTxt = "This is a test"; Assertions.assertEquals(cursorTxt, new Cursor(cursorTxt.getBytes()).toString()); } 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 ff529a018f7..7f01bc605b2 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 @@ -36,10 +36,10 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; -import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.ExecutionContext; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java index b106f396fb0..a88d81c0201 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/ExpressionScriptEngineTest.java @@ -27,8 +27,8 @@ import org.opensearch.script.ScriptEngine; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; -import org.opensearch.sql.opensearch.storage.script.filter.ExpressionFilterScriptFactory; import org.opensearch.sql.expression.serialization.ExpressionSerializer; +import org.opensearch.sql.opensearch.storage.script.filter.ExpressionFilterScriptFactory; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 08ad3b963a1..16b551a46ce 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -57,6 +57,7 @@ import org.opensearch.sql.datasource.DataSourceServiceImpl; import org.opensearch.sql.datasource.model.DataSource; import org.opensearch.sql.datasource.model.DataSourceMetadata; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.legacy.esdomain.LocalClusterState; import org.opensearch.sql.legacy.executor.AsyncRestExecutor; import org.opensearch.sql.legacy.metrics.Metrics; @@ -68,7 +69,6 @@ import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory; import org.opensearch.sql.opensearch.storage.script.ExpressionScriptEngine; -import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.plugin.config.OpenSearchPluginModule; import org.opensearch.sql.plugin.datasource.DataSourceSettings; import org.opensearch.sql.plugin.rest.RestPPLQueryAction; diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index ecac53a9c1b..774143a3484 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -26,6 +26,7 @@ import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.execution.PaginatedQueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; @@ -45,6 +46,9 @@ public class PPLServiceTest { @Mock private QueryService queryService; + @Mock + private PaginatedQueryService paginatedQueryService; + @Mock private ExecutionEngine.Schema schema; @@ -59,7 +63,7 @@ public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); pplService = new PPLService(new PPLSyntaxParser(), queryManager, - new QueryPlanFactory(queryService, paginatedPlanCache)); + new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache)); } @After diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java index cdb0e37ee50..de74e4932f9 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java @@ -39,7 +39,8 @@ public void buildQueryStatement() { "search source=t a=1", new Query( project( - filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFields.of()), 0)); + filter(relation("t"), compare("=", field("a"), + intLiteral(1))), AllFields.of()), 0)); } @Test 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 bd0db9e9a00..62db130fc2b 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 @@ -65,11 +65,12 @@ public class SQLQueryRequest { private boolean sanitize = true; private String cursor = ""; + /** * Constructor of SQLQueryRequest that passes request params. */ - public SQLQueryRequest( - JSONObject jsonContent, String query, String path, Map params, String cursor) { + public SQLQueryRequest(JSONObject jsonContent, String query, String path, + Map params, String cursor) { this.jsonContent = jsonContent; this.query = query; this.path = path; @@ -77,7 +78,7 @@ public SQLQueryRequest( this.format = getFormat(params); this.sanitize = shouldSanitize(params); // TODO hack - this.cursor = cursor == null? "" : cursor; + this.cursor = cursor == null ? "" : cursor; } /** 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 9fcf1fa0984..62718800b19 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -28,6 +28,7 @@ import org.opensearch.sql.executor.ExecutionEngine.ExplainResponseNode; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.execution.PaginatedQueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; @@ -47,6 +48,9 @@ class SQLServiceTest { @Mock private QueryService queryService; + @Mock + private PaginatedQueryService paginatedQueryService; + @Mock private ExecutionEngine.Schema schema; @@ -57,7 +61,7 @@ class SQLServiceTest { public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); sqlService = new SQLService(new SQLSyntaxParser(), queryManager, - new QueryPlanFactory(queryService, paginatedPlanCache)); + new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache)); } @AfterEach From 2d29549e489dee09e4f726f3fa097f8640d551d8 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 7 Mar 2023 17:10:31 -0800 Subject: [PATCH 18/46] Pagination, phase 1: Add unit tests for `:core` module with coverage. (#230) * Add unit tests for `:core` module with coverage. Uncovered: `toCursor`, because it is will be changed soon. Signed-off-by: Yury-Fridlyand --- core/build.gradle | 1 + .../sql/executor/CanPaginateVisitor.java | 2 +- .../sql/executor/PaginatedPlanCache.java | 4 +- .../execution/ContinuePaginatedPlan.java | 3 +- .../sql/executor/execution/PaginatedPlan.java | 4 +- .../execution/PaginatedQueryService.java | 6 +- .../sql/planner/PaginateOperator.java | 6 +- .../opensearch/sql/analysis/AnalyzerTest.java | 9 + .../sql/executor/CanPaginateVisitorTest.java | 131 ++++++ .../sql/executor/PaginatedPlanCacheTest.java | 417 +++++++++++++++++- .../execution/ContinuePaginatedPlanTest.java | 120 +++++ .../executor/execution/PaginatedPlanTest.java | 99 +++++ .../execution/QueryPlanFactoryTest.java | 27 ++ .../DefaultExpressionSerializerTest.java | 0 .../sql/planner/DefaultImplementorTest.java | 30 +- .../logical/LogicalPlanNodeVisitorTest.java | 160 +++---- .../optimizer/LogicalPlanOptimizerTest.java | 62 ++- .../optimizer/pattern/PatternsTest.java | 40 +- .../physical/PaginateOperatorTest.java | 86 ++++ .../physical/PhysicalPlanNodeVisitorTest.java | 9 + .../planner/physical/RemoveOperatorTest.java | 2 +- .../sql/storage/StorageEngineTest.java | 7 +- .../org/opensearch/sql/storage/TableTest.java | 25 ++ .../scan/OpenSearchIndexScanBuilder.java | 4 +- .../OpenSearchExecutionEngineTest.java | 78 ++++ .../opensearch/sql/sql/SQLServiceTest.java | 3 + 26 files changed, 1175 insertions(+), 160 deletions(-) create mode 100644 core/src/test/java/org/opensearch/sql/executor/CanPaginateVisitorTest.java create mode 100644 core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java create mode 100644 core/src/test/java/org/opensearch/sql/executor/execution/PaginatedPlanTest.java rename {opensearch/src/test/java/org/opensearch/sql/opensearch/storage => core/src/test/java/org/opensearch/sql/expression}/serialization/DefaultExpressionSerializerTest.java (100%) create mode 100644 core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java create mode 100644 core/src/test/java/org/opensearch/sql/storage/TableTest.java diff --git a/core/build.gradle b/core/build.gradle index 24384705184..e025b672f39 100644 --- a/core/build.gradle +++ b/core/build.gradle @@ -56,6 +56,7 @@ dependencies { testImplementation('org.junit.jupiter:junit-jupiter:5.6.2') testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: '2.1' testImplementation group: 'org.mockito', name: 'mockito-core', version: '3.12.4' + testImplementation group: 'org.mockito', name: 'mockito-inline', version: '3.12.4' testImplementation group: 'org.mockito', name: 'mockito-junit-jupiter', version: '3.12.4' } diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index e3f7c7ad603..d4c1c2f3005 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -39,13 +39,13 @@ public Boolean visitRelation(Relation node, Object context) { return Boolean.TRUE; } + /* private Boolean canPaginate(Node node, Object context) { AtomicBoolean result = new AtomicBoolean(true); node.getChild().forEach(n -> result.set(result.get() && n.accept(this, context))); return result.get(); } - /* For queries without `FROM` clause. Required to overload `toCursor` function in `ValuesOperator` and modify cursor parsing. @Override diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 49cd02dcf21..26294e3be15 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -47,7 +47,7 @@ static class SerializationContext { public Cursor convertToCursor(PhysicalPlan plan) { if (plan instanceof PaginateOperator) { var cursor = plan.toCursor(); - if (cursor == null || cursor.isEmpty()) { + if (cursor == null) { return Cursor.None; } var raw = CURSOR_PREFIX + compress(cursor); @@ -67,7 +67,6 @@ public static String compress(String str) { if (str == null || str.length() == 0) { return null; } - ByteArrayOutputStream out = new ByteArrayOutputStream(); GZIPOutputStream gzip = new GZIPOutputStream(out); gzip.write(str.getBytes()); @@ -123,6 +122,7 @@ public PhysicalPlan convertToPlan(String cursor) { if (!cursor.startsWith("(Paginate,")) { throw new UnsupportedOperationException("Unsupported cursor"); } + // TODO add checks for > 0 cursor = cursor.substring(cursor.indexOf(',') + 1); final int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java index adefae778e1..d61747d0eb7 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java @@ -49,7 +49,8 @@ public void execute() { } @Override + // TODO why can't use listener given in the constructor? public void explain(ResponseListener listener) { - throw new NotImplementedException("Explain of query continuation is not supported"); + throw new UnsupportedOperationException("Explain of query continuation is not supported"); } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java index 36b1b23e5e8..65a1cac5b75 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedPlan.java @@ -5,6 +5,7 @@ package org.opensearch.sql.executor.execution; +import org.apache.commons.lang3.NotImplementedException; import org.opensearch.sql.ast.tree.Paginate; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; @@ -38,8 +39,9 @@ public void execute() { } @Override + // TODO why can't use listener given in the constructor? public void explain(ResponseListener listener) { - listener.onFailure(new UnsupportedOperationException( + listener.onFailure(new NotImplementedException( "`explain` feature for paginated requests is not implemented yet.")); } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java index e08c6dc59da..9e2ad73336b 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java @@ -46,11 +46,7 @@ public void executePlan(LogicalPlan plan, */ public void executePlan(PhysicalPlan plan, ResponseListener listener) { - try { - executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); - } catch (Exception e) { - listener.onFailure(e); - } + executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); } public LogicalPlan analyze(UnresolvedPlan plan) { diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java index a99beea8365..d867674e053 100644 --- a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java @@ -54,6 +54,7 @@ public boolean hasNext() { @Override public void open() { super.open(); + // TODO numReturned set to 0 for each new object. Do plans support re-opening? numReturned = 0; } @@ -69,7 +70,10 @@ public List getChild() { @Override public ExecutionEngine.Schema schema() { - assert input instanceof ProjectOperator; + // TODO remove assert or do in constructor + if (!(input instanceof ProjectOperator)) { + throw new UnsupportedOperationException(); + } return input.schema(); } diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index 1db29a6a42c..01e2091da93 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -75,6 +75,7 @@ import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Kmeans; import org.opensearch.sql.ast.tree.ML; +import org.opensearch.sql.ast.tree.Paginate; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.exception.ExpressionEvaluationException; import org.opensearch.sql.exception.SemanticCheckException; @@ -83,6 +84,7 @@ import org.opensearch.sql.expression.window.WindowDefinition; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalMLCommons; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalProject; @@ -1189,4 +1191,11 @@ public void ml_relation_predict_rcf_without_time_field() { assertTrue(((LogicalProject) actual).getProjectList() .contains(DSL.named(RCF_ANOMALOUS, DSL.ref(RCF_ANOMALOUS, BOOLEAN)))); } + + @Test + public void visit_paginate() { + LogicalPlan actual = analyze(new Paginate(10, AstDSL.relation("dummy"))); + assertTrue(actual instanceof LogicalPaginate); + assertEquals(10, ((LogicalPaginate) actual).getPageSize()); + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/CanPaginateVisitorTest.java b/core/src/test/java/org/opensearch/sql/executor/CanPaginateVisitorTest.java new file mode 100644 index 00000000000..c915685ba80 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/executor/CanPaginateVisitorTest.java @@ -0,0 +1,131 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +import java.util.List; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.ast.dsl.AstDSL; +import org.opensearch.sql.ast.tree.Project; +import org.opensearch.sql.ast.tree.Relation; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class CanPaginateVisitorTest { + + static final CanPaginateVisitor visitor = new CanPaginateVisitor(); + + @Test + // select * from y + public void accept_query_with_select_star_and_from() { + var plan = AstDSL.project(AstDSL.relation("dummy"), AstDSL.allFields()); + assertTrue(plan.accept(visitor, null)); + } + + @Test + // select x from y + public void reject_query_with_select_field_and_from() { + var plan = AstDSL.project(AstDSL.relation("dummy"), AstDSL.field("pewpew")); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select x,z from y + public void reject_query_with_select_fields_and_from() { + var plan = AstDSL.project(AstDSL.relation("dummy"), + AstDSL.field("pewpew"), AstDSL.field("pewpew")); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select x + public void reject_query_without_from() { + var plan = AstDSL.project(AstDSL.values(List.of(AstDSL.intLiteral(1))), + AstDSL.alias("1",AstDSL.intLiteral(1))); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select * from y limit z + public void reject_query_with_limit() { + var plan = AstDSL.project(AstDSL.limit(AstDSL.relation("dummy"), 1, 2), AstDSL.allFields()); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select * from y where z + public void reject_query_with_where() { + var plan = AstDSL.project(AstDSL.filter(AstDSL.relation("dummy"), + AstDSL.booleanLiteral(true)), AstDSL.allFields()); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select * from y order by z + public void reject_query_with_order_by() { + var plan = AstDSL.project(AstDSL.sort(AstDSL.relation("dummy"), AstDSL.field("1")), + AstDSL.allFields()); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select * from y group by z + public void reject_query_with_group_by() { + var plan = AstDSL.project(AstDSL.agg( + AstDSL.relation("dummy"), List.of(), List.of(), List.of(AstDSL.field("1")), List.of()), + AstDSL.allFields()); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select agg(x) from y + public void reject_query_with_aggregation_function() { + var plan = AstDSL.project(AstDSL.agg( + AstDSL.relation("dummy"), + List.of(AstDSL.alias("agg", AstDSL.aggregate("func", AstDSL.field("pewpew")))), + List.of(), List.of(), List.of()), + AstDSL.allFields()); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select window(x) from y + public void reject_query_with_window_function() { + var plan = AstDSL.project(AstDSL.relation("dummy"), + AstDSL.alias("pewpew", + AstDSL.window( + AstDSL.aggregate("func", AstDSL.field("pewpew")), + List.of(AstDSL.qualifiedName("1")), List.of()))); + assertFalse(plan.accept(visitor, null)); + } + + @Test + // select * from y, z + public void reject_query_with_select_from_multiple_indices() { + var plan = mock(Project.class); + when(plan.getChild()).thenReturn(List.of(AstDSL.relation("dummy"), AstDSL.relation("pummy"))); + when(plan.getProjectList()).thenReturn(List.of(AstDSL.allFields())); + assertFalse(visitor.visitProject(plan, null)); + } + + @Test + // unreal case, added for coverage only + public void reject_project_when_relation_has_child() { + var relation = mock(Relation.class, withSettings().useConstructor(AstDSL.qualifiedName("42"))); + when(relation.getChild()).thenReturn(List.of(AstDSL.relation("pewpew"))); + when(relation.accept(visitor, null)).thenCallRealMethod(); + var plan = mock(Project.class); + when(plan.getChild()).thenReturn(List.of(relation)); + when(plan.getProjectList()).thenReturn(List.of(AstDSL.allFields())); + assertFalse(visitor.visitProject((Project) plan, null)); + } +} diff --git a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java index 0b879a8c2de..fe798461ba6 100644 --- a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java @@ -5,34 +5,292 @@ package org.opensearch.sql.executor; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.Map; +import java.util.stream.Stream; +import java.util.zip.GZIPOutputStream; +import lombok.SneakyThrows; +import org.apache.commons.lang3.reflect.FieldUtils; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.Answers; import org.mockito.Mock; +import org.mockito.Mockito; import org.opensearch.sql.ast.dsl.AstDSL; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.opensearch.executor.Cursor; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.sql.storage.TableScanOperator; -class PaginatedPlanCacheTest { +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class PaginatedPlanCacheTest { - @Mock StorageEngine storageEngine; PaginatedPlanCache planCache; + // encoded query 'select * from cacls' o_O + static final String testCursor = "(Paginate,1,2,(Project," + + "(namedParseExpressions,),(projectList,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5" + + "OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVk" + + "dAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZ" + + "y5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH" + + "4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHl" + + "wZS9FeHByVHlwZTt4cHQABWJvb2wzc3IAGmphdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFh" + + "dAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAf" + + "gAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS" + + "5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAHQk9PTEVBTnEAfgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZX" + + "hwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAA" + + "JZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgAB" + + "eHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA" + + "0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3" + + "FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABGludDBzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYg" + + "G0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAA" + + "eHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAA" + + "HhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAdJTlRFR0VScQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlY" + + "XJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy" + + "9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAA" + + "EbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274" + + "AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZ" + + "W5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABXRpbWUxc3IAGmphdmEudXRpbC5BcnJheXMkQXJyYX" + + "lMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW5nLlN0cmluZzu" + + "t0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBl" + + "AAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAJVElNRVNUQU1QcQB+AAg=,rO0ABXNy" + + "AC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzd" + + "AASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0" + + "V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5" + + "jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5" + + "cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABWJvb2wyc3IAGmphdmEudXRpb" + + "C5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS" + + "5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGU" + + "uRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAHQk9PTEVBTnEA" + + "fgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIA" + + "A0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9le" + + "HByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW" + + "9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9" + + "MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABGludDJzcgAa" + + "amF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1c" + + "gATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLm" + + "RhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAd" + + "JTlRFR0VScQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb27" + + "4hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2Vh" + + "cmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxb" + + "C5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTG" + + "phdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQ" + + "ABGludDFzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9P" + + "YmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZ" + + "WFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAA" + + "AAEgAAeHB0AAdJTlRFR0VScQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZE" + + "V4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9" + + "yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVu" + + "c2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwAB" + + "XBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeH" + + "ByVHlwZTt4cHQABHN0cjNzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGp" + + "hdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgAp" + + "b3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuR" + + "W51bQAAAAAAAAAAEgAAeHB0AAZTVFJJTkdxAH4ACA==,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc" + + "2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZW" + + "dhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3I" + + "AMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0" + + "dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2Rhd" + + "GEvdHlwZS9FeHByVHlwZTt4cHQABGludDNzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAV" + + "sAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAA" + + "BcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5q" + + "YXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAdJTlRFR0VScQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5z" + + "cWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpb" + + "mc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZX" + + "EAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWv" + + "MkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFy" + + "Y2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABHN0cjFzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZp" + + "Dy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHX" + + "tHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAA" + + "AABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZTVFJJTkdxAH4ACA==,rO0ABXNyAC1vcmcub3B" + + "lbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEv" + + "bGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb" + + "247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3" + + "Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3J" + + "nL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABHN0cjJzcgAaamF2YS51dGlsLkFycmF5cyRB" + + "cnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3Rya" + + "W5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZV" + + "R5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZTVFJJTkdxAH4ACA==,rO0ABX" + + "NyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWF" + + "zdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9u" + + "L0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZ" + + "W5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABH" + + "R5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABXRpbWUwc3IAGmphdmEudXR" + + "pbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2" + + "YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5c" + + "GUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAJVElNRVNUQU" + + "1QcQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q" + + "2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3N" + + "xbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHBy" + + "ZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvd" + + "XRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQACWRhdG" + + "V0aW1lMHNyABpqYXZhLnV0aWwuQXJyYXlzJEFycmF5TGlzdNmkPL7NiAbSAgABWwABYXQAE1tMamF2YS9sYW5nL09" + + "iamVjdDt4cHVyABNbTGphdmEubGFuZy5TdHJpbmc7rdJW5+kde0cCAAB4cAAAAAFxAH4ACH5yAClvcmcub3BlbnNl" + + "YXJjaC5zcWwuZGF0YS50eXBlLkV4cHJDb3JlVHlwZQAAAAAAAAAAEgAAeHIADmphdmEubGFuZy5FbnVtAAAAAAAAA" + + "AASAAB4cHQACVRJTUVTVEFNUHEAfgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZ" + + "EV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG" + + "9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGV" + + "uc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwA" + + "BXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9Fe" + + "HByVHlwZTt4cHQABG51bTFzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTG" + + "phdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgA" + + "pb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcu" + + "RW51bQAAAAAAAAAAEgAAeHB0AAZET1VCTEVxAH4ACA==,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVz" + + "c2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZ" + + "WdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3" + + "IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF" + + "0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2Rh" + + "dGEvdHlwZS9FeHByVHlwZTt4cHQABG51bTBzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAA" + + "VsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAA" + + "ABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5" + + "qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZET1VCTEVxAH4ACA==,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5" + + "zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJp" + + "bmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZ" + + "XEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxW" + + "vMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWF" + + "yY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQACWRhdGV0aW1lMXNyABpqYXZhLnV0aWwuQXJyYXlzJEFycmF5" + + "TGlzdNmkPL7NiAbSAgABWwABYXQAE1tMamF2YS9sYW5nL09iamVjdDt4cHVyABNbTGphdmEubGFuZy5TdHJpbmc7r" + + "dJW5+kde0cCAAB4cAAAAAFxAH4ACH5yAClvcmcub3BlbnNlYXJjaC5zcWwuZGF0YS50eXBlLkV4cHJDb3JlVHlwZQ" + + "AAAAAAAAAAEgAAeHIADmphdmEubGFuZy5FbnVtAAAAAAAAAAASAAB4cHQACVRJTUVTVEFNUHEAfgAI,rO0ABXNyAC" + + "1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAA" + + "STGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4" + + "cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZ" + + "UV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cG" + + "V0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABG51bTRzcgAaamF2YS51dGlsLkF" + + "ycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxh" + + "bmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5Fe" + + "HByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZET1VCTEVxAH4ACA" + + "==,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0" + + "wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHB" + + "yZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9u" + + "LlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9Ma" + + "XN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABWJvb2wxc3IAGm" + + "phdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXI" + + "AE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5k" + + "YXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAHQ" + + "k9PTEVBTnEAfgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274h" + + "hKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2Vhcm" + + "NoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5" + + "leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGph" + + "dmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQAA" + + "2tleXNyABpqYXZhLnV0aWwuQXJyYXlzJEFycmF5TGlzdNmkPL7NiAbSAgABWwABYXQAE1tMamF2YS9sYW5nL09iam" + + "VjdDt4cHVyABNbTGphdmEubGFuZy5TdHJpbmc7rdJW5+kde0cCAAB4cAAAAAFxAH4ACH5yAClvcmcub3BlbnNlYXJ" + + "jaC5zcWwuZGF0YS50eXBlLkV4cHJDb3JlVHlwZQAAAAAAAAAAEgAAeHIADmphdmEubGFuZy5FbnVtAAAAAAAAAAAS" + + "AAB4cHQABlNUUklOR3EAfgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJl" + + "c3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vc" + + "GVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2Vhcm" + + "NoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGh" + + "zdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlw" + + "ZTt4cHQABG51bTNzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvb" + + "GFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm" + + "9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQA" + + "AAAAAAAAAEgAAeHB0AAZET1VCTEVxAH4ACA==,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5" + + "OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVk" + + "dAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZ" + + "y5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH" + + "4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHl" + + "wZS9FeHByVHlwZTt4cHQABWJvb2wwc3IAGmphdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFh" + + "dAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAf" + + "gAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS" + + "5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAHQk9PTEVBTnEAfgAI,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZX" + + "hwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAA" + + "JZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgAB" + + "eHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA" + + "0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3" + + "FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABG51bTJzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheUxpc3TZpDy+zYg" + + "G0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63SVufpHXtHAgAA" + + "eHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAA" + + "HhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZET1VCTEVxAH4ACA==,rO0ABXNyAC1vcmcub3BlbnNlY" + + "XJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy" + + "9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAA" + + "EbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274" + + "AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZ" + + "W5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABHN0cjBzcgAaamF2YS51dGlsLkFycmF5cyRBcnJheU" + + "xpc3TZpDy+zYgG0gIAAVsAAWF0ABNbTGphdmEvbGFuZy9PYmplY3Q7eHB1cgATW0xqYXZhLmxhbmcuU3RyaW5nO63" + + "SVufpHXtHAgAAeHAAAAABcQB+AAh+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUA" + + "AAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAZTVFJJTkdxAH4ACA==,rO0ABXNyAC1v" + + "cmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAAST" + + "GphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cH" + + "Jlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV" + + "4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0" + + "ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABWRhdGUzc3IAGmphdmEudXRpbC5Bc" + + "nJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW" + + "5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXh" + + "wckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAJVElNRVNUQU1QcQB+" + + "AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIA" + + "A0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9le" + + "HByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW" + + "9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9" + + "MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQABWRhdGUyc3IA" + + "GmphdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwd" + + "XIAE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC" + + "5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAA" + + "JVElNRVNUQU1QcQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3N" + + "pb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVu" + + "c2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoL" + + "nNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdA" + + "AQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt" + + "4cHQABWRhdGUxc3IAGmphdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAFhdAATW0xqYXZhL2xh" + + "bmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEAfgAIfnIAKW9yZy5vc" + + "GVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAA" + + "AAAAAAABIAAHhwdAAJVElNRVNUQU1QcQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zcWwuZXhwcmVzc2lvbi" + + "5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbmc7TAAJZGVsZWdhdGV" + + "kdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXEAfgABeHBwc3IAMW9y" + + "Zy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvMkAIAA0wABGF0dHJxA" + + "H4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdH" + + "lwZS9FeHByVHlwZTt4cHQABWRhdGUwc3IAGmphdmEudXRpbC5BcnJheXMkQXJyYXlMaXN02aQ8vs2IBtICAAFbAAF" + + "hdAATW0xqYXZhL2xhbmcvT2JqZWN0O3hwdXIAE1tMamF2YS5sYW5nLlN0cmluZzut0lbn6R17RwIAAHhwAAAAAXEA" + + "fgAIfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2Y" + + "S5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAJVElNRVNUQU1QcQB+AAg=,rO0ABXNyAC1vcmcub3BlbnNlYXJjaC5zc" + + "WwuZXhwcmVzc2lvbi5OYW1lZEV4cHJlc3Npb274hhKW/q2YQQIAA0wABWFsaWFzdAASTGphdmEvbGFuZy9TdHJpbm" + + "c7TAAJZGVsZWdhdGVkdAAqTG9yZy9vcGVuc2VhcmNoL3NxbC9leHByZXNzaW9uL0V4cHJlc3Npb247TAAEbmFtZXE" + + "AfgABeHBwc3IAMW9yZy5vcGVuc2VhcmNoLnNxbC5leHByZXNzaW9uLlJlZmVyZW5jZUV4cHJlc3Npb274AO0rxWvM" + + "kAIAA0wABGF0dHJxAH4AAUwABXBhdGhzdAAQTGphdmEvdXRpbC9MaXN0O0wABHR5cGV0ACdMb3JnL29wZW5zZWFyY" + + "2gvc3FsL2RhdGEvdHlwZS9FeHByVHlwZTt4cHQAA3p6enNyABpqYXZhLnV0aWwuQXJyYXlzJEFycmF5TGlzdNmkPL" + + "7NiAbSAgABWwABYXQAE1tMamF2YS9sYW5nL09iamVjdDt4cHVyABNbTGphdmEubGFuZy5TdHJpbmc7rdJW5+kde0c" + + "CAAB4cAAAAAFxAH4ACH5yAClvcmcub3BlbnNlYXJjaC5zcWwuZGF0YS50eXBlLkV4cHJDb3JlVHlwZQAAAAAAAAAA" + + "EgAAeHIADmphdmEubGFuZy5FbnVtAAAAAAAAAAASAAB4cHQABlNUUklOR3EAfgAI),(OpenSearchPagedIndexSc" + + "an,calcs,FGluY2x1ZGVfY29udGV4dF91dWlkDXF1ZXJ5QW5kRmV0Y2gBFndYQmJZcHpxU3dtc1hUVkhhYU1uLVEA" + + "AAAAAAAADRY4RzRudHZqbFI0dTBFdkJNZEpCaDd3)))"; + + private static final String testIndexName = "dummyIndex"; + private static final String testScroll = "dummyScroll"; + @BeforeEach void setUp() { + storageEngine = mock(StorageEngine.class); + when(storageEngine.getTableScan(anyString(), anyString())) + .thenReturn(new MockedTableScanOperator()); planCache = new PaginatedPlanCache(storageEngine); } @Test void canConvertToCursor_relation() { - Assertions.assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); + assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); } @Test void canConvertToCursor_project_allFields_relation() { var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.allFields()); - Assertions.assertTrue(planCache.canConvertToCursor(unresolvedPlan)); + assertTrue(planCache.canConvertToCursor(unresolvedPlan)); } @Test @@ -40,4 +298,155 @@ void canConvertToCursor_project_some_fields_relation() { var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.field("rando")); Assertions.assertFalse(planCache.canConvertToCursor(unresolvedPlan)); } + + @ParameterizedTest + @ValueSource(strings = {"pewpew", "asdkfhashdfjkgakgfwuigfaijkb", testCursor}) + void compress_decompress(String input) { + var compressed = PaginatedPlanCache.compress(input); + assertEquals(input, PaginatedPlanCache.decompress(compressed)); + if (input.length() > 200) { + // Compression of short strings isn't profitable, because encoding into string and gzip + // headers add more bytes than input string has. + assertTrue(compressed.length() < input.length()); + } + } + + @Test + // should never happen actually, at least for compress + void compress_decompress_null_or_empty_string() { + assertAll( + () -> assertNull(PaginatedPlanCache.compress(null)), + () -> assertNull(PaginatedPlanCache.compress("")), + () -> assertNull(PaginatedPlanCache.decompress(null)), + () -> assertNull(PaginatedPlanCache.decompress("")) + ); + } + + @Test + // test added for coverage only + void compress_throws() { + var mock = Mockito.mockConstructionWithAnswer(GZIPOutputStream.class, invocation -> null); + assertThrows(Throwable.class, () -> PaginatedPlanCache.compress("\\_(`v`)_/")); + mock.close(); + } + + @Test + void decompress_throws() { + assertAll( + // from gzip - damaged header + () -> assertThrows(Throwable.class, () -> PaginatedPlanCache.decompress("00")), + // from HashCode::fromString + () -> assertThrows(Throwable.class, () -> PaginatedPlanCache.decompress("000")) + ); + } + + @Test + @SneakyThrows + void convert_deconvert_cursor() { + var cursor = buildCursor(Map.of()); + var plan = planCache.convertToPlan(cursor); + // `PaginateOperator::toCursor` shifts cursor to the next page. To have this test consistent + // we have to enforce it staying on the same page. This allows us to get same cursor strings. + var pageNum = (int)FieldUtils.readField(plan, "pageIndex", true); + FieldUtils.writeField(plan, "pageIndex", pageNum - 1, true); + var convertedCursor = planCache.convertToCursor(plan).toString(); + // Then we have to restore page num into the plan, otherwise comparison would fail due to this. + FieldUtils.writeField(plan, "pageIndex", pageNum, true); + var convertedPlan = planCache.convertToPlan(convertedCursor); + assertEquals(cursor, convertedCursor); + // TODO compare plans + } + + @Test + void convertToCursor_cant_convert() { + var plan = mock(MockedTableScanOperator.class); + assertEquals(Cursor.None, planCache.convertToCursor(plan)); + when(plan.toCursor()).thenReturn(""); + assertEquals(Cursor.None, planCache.convertToCursor( + new PaginateOperator(plan, 1, 2))); + } + + @Test + void converted_plan_is_executable() { + // planCache.convertToPlan(buildCursor(Map.of())); + var plan = planCache.convertToPlan("n:" + PaginatedPlanCache.compress(testCursor)); + // TODO + } + + @ParameterizedTest + @MethodSource("generateIncorrectCursors") + void throws_on_parsing_damaged_cursor(String cursor) { + assertThrows(Throwable.class, () -> planCache.convertToPlan(cursor)); + } + + private static Stream generateIncorrectCursors() { + return Stream.of( + PaginatedPlanCache.compress(testCursor), // a valid cursor, but without "n:" prefix + "n:" + testCursor, // a valid, but uncompressed cursor + buildCursor(Map.of("prefix", "g:")), // incorrect prefix + buildCursor(Map.of("header: paginate", "ORDER BY")), // incorrect header + buildCursor(Map.of("pageIndex", "")), // incorrect page # + buildCursor(Map.of("pageIndex", "abc")), // incorrect page # + buildCursor(Map.of("pageSize", "abc")), // incorrect page size + buildCursor(Map.of("pageSize", "null")), // incorrect page size + buildCursor(Map.of("pageSize", "10 ")), // incorrect page size + buildCursor(Map.of("header: project", "")), // incorrect header + buildCursor(Map.of("header: namedParseExpressions", "ololo")), // incorrect header + buildCursor(Map.of("namedParseExpressions", "pewpew")), // incorrect (unparsable) npes + buildCursor(Map.of("namedParseExpressions", "rO0ABXA=,")), // incorrect npes (extra comma) + buildCursor(Map.of("header: projectList", "")), // incorrect header + buildCursor(Map.of("projectList", "\0\0\0\0")), // incorrect project + buildCursor(Map.of("header: OpenSearchPagedIndexScan", "42")) // incorrect header + ).map(Arguments::of); + } + + + /** + * Function puts default valid values into generated cursor string. + * Values could be redefined. + * @param values A map of non-default values to use. + * @return A compressed cursor string. + */ + public static String buildCursor(Map values) { + String prefix = values.getOrDefault("prefix", "n:"); + String headerPaginate = values.getOrDefault("header: paginate", "Paginate"); + String pageIndex = values.getOrDefault("pageIndex", "1"); + String pageSize = values.getOrDefault("pageSize", "2"); + String headerProject = values.getOrDefault("header: project", "Project"); + String headerNpes = values.getOrDefault("header: namedParseExpressions", + "namedParseExpressions"); + String namedParseExpressions = values.getOrDefault("namedParseExpressions", ""); + String headerProjectList = values.getOrDefault("header: projectList", "projectList"); + String projectList = values.getOrDefault("projectList", "rO0ABXA="); // serialized `null` + String headerOspis = values.getOrDefault("header: OpenSearchPagedIndexScan", + "OpenSearchPagedIndexScan"); + String indexName = values.getOrDefault("indexName", testIndexName); + String scrollId = values.getOrDefault("scrollId", testScroll); + var cursor = String.format("(%s,%s,%s,(%s,(%s,%s),(%s,%s),(%s,%s,%s)))", headerPaginate, + pageIndex, pageSize, headerProject, headerNpes, namedParseExpressions, headerProjectList, + projectList, headerOspis, indexName, scrollId); + return prefix + PaginatedPlanCache.compress(cursor); + } + + private static class MockedTableScanOperator extends TableScanOperator { + @Override + public boolean hasNext() { + return false; + } + + @Override + public ExprValue next() { + return null; + } + + @Override + public String explain() { + return null; + } + + @Override + public String toCursor() { + return createSection("OpenSearchPagedIndexScan", testIndexName, testScroll); + } + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java new file mode 100644 index 00000000000..c8c95a0ae68 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java @@ -0,0 +1,120 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor.execution; + +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.executor.PaginatedPlanCacheTest.buildCursor; + +import java.util.Map; +import org.apache.commons.lang3.reflect.FieldUtils; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.DefaultExecutionEngine; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.executor.QueryId; +import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.sql.storage.TableScanOperator; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class ContinuePaginatedPlanTest { + + private static PaginatedPlanCache paginatedPlanCache; + + private static PaginatedQueryService paginatedQueryService; + + /** + * Initialize the mocks. + */ + @BeforeAll + public static void setUp() { + var storageEngine = mock(StorageEngine.class); + when(storageEngine.getTableScan(anyString(), anyString())) + .thenReturn(mock(TableScanOperator.class)); + paginatedPlanCache = new PaginatedPlanCache(storageEngine); + paginatedQueryService = new PaginatedQueryService( + null, new DefaultExecutionEngine(), null); + } + + @Test + public void none_plan_is_empty() { + var plan = ContinuePaginatedPlan.None; + + assertAll( + () -> assertTrue(plan.getQueryId().getQueryId().isEmpty()), + () -> { + var cursor = (String) FieldUtils.readField(plan, "cursor", true); + assertTrue(cursor.isEmpty()); + }, + () -> { + var pqs = (PaginatedQueryService) FieldUtils.readField(plan, "queryService", true); + assertNull(pqs); + }, + () -> { + var ppc = (PaginatedPlanCache) FieldUtils.readField(plan, "paginatedPlanCache", true); + assertNull(ppc); + }, + () -> { + var rl = (ResponseListener) FieldUtils.readField(plan, "queryResponseListener", true); + assertNull(rl); + } + ); + } + + @Test + public void can_execute_plan() { + var listener = new ResponseListener() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + fail(); + } + }; + var plan = new ContinuePaginatedPlan(QueryId.None, buildCursor(Map.of()), + paginatedQueryService, paginatedPlanCache, listener); + plan.execute(); + } + + @Test + // Same as previous test, but with malformed cursor + public void can_handle_error_while_executing_plan() { + var listener = new ResponseListener() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + fail(); + } + + @Override + public void onFailure(Exception e) { + assertNotNull(e); + } + }; + var plan = new ContinuePaginatedPlan(QueryId.None, buildCursor(Map.of("pageSize", "abc")), + paginatedQueryService, paginatedPlanCache, listener); + plan.execute(); + } + + @Test + public void explain_is_not_supported() { + assertThrows(UnsupportedOperationException.class, + () -> ContinuePaginatedPlan.None.explain(null)); + } +} diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/PaginatedPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/PaginatedPlanTest.java new file mode 100644 index 00000000000..f62391afad3 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/executor/execution/PaginatedPlanTest.java @@ -0,0 +1,99 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor.execution; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import org.apache.commons.lang3.NotImplementedException; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.analysis.Analyzer; +import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.DefaultExecutionEngine; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.QueryId; +import org.opensearch.sql.planner.Planner; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.physical.PhysicalPlan; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class PaginatedPlanTest { + + private static PaginatedQueryService paginatedQueryService; + + /** + * Initialize the mocks. + */ + @BeforeAll + public static void setUp() { + var analyzer = mock(Analyzer.class); + when(analyzer.analyze(any(), any())).thenReturn(mock(LogicalPlan.class)); + var planner = mock(Planner.class); + when(planner.plan(any())).thenReturn(mock(PhysicalPlan.class)); + paginatedQueryService = new PaginatedQueryService( + analyzer, new DefaultExecutionEngine(), planner); + } + + @Test + public void can_execute_plan() { + var listener = new ResponseListener() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + fail(); + } + }; + var plan = new PaginatedPlan(QueryId.None, mock(UnresolvedPlan.class), 10, + paginatedQueryService, listener); + plan.execute(); + } + + @Test + // Same as previous test, but with incomplete PaginatedQueryService + public void can_handle_error_while_executing_plan() { + var listener = new ResponseListener() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + fail(); + } + + @Override + public void onFailure(Exception e) { + assertNotNull(e); + } + }; + var plan = new PaginatedPlan(QueryId.None, mock(UnresolvedPlan.class), 10, + new PaginatedQueryService(null, new DefaultExecutionEngine(), null), listener); + plan.execute(); + } + + @Test + public void explain_is_not_supported() { + new PaginatedPlan(null, null, 0, null, null).explain(new ResponseListener<>() { + @Override + public void onResponse(ExecutionEngine.ExplainResponse response) { + fail(); + } + + @Override + public void onFailure(Exception e) { + assertTrue(e instanceof NotImplementedException); + } + }); + } +} diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index f6abaa0d541..72225f0884f 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -11,6 +11,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; import java.util.Optional; @@ -27,6 +28,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.legacy.plugin.UnsupportCursorRequestException; @ExtendWith(MockitoExtension.class) class QueryPlanFactoryTest { @@ -74,6 +76,12 @@ public void createFromExplainShouldSuccess() { assertTrue(queryExecution instanceof ExplainPlan); } + @Test + public void createFromCursorShouldSuccess() { + AbstractPlan queryExecution = factory.create("", queryListener); + assertTrue(queryExecution instanceof ContinuePaginatedPlan); + } + @Test public void createFromQueryWithoutQueryListenerShouldThrowException() { Statement query = new Query(plan, 0); @@ -110,4 +118,23 @@ public void noConsumerResponseChannel() { assertEquals( "[BUG] exception response should not sent to unexpected channel", exception.getMessage()); } + + @Test + public void createQueryWithFetchSizeWhichCanBePaged() { + when(paginatedPlanCache.canConvertToCursor(plan)).thenReturn(true); + factory = new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache); + Statement query = new Query(plan, 10); + AbstractPlan queryExecution = + factory.create(query, Optional.of(queryListener), Optional.empty()); + assertTrue(queryExecution instanceof PaginatedPlan); + } + + @Test + public void createQueryWithFetchSizeWhichCannotBePaged() { + when(paginatedPlanCache.canConvertToCursor(plan)).thenReturn(false); + factory = new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache); + Statement query = new Query(plan, 10); + assertThrows(UnsupportCursorRequestException.class, + () -> factory.create(query, Optional.of(queryListener), Optional.empty())); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java b/core/src/test/java/org/opensearch/sql/expression/serialization/DefaultExpressionSerializerTest.java similarity index 100% rename from opensearch/src/test/java/org/opensearch/sql/opensearch/storage/serialization/DefaultExpressionSerializerTest.java rename to core/src/test/java/org/opensearch/sql/expression/serialization/DefaultExpressionSerializerTest.java diff --git a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java index 017cfb60ea0..a426bbb1a94 100644 --- a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java @@ -33,6 +33,8 @@ import java.util.Map; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -50,6 +52,7 @@ import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.window.WindowDefinition; import org.opensearch.sql.expression.window.ranking.RowNumberFunction; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalRelation; @@ -62,24 +65,16 @@ import org.opensearch.sql.storage.write.TableWriteOperator; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class DefaultImplementorTest { - @Mock - private Expression filter; - - @Mock - private NamedAggregator aggregator; - - @Mock - private NamedExpression groupBy; - @Mock private Table table; private final DefaultImplementor implementor = new DefaultImplementor<>(); @Test - public void visitShouldReturnDefaultPhysicalOperator() { + public void visit_should_return_default_physical_operator() { String indexName = "test"; NamedExpression include = named("age", ref("age", INTEGER)); ReferenceExpression exclude = ref("name", STRING); @@ -157,14 +152,14 @@ public void visitShouldReturnDefaultPhysicalOperator() { } @Test - public void visitRelationShouldThrowException() { + public void visitRelation_should_throw_an_exception() { assertThrows(UnsupportedOperationException.class, () -> new LogicalRelation("test", table).accept(implementor, null)); } @SuppressWarnings({"rawtypes", "unchecked"}) @Test - public void visitWindowOperatorShouldReturnPhysicalWindowOperator() { + public void visitWindowOperator_should_return_PhysicalWindowOperator() { NamedExpression windowFunction = named(new RowNumberFunction()); WindowDefinition windowDefinition = new WindowDefinition( Collections.singletonList(ref("state", STRING)), @@ -204,7 +199,7 @@ public void visitWindowOperatorShouldReturnPhysicalWindowOperator() { } @Test - public void visitTableScanBuilderShouldBuildTableScanOperator() { + public void visitTableScanBuilder_should_build_TableScanOperator() { TableScanOperator tableScanOperator = Mockito.mock(TableScanOperator.class); TableScanBuilder tableScanBuilder = new TableScanBuilder() { @Override @@ -216,7 +211,7 @@ public TableScanOperator build() { } @Test - public void visitTableWriteBuilderShouldBuildTableWriteOperator() { + public void visitTableWriteBuilder_should_build_TableWriteOperator() { LogicalPlan child = values(); TableWriteOperator tableWriteOperator = Mockito.mock(TableWriteOperator.class); TableWriteBuilder logicalPlan = new TableWriteBuilder(child) { @@ -227,4 +222,11 @@ public TableWriteOperator build(PhysicalPlan child) { }; assertEquals(tableWriteOperator, logicalPlan.accept(implementor, null)); } + + @Test + public void visitPaginate_should_build_PaginateOperator_and_keep_page_size() { + var paginate = new LogicalPaginate(42, List.of(values())); + var plan = paginate.accept(implementor, null); + assertEquals(paginate.getPageSize(), ((PaginateOperator) plan).getPageSize()); + } } diff --git a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java index 341bcbc29e4..c9d74fa8712 100644 --- a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java @@ -8,21 +8,23 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.mock; import static org.opensearch.sql.expression.DSL.named; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.util.Collections; -import java.util.HashMap; +import java.util.List; import java.util.Map; -import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.commons.lang3.tuple.Pair; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.sql.ast.expression.DataType; -import org.opensearch.sql.ast.expression.Literal; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.ast.tree.Sort.SortOption; import org.opensearch.sql.data.model.ExprValueUtils; @@ -42,20 +44,24 @@ /** * Todo. Temporary added for UT coverage, Will be removed. */ -@ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class LogicalPlanNodeVisitorTest { - @Mock - Expression expression; - @Mock - ReferenceExpression ref; - @Mock - Aggregator aggregator; - @Mock - Table table; + static Expression expression; + static ReferenceExpression ref; + static Aggregator aggregator; + static Table table; + + @BeforeAll + private static void initMocks() { + expression = mock(Expression.class); + ref = mock(ReferenceExpression.class); + aggregator = mock(Aggregator.class); + table = mock(Table.class); + } @Test - public void logicalPlanShouldTraversable() { + public void logical_plan_should_be_traversable() { LogicalPlan logicalPlan = LogicalPlanDSL.rename( LogicalPlanDSL.aggregation( @@ -72,119 +78,57 @@ public void logicalPlanShouldTraversable() { assertEquals(5, result); } - @Test - public void testAbstractPlanNodeVisitorShouldReturnNull() { + @SuppressWarnings("unchecked") + private static Stream getLogicalPlansForVisitorTest() { LogicalPlan relation = LogicalPlanDSL.relation("schema", table); - assertNull(relation.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan tableScanBuilder = new TableScanBuilder() { @Override public TableScanOperator build() { return null; } }; - assertNull(tableScanBuilder.accept(new LogicalPlanNodeVisitor() { - }, null)); - - LogicalPlan write = LogicalPlanDSL.write(null, table, Collections.emptyList()); - assertNull(write.accept(new LogicalPlanNodeVisitor() { - }, null)); - TableWriteBuilder tableWriteBuilder = new TableWriteBuilder(null) { @Override public TableWriteOperator build(PhysicalPlan child) { return null; } }; - assertNull(tableWriteBuilder.accept(new LogicalPlanNodeVisitor() { - }, null)); - + LogicalPlan write = LogicalPlanDSL.write(null, table, Collections.emptyList()); LogicalPlan filter = LogicalPlanDSL.filter(relation, expression); - assertNull(filter.accept(new LogicalPlanNodeVisitor() { - }, null)); - - LogicalPlan aggregation = - LogicalPlanDSL.aggregation( - filter, ImmutableList.of(DSL.named("avg", aggregator)), ImmutableList.of(DSL.named( - "group", expression))); - assertNull(aggregation.accept(new LogicalPlanNodeVisitor() { - }, null)); - + LogicalPlan aggregation = LogicalPlanDSL.aggregation( + filter, ImmutableList.of(DSL.named("avg", aggregator)), ImmutableList.of(DSL.named( + "group", expression))); LogicalPlan rename = LogicalPlanDSL.rename(aggregation, ImmutableMap.of(ref, ref)); - assertNull(rename.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan project = LogicalPlanDSL.project(relation, named("ref", ref)); - assertNull(project.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan remove = LogicalPlanDSL.remove(relation, ref); - assertNull(remove.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan eval = LogicalPlanDSL.eval(relation, Pair.of(ref, expression)); - assertNull(eval.accept(new LogicalPlanNodeVisitor() { - }, null)); - - LogicalPlan sort = LogicalPlanDSL.sort(relation, - Pair.of(SortOption.DEFAULT_ASC, expression)); - assertNull(sort.accept(new LogicalPlanNodeVisitor() { - }, null)); - + LogicalPlan sort = LogicalPlanDSL.sort(relation, Pair.of(SortOption.DEFAULT_ASC, expression)); LogicalPlan dedup = LogicalPlanDSL.dedupe(relation, 1, false, false, expression); - assertNull(dedup.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan window = LogicalPlanDSL.window(relation, named(expression), new WindowDefinition( ImmutableList.of(ref), ImmutableList.of(Pair.of(SortOption.DEFAULT_ASC, expression)))); - assertNull(window.accept(new LogicalPlanNodeVisitor() { - }, null)); - LogicalPlan rareTopN = LogicalPlanDSL.rareTopN( relation, CommandType.TOP, ImmutableList.of(expression), expression); - assertNull(rareTopN.accept(new LogicalPlanNodeVisitor() { - }, null)); - - Map args = new HashMap<>(); LogicalPlan highlight = new LogicalHighlight(filter, - new LiteralExpression(ExprValueUtils.stringValue("fieldA")), args); - assertNull(highlight.accept(new LogicalPlanNodeVisitor() { - }, null)); - - LogicalPlan mlCommons = new LogicalMLCommons(LogicalPlanDSL.relation("schema", table), - "kmeans", - ImmutableMap.builder() - .put("centroids", new Literal(3, DataType.INTEGER)) - .put("iterations", new Literal(3, DataType.DOUBLE)) - .put("distance_type", new Literal(null, DataType.STRING)) - .build()); - assertNull(mlCommons.accept(new LogicalPlanNodeVisitor() { - }, null)); - - LogicalPlan ad = new LogicalAD(LogicalPlanDSL.relation("schema", table), - new HashMap() {{ - put("shingle_size", new Literal(8, DataType.INTEGER)); - put("time_decay", new Literal(0.0001, DataType.DOUBLE)); - put("time_field", new Literal(null, DataType.STRING)); - } - }); - assertNull(ad.accept(new LogicalPlanNodeVisitor() { - }, null)); + new LiteralExpression(ExprValueUtils.stringValue("fieldA")), Map.of()); + LogicalPlan mlCommons = new LogicalMLCommons(relation, "kmeans", Map.of()); + LogicalPlan ad = new LogicalAD(relation, Map.of()); + LogicalPlan ml = new LogicalML(relation, Map.of()); + LogicalPlan paginate = new LogicalPaginate(42, List.of(relation)); + + return Stream.of( + relation, tableScanBuilder, write, tableWriteBuilder, filter, aggregation, rename, project, + remove, eval, sort, dedup, window, rareTopN, highlight, mlCommons, ad, ml, paginate + ).map(Arguments::of); + } - LogicalPlan ml = new LogicalML(LogicalPlanDSL.relation("schema", table), - new HashMap() {{ - put("action", new Literal("train", DataType.STRING)); - put("algorithm", new Literal("rcf", DataType.STRING)); - put("shingle_size", new Literal(8, DataType.INTEGER)); - put("time_decay", new Literal(0.0001, DataType.DOUBLE)); - put("time_field", new Literal(null, DataType.STRING)); - } - }); - assertNull(ml.accept(new LogicalPlanNodeVisitor() { + @ParameterizedTest + @MethodSource("getLogicalPlansForVisitorTest") + public void abstract_plan_node_visitor_should_return_null(LogicalPlan plan) { + assertNull(plan.accept(new LogicalPlanNodeVisitor() { }, null)); } + private static class NodesCount extends LogicalPlanNodeVisitor { @Override public Integer visitRelation(LogicalRelation plan, Object context) { @@ -195,32 +139,28 @@ public Integer visitRelation(LogicalRelation plan, Object context) { public Integer visitFilter(LogicalFilter plan, Object context) { return 1 + plan.getChild().stream() - .map(child -> child.accept(this, context)) - .collect(Collectors.summingInt(Integer::intValue)); + .map(child -> child.accept(this, context)).mapToInt(Integer::intValue).sum(); } @Override public Integer visitAggregation(LogicalAggregation plan, Object context) { return 1 + plan.getChild().stream() - .map(child -> child.accept(this, context)) - .collect(Collectors.summingInt(Integer::intValue)); + .map(child -> child.accept(this, context)).mapToInt(Integer::intValue).sum(); } @Override public Integer visitRename(LogicalRename plan, Object context) { return 1 + plan.getChild().stream() - .map(child -> child.accept(this, context)) - .collect(Collectors.summingInt(Integer::intValue)); + .map(child -> child.accept(this, context)).mapToInt(Integer::intValue).sum(); } @Override public Integer visitRareTopN(LogicalRareTopN plan, Object context) { return 1 + plan.getChild().stream() - .map(child -> child.accept(this, context)) - .collect(Collectors.summingInt(Integer::intValue)); + .map(child -> child.accept(this, context)).mapToInt(Integer::intValue).sum(); } } } diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index 7516aa18091..9d3620d1116 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -7,8 +7,11 @@ package org.opensearch.sql.planner.optimizer; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.when; import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; import static org.opensearch.sql.data.model.ExprValueUtils.longValue; @@ -26,9 +29,12 @@ import com.google.common.collect.ImmutableList; import java.util.Collections; +import java.util.List; import java.util.Map; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -38,13 +44,16 @@ import org.opensearch.sql.ast.tree.Sort; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.DSL; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalRelation; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class LogicalPlanOptimizerTest { @Mock @@ -55,7 +64,7 @@ class LogicalPlanOptimizerTest { @BeforeEach void setUp() { - when(table.createScanBuilder()).thenReturn(tableScanBuilder); + lenient().when(table.createScanBuilder()).thenReturn(tableScanBuilder); } /** @@ -255,7 +264,6 @@ void table_scan_builder_support_highlight_push_down_can_apply_its_rule() { @Test void table_not_support_scan_builder_should_not_be_impact() { - Mockito.reset(table, tableScanBuilder); Table table = new Table() { @Override public Map getFieldTypes() { @@ -276,7 +284,6 @@ public PhysicalPlan implement(LogicalPlan plan) { @Test void table_support_write_builder_should_be_replaced() { - Mockito.reset(table, tableScanBuilder); TableWriteBuilder writeBuilder = Mockito.mock(TableWriteBuilder.class); when(table.createWriteBuilder(any())).thenReturn(writeBuilder); @@ -288,7 +295,6 @@ void table_support_write_builder_should_be_replaced() { @Test void table_not_support_write_builder_should_report_error() { - Mockito.reset(table, tableScanBuilder); Table table = new Table() { @Override public Map getFieldTypes() { @@ -305,6 +311,54 @@ public PhysicalPlan implement(LogicalPlan plan) { () -> table.createWriteBuilder(null)); } + @Test + void paged_table_scan_builder_support_project_push_down_can_apply_its_rule() { + when(tableScanBuilder.pushDownProject(any())).thenReturn(true); + when(table.createPagedScanBuilder(anyInt())).thenReturn(tableScanBuilder); + + var relation = new LogicalRelation("schema", table); + relation.setPageSize(anyInt()); + + assertEquals( + tableScanBuilder, + LogicalPlanOptimizer.paginationCreate().optimize(project(relation)) + ); + } + + @Test + void push_page_size() { + var relation = new LogicalRelation("schema", table); + var paginate = new LogicalPaginate(42, List.of(project(relation))); + assertNull(relation.getPageSize()); + LogicalPlanOptimizer.paginationCreate().optimize(paginate); + assertEquals(42, relation.getPageSize()); + } + + @Test + void push_page_size_noop_if_no_relation() { + var paginate = new LogicalPaginate(42, List.of(project(values()))); + LogicalPlanOptimizer.paginationCreate().optimize(paginate); + } + + @Test + void push_page_size_noop_if_no_sub_plans() { + var paginate = new LogicalPaginate(42, List.of()); + LogicalPlanOptimizer.paginationCreate().optimize(paginate); + } + + @Test + void table_scan_builder_support_offset_push_down_can_apply_its_rule() { + // next line is noop, added for coverage only + lenient().when(tableScanBuilder.pushDownOffset(anyInt())).thenReturn(true); + when(table.createPagedScanBuilder(anyInt())).thenReturn(tableScanBuilder); + + var optimized = LogicalPlanOptimizer.paginationCreate() + .optimize(new LogicalPaginate(42, List.of(project(relation("schema", table))))); + // `optimized` structure: LogicalPaginate -> LogicalProject -> TableScanBuilder + // LogicalRelation replaced by a TableScanBuilder instance + assertEquals(tableScanBuilder, optimized.getChild().get(0).getChild().get(0)); + } + private LogicalPlan optimize(LogicalPlan plan) { final LogicalPlanOptimizer optimizer = LogicalPlanOptimizer.create(); final LogicalPlan optimize = optimizer.optimize(plan); diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/pattern/PatternsTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/pattern/PatternsTest.java index 9f90fd8d055..1fd572e7daf 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/pattern/PatternsTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/pattern/PatternsTest.java @@ -6,35 +6,49 @@ package org.opensearch.sql.planner.optimizer.pattern; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import java.util.Collections; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; -@ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class PatternsTest { - @Mock - LogicalPlan plan; - @Test void source_is_empty() { + var plan = mock(LogicalPlan.class); when(plan.getChild()).thenReturn(Collections.emptyList()); - assertFalse(Patterns.source().getFunction().apply(plan).isPresent()); - assertFalse(Patterns.source(null).getProperty().getFunction().apply(plan).isPresent()); + assertAll( + () -> assertFalse(Patterns.source().getFunction().apply(plan).isPresent()), + () -> assertFalse(Patterns.source(null).getProperty().getFunction().apply(plan).isPresent()) + ); } @Test void table_is_empty() { - plan = Mockito.mock(LogicalFilter.class); - assertFalse(Patterns.table().getFunction().apply(plan).isPresent()); - assertFalse(Patterns.writeTable().getFunction().apply(plan).isPresent()); + var plan = mock(LogicalFilter.class); + assertAll( + () -> assertFalse(Patterns.table().getFunction().apply(plan).isPresent()), + () -> assertFalse(Patterns.writeTable().getFunction().apply(plan).isPresent()) + ); + } + + @Test + void pagination() { + assertAll( + () -> assertTrue(Patterns.pagination().getFunction() + .apply(mock(LogicalPaginate.class)).isPresent()), + () -> assertFalse(Patterns.pagination().getFunction() + .apply(mock(LogicalFilter.class)).isPresent()) + ); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java new file mode 100644 index 00000000000..e91766b1a26 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java @@ -0,0 +1,86 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + + +package org.opensearch.sql.planner.physical; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.project; + +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.expression.DSL; +import org.opensearch.sql.planner.PaginateOperator; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class PaginateOperatorTest extends PhysicalPlanTestBase { + + @Test + public void accept() { + var visitor = new PhysicalPlanNodeVisitor() {}; + assertNull(new PaginateOperator(null, 42).accept(visitor, null)); + } + + @Test + public void hasNext_a_page() { + var plan = mock(PhysicalPlan.class); + when(plan.hasNext()).thenReturn(true); + when(plan.next()).thenReturn(null); + var paginate = new PaginateOperator(plan, 1, 1); + assertTrue(paginate.hasNext()); + paginate.next(); + assertFalse(paginate.hasNext()); + } + + @Test + public void hasNext_no_more_entries() { + var plan = mock(PhysicalPlan.class); + when(plan.hasNext()).thenReturn(false); + var paginate = new PaginateOperator(plan, 1, 1); + assertFalse(paginate.hasNext()); + } + + @Test + public void getChild() { + var plan = mock(PhysicalPlan.class); + var paginate = new PaginateOperator(plan, 1); + assertSame(plan, paginate.getChild().get(0)); + } + + @Test + public void open() { + var plan = mock(PhysicalPlan.class); + doNothing().when(plan).open(); + new PaginateOperator(plan, 1).open(); + verify(plan, times(1)).open(); + } + + @Test + public void schema() { + PhysicalPlan project = project(null, + DSL.named("response", DSL.ref("response", INTEGER)), + DSL.named("action", DSL.ref("action", STRING), "act")); + assertEquals(project.schema(), new PaginateOperator(project, 42).schema()); + } + + @Test + public void schema_assert() { + assertThrows(Throwable.class, + () -> new PaginateOperator(mock(PhysicalPlan.class), 42).schema()); + } +} diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java index 735b914d3e1..27d7a43bfdc 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java @@ -26,6 +26,7 @@ import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.window.WindowDefinition; +import org.opensearch.sql.planner.PaginateOperator; /** * Todo, testing purpose, delete later. @@ -158,6 +159,14 @@ public void test_visitML() { assertNull(physicalPlanNodeVisitor.visitML(plan, null)); } + @Test + public void test_visitPaginate() { + PhysicalPlanNodeVisitor physicalPlanNodeVisitor = + new PhysicalPlanNodeVisitor() {}; + + assertNull(physicalPlanNodeVisitor.visitPaginate(new PaginateOperator(plan, 42), null)); + } + public static class PhysicalPlanPrinter extends PhysicalPlanNodeVisitor { public String print(PhysicalPlan node) { diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java index 1cc7d5532fb..ec950e6016b 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/RemoveOperatorTest.java @@ -113,7 +113,7 @@ public void remove_nothing_with_none_tuple_value() { @Test public void invalid_to_retrieve_schema_from_remove() { - PhysicalPlan plan = remove(inputPlan, DSL.ref("response", STRING), DSL.ref("referer", STRING)); + PhysicalPlan plan = remove(inputPlan); IllegalStateException exception = assertThrows(IllegalStateException.class, () -> plan.schema()); assertEquals( diff --git a/core/src/test/java/org/opensearch/sql/storage/StorageEngineTest.java b/core/src/test/java/org/opensearch/sql/storage/StorageEngineTest.java index 0e969c6dac3..9c96459d061 100644 --- a/core/src/test/java/org/opensearch/sql/storage/StorageEngineTest.java +++ b/core/src/test/java/org/opensearch/sql/storage/StorageEngineTest.java @@ -13,11 +13,16 @@ public class StorageEngineTest { - @Test void testFunctionsMethod() { StorageEngine k = (dataSourceSchemaName, tableName) -> null; Assertions.assertEquals(Collections.emptyList(), k.getFunctions()); } + @Test + void getTableScan() { + StorageEngine k = (dataSourceSchemaName, tableName) -> null; + Assertions.assertThrows(UnsupportedOperationException.class, + () -> k.getTableScan("indexName", "scrollId")); + } } diff --git a/core/src/test/java/org/opensearch/sql/storage/TableTest.java b/core/src/test/java/org/opensearch/sql/storage/TableTest.java new file mode 100644 index 00000000000..2a2b5550145 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/storage/TableTest.java @@ -0,0 +1,25 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.storage; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.withSettings; + +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.mockito.invocation.InvocationOnMock; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class TableTest { + + @Test + public void createPagedScanBuilder_throws() { + var table = mock(Table.class, withSettings().defaultAnswer(InvocationOnMock::callRealMethod)); + assertThrows(Throwable.class, () -> table.createPagedScanBuilder(0)); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index 72f21049490..0ce3a077077 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -93,8 +93,8 @@ public boolean pushDownProject(LogicalProject project) { } @Override - public void pushDownOffset(int i) { - delegate.pushDownOffset(i); + public boolean pushDownOffset(int i) { + return delegate.pushDownOffset(i); } @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 7f01bc605b2..2441b2d3b21 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 @@ -45,6 +45,7 @@ import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; @@ -105,6 +106,35 @@ public void onFailure(Exception e) { assertTrue(plan.hasClosed); } + @Test + void executeWithCursor() { + List expected = + Arrays.asList( + tupleValue(of("name", "John", "age", 20)), tupleValue(of("name", "Allen", "age", 30))); + FakePaginatePlan plan = new FakePaginatePlan(new FakePhysicalPlan(expected.iterator()), 10, 0); + when(protector.protect(plan)).thenReturn(plan); + + OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, + PaginatedPlanCache.None); + List actual = new ArrayList<>(); + executor.execute( + plan, + new ResponseListener() { + @Override + public void onResponse(QueryResponse response) { + actual.addAll(response.getResults()); + assertTrue(response.getCursor().toString().startsWith("n:")); + } + + @Override + public void onFailure(Exception e) { + fail("Error occurred during execution", e); + } + }); + + assertEquals(expected, actual); + } + @Test void executeWithFailure() { PhysicalPlan plan = mock(PhysicalPlan.class); @@ -214,6 +244,54 @@ public void onFailure(Exception e) { assertTrue(plan.hasClosed); } + private static class FakePaginatePlan extends PaginateOperator { + private final PhysicalPlan input; + private final int pageSize; + private final int pageIndex; + + public FakePaginatePlan(PhysicalPlan input, int pageSize, int pageIndex) { + super(input, pageSize, pageIndex); + this.input = input; + this.pageSize = pageSize; + this.pageIndex = pageIndex; + } + + @Override + public void open() { + input.open(); + } + + @Override + public void close() { + input.close(); + } + + @Override + public void add(Split split) { + input.add(split); + } + + @Override + public boolean hasNext() { + return input.hasNext(); + } + + @Override + public ExprValue next() { + return input.next(); + } + + @Override + public ExecutionEngine.Schema schema() { + return input.schema(); + } + + @Override + public String toCursor() { + return "FakePaginatePlan"; + } + } + @RequiredArgsConstructor private static class FakePhysicalPlan extends TableScanOperator { private final Iterator it; 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 62718800b19..cbe8eb8dfbd 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -57,6 +57,9 @@ class SQLServiceTest { @Mock private PaginatedPlanCache paginatedPlanCache; + @Mock + private PaginatedQueryService paginatedQueryService; + @BeforeEach public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); From 70ccfcbfb7910d13c9706c634284c634f21a6c22 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 9 Mar 2023 18:41:00 -0800 Subject: [PATCH 19/46] Pagination, phase 1: Add unit tests for SQL module with coverage. (#239) * Add unit tests for SQL module with coverage. Signed-off-by: Yury-Fridlyand * Update sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java Signed-off-by: Yury-Fridlyand Co-authored-by: GabeFernandez310 --------- Signed-off-by: Yury-Fridlyand Co-authored-by: GabeFernandez310 --- .../org/opensearch/sql/sql/SQLService.java | 11 ++ .../sql/sql/domain/SQLQueryRequest.java | 40 ++-- .../opensearch/sql/sql/SQLServiceTest.java | 66 ++++--- .../sql/sql/domain/SQLQueryRequestTest.java | 178 ++++++++++++++---- 4 files changed, 210 insertions(+), 85 deletions(-) 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 908f03e4f26..2acc44de63e 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -13,6 +13,7 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.execution.AbstractPlan; import org.opensearch.sql.executor.execution.QueryPlanFactory; @@ -68,9 +69,19 @@ private AbstractPlan plan( if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. if (queryListener.isEmpty() && explainListener.isPresent()) { // explain request + // TODO explain should be processed inside the plan explainListener.get().onFailure(new UnsupportedOperationException( "`explain` request for cursor requests is not supported. " + "Use `explain` for the initial query request.")); + return new AbstractPlan(QueryId.queryId()) { + @Override + public void execute() { + } + + @Override + public void explain(ResponseListener listener) { + } + }; } // non-explain request return queryExecutionFactory.create(request.getCursor().get(), queryListener.get()); 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 62db130fc2b..0d15abf54a5 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 @@ -6,13 +6,12 @@ package org.opensearch.sql.sql.domain; -import com.google.common.base.Strings; -import com.google.common.collect.ImmutableSet; import java.util.Collections; import java.util.Locale; import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.stream.Stream; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; @@ -64,7 +63,7 @@ public class SQLQueryRequest { @Accessors(fluent = true) private boolean sanitize = true; - private String cursor = ""; + private String cursor; /** * Constructor of SQLQueryRequest that passes request params. @@ -77,8 +76,7 @@ public SQLQueryRequest(JSONObject jsonContent, String query, String path, this.params = params; this.format = getFormat(params); this.sanitize = shouldSanitize(params); - // TODO hack - this.cursor = cursor == null ? "" : cursor; + this.cursor = cursor; } /** @@ -86,20 +84,27 @@ public SQLQueryRequest(JSONObject jsonContent, String query, String path, * 1.Only supported fields present in request body, ex. "filter" and "cursor" are not supported * 2.Response format is default or can be supported. * - * @return true if supported. + * @return true if supported. */ public boolean isSupported() { - return (isCursor() || isOnlySupportedFieldInPayload()) - && isSupportedFormat(); + var noCursor = !isCursor(); + var noQuery = query == null; + var noParams = params.isEmpty(); + var noContent = jsonContent == null || jsonContent.isEmpty(); + + return ((!noCursor && noQuery && noParams && noContent) // if cursor is given, but other things + || (noCursor && !noQuery)) // or if cursor is not given, but query + && isOnlySupportedFieldInPayload() // and request has supported fields only + && isSupportedFormat(); // and request is in supported format } private boolean isCursor() { - return cursor != null && cursor.isEmpty() == false; + return cursor != null && !cursor.isEmpty(); } /** * Check if request is to explain rather than execute the query. - * @return true if it is a explain request + * @return true if it is an explain request */ public boolean isExplainRequest() { return path.endsWith("/_explain"); @@ -122,13 +127,8 @@ private boolean isOnlySupportedFieldInPayload() { return jsonContent == null || SUPPORTED_FIELDS.containsAll(jsonContent.keySet()); } - public Optional getCursor() { - return cursor != "" ? Optional.of(cursor) : Optional.empty(); - } - - public boolean mayReturnCursor() { - return cursor != "" || getFetchSize() > 0; + return Optional.ofNullable(cursor); } public int getFetchSize() { @@ -136,15 +136,11 @@ public int getFetchSize() { } private boolean isSupportedFormat() { - return Strings.isNullOrEmpty(format) || "jdbc".equalsIgnoreCase(format) - || "csv".equalsIgnoreCase(format) || "raw".equalsIgnoreCase(format); + return Stream.of("csv", "jdbc", "raw").anyMatch(format::equalsIgnoreCase); } private String getFormat(Map params) { - if (params.containsKey(QUERY_PARAMS_FORMAT)) { - return params.get(QUERY_PARAMS_FORMAT); - } - return "jdbc"; + return params.getOrDefault(QUERY_PARAMS_FORMAT, "jdbc"); } private boolean shouldSanitize(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 cbe8eb8dfbd..cb1159bc17f 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -7,16 +7,19 @@ package org.opensearch.sql.sql; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.opensearch.sql.executor.ExecutionEngine.QueryResponse; -import java.util.Collections; +import java.util.Map; import java.util.concurrent.TimeUnit; import org.json.JSONObject; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -30,11 +33,11 @@ import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.PaginatedQueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; -import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.sql.domain.SQLQueryRequest; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class SQLServiceTest { private static String QUERY = "/_plugins/_sql"; @@ -51,9 +54,6 @@ class SQLServiceTest { @Mock private PaginatedQueryService paginatedQueryService; - @Mock - private ExecutionEngine.Schema schema; - @Mock private PaginatedPlanCache paginatedPlanCache; @@ -73,13 +73,7 @@ public void cleanup() throws InterruptedException { } @Test - public void canExecuteSqlQuery() { - doAnswer(invocation -> { - ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); - return null; - }).when(queryService).execute(any(), any()); - + public void can_execute_sql_query() { sqlService.execute( new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "jdbc"), new ResponseListener<>() { @@ -96,13 +90,24 @@ public void onFailure(Exception e) { } @Test - public void canExecuteCsvFormatRequest() { - doAnswer(invocation -> { - ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); - return null; - }).when(queryService).execute(any(), any()); + public void can_execute_cursor_query() { + sqlService.execute( + new SQLQueryRequest(new JSONObject(), null, QUERY, Map.of("format", "jdbc"), "n:cursor"), + new ResponseListener<>() { + @Override + public void onResponse(QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + fail(e); + } + }); + } + @Test + public void can_execute_csv_format_request() { sqlService.execute( new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "csv"), new ResponseListener() { @@ -119,7 +124,7 @@ public void onFailure(Exception e) { } @Test - public void canExplainSqlQuery() { + public void can_explain_sql_query() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); listener.onResponse(new ExplainResponse(new ExplainResponseNode("Test"))); @@ -141,7 +146,25 @@ public void onFailure(Exception e) { } @Test - public void canCaptureErrorDuringExecution() { + public void cannot_explain_cursor_query() { + sqlService.explain(new SQLQueryRequest(new JSONObject(), null, EXPLAIN, + Map.of("format", "jdbc"), "n:cursor"), + new ResponseListener() { + @Override + public void onResponse(ExplainResponse response) { + fail(response.toString()); + } + + @Override + public void onFailure(Exception e) { + assertTrue(e.getMessage() + .contains("`explain` request for cursor requests is not supported.")); + } + }); + } + + @Test + public void can_capture_error_during_execution() { sqlService.execute( new SQLQueryRequest(new JSONObject(), "SELECT", QUERY, ""), new ResponseListener() { @@ -158,7 +181,7 @@ public void onFailure(Exception e) { } @Test - public void canCaptureErrorDuringExplain() { + public void can_capture_error_during_explain() { sqlService.explain( new SQLQueryRequest(new JSONObject(), "SELECT", EXPLAIN, ""), new ResponseListener() { @@ -173,5 +196,4 @@ public void onFailure(Exception e) { } }); } - } 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 cf14d40e015..3a4e4ae09ff 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 @@ -6,38 +6,43 @@ package org.opensearch.sql.sql.domain; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.ImmutableMap; +import java.util.HashMap; import java.util.Map; import org.json.JSONObject; -import org.junit.jupiter.api.Assumptions; -import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.opensearch.sql.protocol.response.format.Format; +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class SQLQueryRequestTest { @Test - public void shouldSupportQuery() { + public void should_support_query() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1").build(); assertTrue(request.isSupported()); } @Test - public void shouldSupportQueryWithJDBCFormat() { + public void should_support_query_with_JDBC_format() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1") .format("jdbc") .build(); - assertTrue(request.isSupported()); - assertEquals(request.format(), Format.JDBC); + assertAll( + () -> assertTrue(request.isSupported()), + () -> assertEquals(request.format(), Format.JDBC) + ); } @Test - public void shouldSupportQueryWithQueryFieldOnly() { + public void should_support_query_with_query_field_only() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1") .jsonContent("{\"query\": \"SELECT 1\"}") @@ -46,16 +51,32 @@ public void shouldSupportQueryWithQueryFieldOnly() { } @Test - public void shouldSupportQueryWithParameters() { - SQLQueryRequest request = + public void should_support_query_with_parameters() { + SQLQueryRequest requestWithContent = SQLQueryRequestBuilder.request("SELECT 1") .jsonContent("{\"query\": \"SELECT 1\", \"parameters\":[]}") .build(); - assertTrue(request.isSupported()); + SQLQueryRequest requestWithParams = + SQLQueryRequestBuilder.request("SELECT 1") + .params(Map.of("one", "two")) + .build(); + assertAll( + () -> assertTrue(requestWithContent.isSupported()), + () -> assertTrue(requestWithParams.isSupported()) + ); } @Test - public void shouldSupportQueryWithZeroFetchSize() { + public void should_support_query_without_parameters() { + SQLQueryRequest requestWithNoParams = + SQLQueryRequestBuilder.request("SELECT 1") + .params(Map.of()) + .build(); + assertTrue(requestWithNoParams.isSupported()); + } + + @Test + public void should_support_query_with_zero_fetch_size() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1") .jsonContent("{\"query\": \"SELECT 1\", \"fetch_size\": 0}") @@ -64,7 +85,7 @@ public void shouldSupportQueryWithZeroFetchSize() { } @Test - public void shouldSupportQueryWithParametersAndZeroFetchSize() { + public void should_support_query_with_parameters_and_zero_fetch_size() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1") .jsonContent("{\"query\": \"SELECT 1\", \"fetch_size\": 0, \"parameters\":[]}") @@ -73,71 +94,143 @@ public void shouldSupportQueryWithParametersAndZeroFetchSize() { } @Test - public void shouldSupportExplain() { + public void should_support_explain() { SQLQueryRequest explainRequest = SQLQueryRequestBuilder.request("SELECT 1") .path("_plugins/_sql/_explain") .build(); - assertTrue(explainRequest.isExplainRequest()); - assertTrue(explainRequest.isSupported()); + + assertAll( + () -> assertTrue(explainRequest.isExplainRequest()), + () -> assertTrue(explainRequest.isSupported()) + ); } @Test - @Disabled("SQLQueryRequest does support cursor requests") - public void shouldNotSupportCursorRequest() { + public void should_support_cursor_request() { SQLQueryRequest fetchSizeRequest = SQLQueryRequestBuilder.request("SELECT 1") .jsonContent("{\"query\": \"SELECT 1\", \"fetch_size\": 5}") .build(); - assertFalse(fetchSizeRequest.isSupported()); SQLQueryRequest cursorRequest = + SQLQueryRequestBuilder.request(null) + .cursor("abcdefgh...") + .build(); + + assertAll( + () -> assertTrue(fetchSizeRequest.isSupported()), + () -> assertTrue(cursorRequest.isSupported()) + ); + } + + @Test + public void should_not_support_request_with_empty_cursor() { + SQLQueryRequest requestWithEmptyCursor = + SQLQueryRequestBuilder.request(null) + .cursor("") + .build(); + SQLQueryRequest requestWithNullCursor = + SQLQueryRequestBuilder.request(null) + .cursor(null) + .build(); + assertAll( + () -> assertFalse(requestWithEmptyCursor.isSupported()), + () -> assertFalse(requestWithNullCursor.isSupported()) + ); + } + + @Test + public void should_not_support_request_with_unknown_field() { + SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1") - .jsonContent("{\"cursor\": \"abcdefgh...\"}") + .jsonContent("{\"pewpew\": 42}") .build(); - assertFalse(cursorRequest.isSupported()); + assertFalse(request.isSupported()); } @Test - public void shouldUseJDBCFormatByDefault() { + public void should_not_support_request_with_cursor_and_something_else() { + SQLQueryRequest requestWithQuery = + SQLQueryRequestBuilder.request("SELECT 1") + .cursor("n:12356") + .build(); + SQLQueryRequest requestWithParams = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .params(Map.of("one", "two")) + .build(); + SQLQueryRequest requestWithFetchSize = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .jsonContent("{\"fetch_size\": 5}") + .build(); + SQLQueryRequest requestWithNoParams = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .params(Map.of()) + .build(); + SQLQueryRequest requestWithNoContent = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .jsonContent("{}") + .build(); + assertAll( + () -> assertFalse(requestWithQuery.isSupported()), + () -> assertFalse(requestWithParams.isSupported()), + () -> assertFalse(requestWithFetchSize.isSupported()), + () -> assertTrue(requestWithNoParams.isSupported()), + () -> assertTrue(requestWithNoContent.isSupported()) + ); + } + + @Test + public void should_use_JDBC_format_by_default() { SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1").params(ImmutableMap.of()).build(); assertEquals(request.format(), Format.JDBC); } @Test - public void shouldSupportCSVFormatAndSanitize() { + public void should_support_CSV_format_and_sanitize() { SQLQueryRequest csvRequest = SQLQueryRequestBuilder.request("SELECT 1") .format("csv") .build(); - assertTrue(csvRequest.isSupported()); - assertEquals(csvRequest.format(), Format.CSV); - assertTrue(csvRequest.sanitize()); + assertAll( + () -> assertTrue(csvRequest.isSupported()), + () -> assertEquals(csvRequest.format(), Format.CSV), + () -> assertTrue(csvRequest.sanitize()) + ); } @Test - public void shouldSkipSanitizeIfSetFalse() { + public void should_skip_sanitize_if_set_false() { ImmutableMap.Builder builder = ImmutableMap.builder(); Map params = builder.put("format", "csv").put("sanitize", "false").build(); SQLQueryRequest csvRequest = SQLQueryRequestBuilder.request("SELECT 1").params(params).build(); - assertEquals(csvRequest.format(), Format.CSV); - assertFalse(csvRequest.sanitize()); + assertAll( + () -> assertEquals(csvRequest.format(), Format.CSV), + () -> assertFalse(csvRequest.sanitize()) + ); } @Test - public void shouldNotSupportOtherFormat() { + public void should_not_support_other_format() { SQLQueryRequest csvRequest = SQLQueryRequestBuilder.request("SELECT 1") .format("other") .build(); - assertFalse(csvRequest.isSupported()); - assertThrows(IllegalArgumentException.class, csvRequest::format, - "response in other format is not supported."); + + assertAll( + () -> assertFalse(csvRequest.isSupported()), + () -> assertEquals("response in other format is not supported.", + assertThrows(IllegalArgumentException.class, csvRequest::format).getMessage()) + ); } @Test - public void shouldSupportRawFormat() { + public void should_support_raw_format() { SQLQueryRequest csvRequest = SQLQueryRequestBuilder.request("SELECT 1") .format("raw") @@ -153,7 +246,8 @@ private static class SQLQueryRequestBuilder { private String query; private String path = "_plugins/_sql"; private String format; - private Map params; + private String cursor; + private Map params = new HashMap<>(); static SQLQueryRequestBuilder request(String query) { SQLQueryRequestBuilder builder = new SQLQueryRequestBuilder(); @@ -181,15 +275,17 @@ SQLQueryRequestBuilder params(Map params) { return this; } + SQLQueryRequestBuilder cursor(String cursor) { + this.cursor = cursor; + return this; + } + SQLQueryRequest build() { - if (jsonContent == null) { - jsonContent = "{\"query\": \"" + query + "\"}"; - } - if (params != null) { - return new SQLQueryRequest(new JSONObject(jsonContent), query, path, params, - ""); + if (format != null) { + params.put("format", format); } - return new SQLQueryRequest(new JSONObject(jsonContent), query, path, format); + return new SQLQueryRequest(jsonContent == null ? null : new JSONObject(jsonContent), + query, path, params, cursor); } } From 803f50e1afec2acf8147623c0737408ffcae8671 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 9 Mar 2023 18:44:18 -0800 Subject: [PATCH 20/46] Pagination, phase 1: Add unit tests for `:opensearch` module with coverage. (#233) * Add UT for `:opensearch` module with full coverage, except `toCursor`. Signed-off-by: Yury-Fridlyand * Fix checkstyle. Signed-off-by: Yury-Fridlyand --------- Signed-off-by: Yury-Fridlyand --- .../sql/planner/PaginateOperator.java | 2 + .../sql/storage/read/TableScanBuilder.java | 3 - .../ContinueScrollRequest.java | 8 +- .../InitialPageRequestBuilder.java | 17 ++- .../request/OpenSearchQueryRequest.java | 8 +- .../request/OpenSearchRequestBuilder.java | 1 + .../PagedRequestBuilder.java | 5 +- .../SubsequentPageRequestBuilder.java | 17 ++- .../opensearch/storage/OpenSearchIndex.java | 6 +- .../storage/OpenSearchStorageEngine.java | 3 + .../{ => scan}/OpenSearchIndexScan.java | 2 +- ...OpenSearchIndexScanAggregationBuilder.java | 1 - .../scan/OpenSearchIndexScanBuilder.java | 6 - .../scan/OpenSearchIndexScanQueryBuilder.java | 5 - .../{ => scan}/OpenSearchPagedIndexScan.java | 6 +- .../OpenSearchPagedIndexScanBuilder.java} | 11 +- .../client/OpenSearchNodeClientTest.java | 12 +- .../client/OpenSearchRestClientTest.java | 9 +- .../OpenSearchExecutionEngineTest.java | 3 +- .../OpenSearchExecutionProtectorTest.java | 10 +- .../request/ContinueScrollRequestTest.java | 117 ++++++++++++++++++ .../InitialPageRequestBuilderTest.java | 107 ++++++++++++++++ .../request/OpenSearchQueryRequestTest.java | 3 +- .../request/OpenSearchScrollRequestTest.java | 3 + .../SubsequentPageRequestBuilderTest.java | 48 +++++++ .../storage/OpenSearchIndexTest.java | 16 +++ .../storage/OpenSearchStorageEngineTest.java | 37 +++++- .../OpenSearchIndexScanOptimizationTest.java | 1 - .../{ => scan}/OpenSearchIndexScanTest.java | 20 +-- .../scan/OpenSearchPagedIndexScanTest.java | 117 ++++++++++++++++++ 30 files changed, 521 insertions(+), 83 deletions(-) rename opensearch/src/main/java/org/opensearch/sql/opensearch/{storage => request}/ContinueScrollRequest.java (90%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/{storage => request}/InitialPageRequestBuilder.java (84%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/{storage => request}/PagedRequestBuilder.java (52%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/{storage => request}/SubsequentPageRequestBuilder.java (56%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/storage/{ => scan}/OpenSearchIndexScan.java (98%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/storage/{ => scan}/OpenSearchPagedIndexScan.java (88%) rename opensearch/src/main/java/org/opensearch/sql/opensearch/storage/{OpenSearchPagedScanBuilder.java => scan/OpenSearchPagedIndexScanBuilder.java} (58%) create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilderTest.java rename opensearch/src/test/java/org/opensearch/sql/opensearch/storage/{ => scan}/OpenSearchIndexScanTest.java (96%) create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java index d867674e053..68e94ac6b21 100644 --- a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java @@ -6,6 +6,7 @@ package org.opensearch.sql.planner; import java.util.List; +import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; import org.opensearch.sql.data.model.ExprValue; @@ -15,6 +16,7 @@ import org.opensearch.sql.planner.physical.ProjectOperator; @RequiredArgsConstructor +@EqualsAndHashCode(callSuper = false) public class PaginateOperator extends PhysicalPlan { @Getter private final PhysicalPlan input; diff --git a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java index e05cfad94ee..c0fdf36e709 100644 --- a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java +++ b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java @@ -108,7 +108,4 @@ public boolean pushDownHighlight(LogicalHighlight highlight) { public R accept(LogicalPlanNodeVisitor visitor, C context) { return visitor.visitTableScanBuilder(this, context); } - - public void pushDownOffset(int i) { - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java similarity index 90% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java index b33b7fd9a33..1ec5960b21d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.request; import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.DEFAULT_SCROLL_TIMEOUT; @@ -17,9 +17,9 @@ import org.opensearch.action.search.SearchScrollRequest; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +@EqualsAndHashCode public class ContinueScrollRequest implements OpenSearchRequest { final String initialScrollId; @@ -39,9 +39,7 @@ public ContinueScrollRequest(String scrollId, OpenSearchExprValueFactory exprVal @Override public OpenSearchResponse search(Function searchAction, Function scrollAction) { - SearchResponse openSearchResponse; - - openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId) + SearchResponse openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId) .scroll(DEFAULT_SCROLL_TIMEOUT)); // TODO if terminated_early - something went wrong, e.g. no scroll returned. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java similarity index 84% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java index f2e4fd08b4a..7ef87473816 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java @@ -3,14 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.request; import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.stream.Collectors; import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.index.query.QueryBuilder; @@ -22,8 +21,6 @@ import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.OpenSearchRequest; -import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; public class InitialPageRequestBuilder implements PagedRequestBuilder { @@ -40,15 +37,16 @@ public class InitialPageRequestBuilder implements PagedRequestBuilder { * @param settings other settings * @param exprValueFactory value factory */ + // TODO accept indexName as string (same way as `OpenSearchRequestBuilder` does)? public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, int pageSize, - Settings settings, + Settings settings, // TODO: settings are not used - refactor? OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; - this.sourceBuilder = new SearchSourceBuilder(); this.exprValueFactory = exprValueFactory; this.querySize = pageSize; - sourceBuilder.from(0) + this.sourceBuilder = new SearchSourceBuilder() + .from(0) .size(querySize) .timeout(DEFAULT_QUERY_TIMEOUT); } @@ -85,9 +83,8 @@ public void pushDownHighlight(String field, Map arguments) { * Push down project expression to OpenSearch. */ public void pushDownProjects(Set projects) { - final Set projectsSet = - projects.stream().map(ReferenceExpression::getAttr).collect(Collectors.toSet()); - sourceBuilder.fetchSource(projectsSet.toArray(new String[0]), new String[0]); + sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr) + .distinct().toArray(String[]::new), new String[0]); } public void pushTypeMapping(Map typeMapping) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 6f6fea841b6..0795ce7cdc7 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -6,6 +6,8 @@ package org.opensearch.sql.opensearch.request; +import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; + import com.google.common.annotations.VisibleForTesting; import java.util.function.Consumer; import java.util.function.Function; @@ -15,7 +17,6 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; -import org.opensearch.common.unit.TimeValue; import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; @@ -32,11 +33,6 @@ @ToString public class OpenSearchQueryRequest implements OpenSearchRequest { - /** - * Default query timeout in minutes. - */ - public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); - /** * {@link OpenSearchRequest.IndexName}. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index a432bf1ca8f..7239ea7c0b9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -42,6 +42,7 @@ @EqualsAndHashCode @Getter @ToString +// TODO make an interface which defines all pushDown functions? public class OpenSearchRequestBuilder { /** diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java similarity index 52% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java index ae89a238a0e..365c4a60615 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java @@ -3,10 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; - -import org.opensearch.sql.opensearch.request.OpenSearchRequest; -import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; +package org.opensearch.sql.opensearch.request; public interface PagedRequestBuilder { OpenSearchRequest build(); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java similarity index 56% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java index 89f8f71933f..c2fb1f3431a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java @@ -3,25 +3,22 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.request; +import lombok.Getter; import lombok.RequiredArgsConstructor; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.OpenSearchRequest; @RequiredArgsConstructor public class SubsequentPageRequestBuilder implements PagedRequestBuilder { - final OpenSearchRequest.IndexName indexName; - final String scrollId; - final OpenSearchExprValueFactory exprValueFactory; + + @Getter + private final OpenSearchRequest.IndexName indexName; + private final String scrollId; + private final OpenSearchExprValueFactory exprValueFactory; @Override public OpenSearchRequest build() { return new ContinueScrollRequest(scrollId, exprValueFactory); } - - @Override - public OpenSearchRequest.IndexName getIndexName() { - return indexName; - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 1366dc374bc..c5ec6c364b2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -18,10 +18,14 @@ import org.opensearch.sql.opensearch.planner.physical.ADOperator; import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator; import org.opensearch.sql.opensearch.planner.physical.MLOperator; +import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.request.system.OpenSearchDescribeIndexRequest; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScanBuilder; import org.opensearch.sql.planner.DefaultImplementor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalML; @@ -132,7 +136,7 @@ public TableScanBuilder createPagedScanBuilder(int pageSize) { var requestBuilder = new InitialPageRequestBuilder(indexName, pageSize, settings, new OpenSearchExprValueFactory(getFieldTypes())); var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); - return new OpenSearchPagedScanBuilder(indexScan); + return new OpenSearchPagedIndexScanBuilder(indexScan); } @VisibleForTesting diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java index 0b0e231760d..00a0a16bbd2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java @@ -14,6 +14,8 @@ import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan; import org.opensearch.sql.opensearch.storage.system.OpenSearchSystemIndex; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; @@ -39,6 +41,7 @@ public Table getTable(DataSourceSchemaName dataSourceSchemaName, String name) { @Override public TableScanOperator getTableScan(String indexName, String scrollId) { + // TODO call `getTable` here? var index = new OpenSearchIndex(client, settings, indexName); var requestBuilder = new SubsequentPageRequestBuilder( new OpenSearchRequest.IndexName(indexName), diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java similarity index 98% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 9a1ddcba08e..3ae2e62cfd5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -4,7 +4,7 @@ */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.storage.scan; import java.util.Collections; import java.util.Iterator; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 4e1b20db6e3..4571961e5fe 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -17,7 +17,6 @@ import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalSort; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index 0ce3a077077..41edbfc768a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -8,7 +8,6 @@ import com.google.common.annotations.VisibleForTesting; import lombok.EqualsAndHashCode; import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; @@ -92,11 +91,6 @@ public boolean pushDownProject(LogicalProject project) { return delegate.pushDownProject(project); } - @Override - public boolean pushDownOffset(int i) { - return delegate.pushDownOffset(i); - } - @Override public boolean pushDownHighlight(LogicalHighlight highlight) { return delegate.pushDownHighlight(highlight); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index cb0940c410e..5cfde4abbee 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -21,7 +21,6 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; import org.opensearch.sql.planner.logical.LogicalFilter; @@ -38,10 +37,6 @@ */ @VisibleForTesting class OpenSearchIndexScanQueryBuilder extends TableScanBuilder { - @Override - public void pushDownOffset(int i) { - indexScan.getRequestBuilder().getSourceBuilder().from(i); - } /** OpenSearch index scan to be optimized. */ @EqualsAndHashCode.Include diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java similarity index 88% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index fa4e4bf1065..5ab2fca3935 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -3,15 +3,17 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.storage.scan; import java.util.Collections; import java.util.Iterator; import lombok.EqualsAndHashCode; import lombok.ToString; +import org.apache.commons.lang3.NotImplementedException; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.PagedRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.storage.TableScanOperator; @@ -33,7 +35,7 @@ public OpenSearchPagedIndexScan(OpenSearchClient client, @Override public String explain() { - throw new RuntimeException("Implement OpenSearchPagedIndexScan.explain"); + throw new NotImplementedException("Implement OpenSearchPagedIndexScan.explain"); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java similarity index 58% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java index 09232cb4de7..779df4ebec9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java @@ -3,26 +3,25 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.storage.scan; import lombok.EqualsAndHashCode; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; /** - * Builder for a paged opensearch request. - * Override pushDown* methods from TableScaneBuilder as more features + * Builder for a paged OpenSearch request. + * Override pushDown* methods from TableScanBuilder as more features * support pagination. */ -public class OpenSearchPagedScanBuilder extends TableScanBuilder { +public class OpenSearchPagedIndexScanBuilder extends TableScanBuilder { @EqualsAndHashCode.Include OpenSearchPagedIndexScan indexScan; - public OpenSearchPagedScanBuilder(OpenSearchPagedIndexScan indexScan) { + public OpenSearchPagedIndexScanBuilder(OpenSearchPagedIndexScan indexScan) { this.indexScan = indexScan; } - @Override public TableScanOperator build() { return indexScan; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index ccfd2a57d05..ab4171ad22d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -15,6 +15,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -33,6 +34,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InOrder; @@ -270,7 +272,8 @@ void search() { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(nodeClient.searchScroll(any()).actionGet()).thenReturn(scrollResponse); - when(scrollResponse.getScrollId()).thenReturn("scroll456"); + // TODO commented out because scroll clean-up is disabled + //when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -284,6 +287,7 @@ void search() { assertFalse(hits.hasNext()); // Verify response for second scroll request + request.setScrollId("scroll123"); OpenSearchResponse response2 = client.search(request); assertTrue(response2.isEmpty()); } @@ -302,18 +306,20 @@ void schedule() { void cleanup() { ClearScrollRequestBuilder requestBuilder = mock(ClearScrollRequestBuilder.class); when(nodeClient.prepareClearScroll()).thenReturn(requestBuilder); - when(requestBuilder.addScrollId(any())).thenReturn(requestBuilder); - when(requestBuilder.get()).thenReturn(null); + lenient().when(requestBuilder.addScrollId(any())).thenReturn(requestBuilder); + lenient().when(requestBuilder.get()).thenReturn(null); OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); request.setScrollId("scroll123"); client.cleanup(request); assertFalse(request.isScrollStarted()); + /* TODO: Scroll cleaning is temporary disabled InOrder inOrder = Mockito.inOrder(nodeClient, requestBuilder); inOrder.verify(nodeClient).prepareClearScroll(); inOrder.verify(requestBuilder).addScrollId("scroll123"); inOrder.verify(requestBuilder).get(); + */ } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index 083446adcc2..67c6f51c982 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -29,6 +29,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -251,7 +252,8 @@ void search() throws IOException { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(restClient.scroll(any(), any())).thenReturn(scrollResponse); - when(scrollResponse.getScrollId()).thenReturn("scroll456"); + // TODO commented out because scroll clean-up is disabled + //when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -265,6 +267,7 @@ void search() throws IOException { assertFalse(hits.hasNext()); // Verify response for second scroll request + request.setScrollId("scroll123"); OpenSearchResponse response2 = client.search(request); assertTrue(response2.isEmpty()); } @@ -315,7 +318,8 @@ void cleanup() throws IOException { OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); request.setScrollId("scroll123"); client.cleanup(request); - verify(restClient).clearScroll(any(), any()); + // TODO: Scroll cleaning is temporary disabled + //verify(restClient).clearScroll(any(), any()); assertFalse(request.isScrollStarted()); } @@ -326,6 +330,7 @@ void cleanupWithoutScrollId() throws IOException { verify(restClient, never()).clearScroll(any(), any()); } + @Disabled("TODO: Scroll cleaning is temporary disabled") @Test void cleanupWithIOException() throws IOException { when(restClient.clearScroll(any(), any())).thenThrow(new IOException()); 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 2441b2d3b21..333fe1cfec1 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 @@ -44,8 +44,7 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; -import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index c64d9e4ad96..58da3f3f9aa 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -59,7 +59,8 @@ import org.opensearch.sql.opensearch.planner.physical.MLOperator; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; @@ -319,6 +320,13 @@ public void testVisitML() { executionProtector.visitML(mlOperator, null)); } + @Test + public void visitPaginate() { + var paginate = new PaginateOperator(values(List.of()), 42); + assertEquals(executionProtector.protect(paginate), + executionProtector.visitPaginate(paginate, null)); + } + PhysicalPlan resourceMonitor(PhysicalPlan input) { return new ResourceMonitorPlan(input, resourceMonitor); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java new file mode 100644 index 00000000000..f25553a55f1 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java @@ -0,0 +1,117 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.request; + +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.function.Consumer; +import java.util.function.Function; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; + +@ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class ContinueScrollRequestTest { + + @Mock + private Function searchAction; + + @Mock + private Function scrollAction; + + @Mock + private Consumer cleanAction; + + @Mock + private SearchResponse searchResponse; + + @Mock + private SearchHits searchHits; + + @Mock + private SearchHit searchHit; + + @Mock + private OpenSearchExprValueFactory factory; + + private final String scroll = "scroll"; + private final String nextScroll = "nextScroll"; + + private final ContinueScrollRequest request = new ContinueScrollRequest(scroll, factory); + + @Test + public void search_with_non_empty_response() { + when(scrollAction.apply(any())).thenReturn(searchResponse); + when(searchResponse.getHits()).thenReturn(searchHits); + when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); + when(searchResponse.getScrollId()).thenReturn(nextScroll); + + OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); + assertAll( + () -> assertFalse(searchResponse.isEmpty()), + () -> assertEquals(nextScroll, request.toCursor()), + () -> verify(scrollAction, times(1)).apply(any()), + () -> verify(searchAction, never()).apply(any()) + ); + } + + @Test + // Empty response means scroll search is done and no cursor/scroll should be set + public void search_with_empty_response() { + when(scrollAction.apply(any())).thenReturn(searchResponse); + when(searchResponse.getHits()).thenReturn(searchHits); + when(searchHits.getHits()).thenReturn(null); + lenient().when(searchResponse.getScrollId()).thenReturn(nextScroll); + + OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); + assertAll( + () -> assertTrue(searchResponse.isEmpty()), + () -> assertNull(request.toCursor()), + () -> verify(scrollAction, times(1)).apply(any()), + () -> verify(searchAction, never()).apply(any()) + ); + } + + @Test + public void clean() { + request.clean(cleanAction); + verify(cleanAction, times(1)).accept(any()); + } + + @Test + // Added for coverage only + public void getters() { + factory = mock(); + assertAll( + () -> assertThrows(Throwable.class, request::getSourceBuilder), + () -> assertSame(factory, new ContinueScrollRequest("", factory).getExprValueFactory()) + ); + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java new file mode 100644 index 00000000000..023a1e397a1 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java @@ -0,0 +1,107 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.request; + +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; + +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.expression.DSL; +import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +@ExtendWith(MockitoExtension.class) +public class InitialPageRequestBuilderTest { + + @Mock + private OpenSearchExprValueFactory exprValueFactory; + + @Mock + private Settings settings; + + private final int pageSize = 42; + + private final OpenSearchRequest.IndexName indexName = new OpenSearchRequest.IndexName("test"); + + private InitialPageRequestBuilder requestBuilder; + + @BeforeEach + void setup() { + requestBuilder = new InitialPageRequestBuilder( + indexName, pageSize, settings, exprValueFactory); + } + + @Test + public void build() { + assertEquals( + new OpenSearchScrollRequest(indexName, + new SearchSourceBuilder() + .from(0) + .size(pageSize) + .timeout(DEFAULT_QUERY_TIMEOUT), + exprValueFactory), + requestBuilder.build() + ); + } + + @Test + public void pushDown_not_supported() { + assertAll( + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDown(mock())), + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownAggregation(mock())), + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownSort(mock())), + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownLimit(1, 2)), + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownHighlight("", Map.of())) + ); + } + + @Test + public void pushTypeMapping() { + Map typeMapping = Map.of("intA", INTEGER); + requestBuilder.pushTypeMapping(typeMapping); + + verify(exprValueFactory).setTypeMapping(typeMapping); + } + + @Test + public void pushDownProject() { + Set references = Set.of(DSL.ref("intA", INTEGER)); + requestBuilder.pushDownProjects(references); + + assertEquals( + new OpenSearchScrollRequest(indexName, + new SearchSourceBuilder() + .from(0) + .size(pageSize) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource(new String[]{"intA"}, new String[0]), + exprValueFactory), + requestBuilder.build() + ); + } + + @Test + public void getIndexName() { + assertEquals(indexName, requestBuilder.getIndexName()); + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index 1ba26e33dc0..c6a9a06a70d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -14,6 +14,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; import java.util.function.Consumer; import java.util.function.Function; @@ -85,7 +86,7 @@ void searchRequest() { new SearchRequest() .indices("test") .source(new SearchSourceBuilder() - .timeout(OpenSearchQueryRequest.DEFAULT_QUERY_TIMEOUT) + .timeout(DEFAULT_QUERY_TIMEOUT) .from(0) .size(200) .query(QueryBuilders.termQuery("name", "John"))), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 0fc9c928106..8bb5c1eebe5 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -47,6 +47,9 @@ void isScrollStarted() { request.setScrollId("scroll123"); assertTrue(request.isScrollStarted()); + + request.reset(); + assertFalse(request.isScrollStarted()); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilderTest.java new file mode 100644 index 00000000000..b8f04b8b2c3 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilderTest.java @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.request; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +@ExtendWith(MockitoExtension.class) +public class SubsequentPageRequestBuilderTest { + + @Mock + private OpenSearchExprValueFactory exprValueFactory; + + private final OpenSearchRequest.IndexName indexName = new OpenSearchRequest.IndexName("test"); + private final String scrollId = "scroll"; + + private SubsequentPageRequestBuilder requestBuilder; + + @BeforeEach + void setup() { + requestBuilder = new SubsequentPageRequestBuilder(indexName, scrollId, exprValueFactory); + } + + @Test + public void build() { + assertEquals( + new ContinueScrollRequest(scrollId, exprValueFactory), + requestBuilder.build() + ); + } + + @Test + public void getIndexName() { + assertEquals(indexName, requestBuilder.getIndexName()); + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 4adacb7dd28..6aaee95a00e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -53,7 +53,12 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.mapping.IndexMapping; +import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.request.PagedRequestBuilder; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; @@ -165,6 +170,17 @@ void implementRelationOperatorOnly() { assertEquals(new OpenSearchIndexScan(client, builder), index.implement(plan)); } + @Test + void implementPagedRelationOperatorOnly() { + when(client.getIndexMaxResultWindows("test")).thenReturn(Map.of("test", 10000)); + + LogicalPlan plan = index.createPagedScanBuilder(42); + Integer maxResultWindow = index.getMaxResultWindow(); + PagedRequestBuilder builder = new InitialPageRequestBuilder( + new OpenSearchRequest.IndexName(indexName), maxResultWindow, settings, exprValueFactory); + assertEquals(new OpenSearchPagedIndexScan(client, builder), index.implement(plan)); + } + @Test void implementRelationOperatorWithOptimization() { when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngineTest.java index ab87f4531cf..6a8727e0fbc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngineTest.java @@ -6,11 +6,18 @@ package org.opensearch.sql.opensearch.storage; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME; import static org.opensearch.sql.utils.SystemIndexUtils.TABLE_INFO; +import java.util.Map; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -18,6 +25,8 @@ import org.opensearch.sql.DataSourceSchemaName; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan; import org.opensearch.sql.opensearch.storage.system.OpenSearchSystemIndex; import org.opensearch.sql.storage.Table; @@ -35,7 +44,10 @@ public void getTable() { OpenSearchStorageEngine engine = new OpenSearchStorageEngine(client, settings); Table table = engine.getTable(new DataSourceSchemaName(DEFAULT_DATASOURCE_NAME, "default"), "test"); - assertNotNull(table); + assertAll( + () -> assertNotNull(table), + () -> assertTrue(table instanceof OpenSearchIndex) + ); } @Test @@ -43,7 +55,26 @@ public void getSystemTable() { OpenSearchStorageEngine engine = new OpenSearchStorageEngine(client, settings); Table table = engine.getTable(new DataSourceSchemaName(DEFAULT_DATASOURCE_NAME, "default"), TABLE_INFO); - assertNotNull(table); - assertTrue(table instanceof OpenSearchSystemIndex); + assertAll( + () -> assertNotNull(table), + () -> assertTrue(table instanceof OpenSearchSystemIndex) + ); + } + + @Test + public void getTableScan() { + when(client.getIndexMappings(anyString())).thenReturn(Map.of()); + OpenSearchResponse response = mock(); + when(response.isEmpty()).thenReturn(true); + when(client.search(any())).thenReturn(response); + OpenSearchStorageEngine engine = new OpenSearchStorageEngine(client, settings); + var scan = engine.getTableScan("test", "test"); + assertAll( + () -> assertTrue(scan instanceof OpenSearchPagedIndexScan), + () -> { + scan.open(); + assertFalse(scan.hasNext()); + } + ); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index 363727cbd38..b5de6e30c5b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -64,7 +64,6 @@ import org.opensearch.sql.opensearch.response.agg.CompositeAggregationParser; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; -import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java similarity index 96% rename from opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java rename to opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index bd0d0089fb3..90ad624135b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -4,7 +4,7 @@ */ -package org.opensearch.sql.opensearch.storage; +package org.opensearch.sql.opensearch.storage.scan; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -65,7 +65,7 @@ void setup() { @Test void queryEmptyResult() { - mockResponse(); + mockResponse(client); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("test", 3, settings, exprValueFactory))) { @@ -77,7 +77,7 @@ void queryEmptyResult() { @Test void queryAllResultsWithQuery() { - mockResponse(new ExprValue[]{ + mockResponse(client, new ExprValue[]{ employee(1, "John", "IT"), employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); @@ -105,7 +105,7 @@ void queryAllResultsWithQuery() { @Test void queryAllResultsWithScroll() { - mockResponse( + mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); @@ -130,7 +130,7 @@ void queryAllResultsWithScroll() { @Test void querySomeResultsWithQuery() { - mockResponse(new ExprValue[]{ + mockResponse(client, new ExprValue[]{ employee(1, "John", "IT"), employee(2, "Smith", "HR"), employee(3, "Allen", "IT"), @@ -158,7 +158,7 @@ void querySomeResultsWithQuery() { @Test void querySomeResultsWithScroll() { - mockResponse( + mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); @@ -228,7 +228,7 @@ void pushDownHighlightWithArguments() { @Test void pushDownHighlightWithRepeatingFields() { - mockResponse( + mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); @@ -300,7 +300,7 @@ PushDownAssertion shouldQuery(QueryBuilder expected) { } } - private void mockResponse(ExprValue[]... searchHitBatches) { + public static void mockResponse(OpenSearchClient client, ExprValue[]... searchHitBatches) { when(client.search(any())) .thenAnswer( new Answer() { @@ -324,14 +324,14 @@ public OpenSearchResponse answer(InvocationOnMock invocation) { }); } - protected ExprValue employee(int docId, String name, String department) { + public static ExprValue employee(int docId, String name, String department) { SearchHit hit = new SearchHit(docId); hit.sourceRef( new BytesArray("{\"name\":\"" + name + "\",\"department\":\"" + department + "\"}")); return tupleValue(hit); } - private ExprValue tupleValue(SearchHit hit) { + private static ExprValue tupleValue(SearchHit hit) { return ExprValueUtils.tupleValue(hit.getSourceAsMap()); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java new file mode 100644 index 00000000000..9006a0573d9 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -0,0 +1,117 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.withSettings; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.employee; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; + +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; + +@ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class OpenSearchPagedIndexScanTest { + @Mock + private OpenSearchClient client; + + @Mock + private Settings settings; + + private OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( + ImmutableMap.of("name", STRING, "department", STRING)); + + @Test + void query_empty_result() { + mockResponse(client); + InitialPageRequestBuilder builder = new InitialPageRequestBuilder( + new OpenSearchRequest.IndexName("test"), 3, settings, exprValueFactory); + try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { + indexScan.open(); + assertFalse(indexScan.hasNext()); + } + verify(client).cleanup(any()); + } + + @Test + void query_all_results_initial_scroll_request() { + mockResponse(client, new ExprValue[]{ + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); + + InitialPageRequestBuilder builder = new InitialPageRequestBuilder( + new OpenSearchRequest.IndexName("test"), 3, settings, exprValueFactory); + try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { + indexScan.open(); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(1, "John", "IT"), indexScan.next()); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + + assertFalse(indexScan.hasNext()); + } + verify(client).cleanup(any()); + } + + @Test + void query_all_results_continuation_scroll_request() { + mockResponse(client, new ExprValue[]{ + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); + + SubsequentPageRequestBuilder builder = new SubsequentPageRequestBuilder( + new OpenSearchRequest.IndexName("test"), "scroll", exprValueFactory); + try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { + indexScan.open(); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(1, "John", "IT"), indexScan.next()); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + + assertTrue(indexScan.hasNext()); + assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + + assertFalse(indexScan.hasNext()); + } + verify(client).cleanup(any()); + } + + @Test + void explain_not_implemented() { + assertThrows(Throwable.class, () -> mock(OpenSearchPagedIndexScan.class, + withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); + } +} From 27e1793c5ae7ad837d1766c1a9b530d1a7cc0cd9 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 9 Mar 2023 20:04:17 -0800 Subject: [PATCH 21/46] Fix the merges. Signed-off-by: Yury-Fridlyand --- .../sql/planner/optimizer/LogicalPlanOptimizerTest.java | 2 -- .../sql/opensearch/executor/OpenSearchExecutionEngineTest.java | 1 + sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java | 3 --- 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index 9d3620d1116..1ee9b9aa3b1 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -348,8 +348,6 @@ void push_page_size_noop_if_no_sub_plans() { @Test void table_scan_builder_support_offset_push_down_can_apply_its_rule() { - // next line is noop, added for coverage only - lenient().when(tableScanBuilder.pushDownOffset(anyInt())).thenReturn(true); when(table.createPagedScanBuilder(anyInt())).thenReturn(tableScanBuilder); var optimized = LogicalPlanOptimizer.paginationCreate() 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 333fe1cfec1..913a774e1b7 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 @@ -45,6 +45,7 @@ import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; +import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; 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 cb1159bc17f..9eaf4f56ba1 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -57,9 +57,6 @@ class SQLServiceTest { @Mock private PaginatedPlanCache paginatedPlanCache; - @Mock - private PaginatedQueryService paginatedQueryService; - @BeforeEach public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); From 304616d496d4d3e94b037fbe24228e7552f5d004 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 9 Mar 2023 20:05:52 -0800 Subject: [PATCH 22/46] Fix explain. Signed-off-by: Yury-Fridlyand --- .../execution/ContinuePaginatedPlan.java | 5 +++-- .../executor/execution/QueryPlanFactory.java | 8 +++++--- .../execution/ContinuePaginatedPlanTest.java | 7 +++++-- .../execution/QueryPlanFactoryTest.java | 9 +++++++-- .../org/opensearch/sql/sql/SQLService.java | 19 ++----------------- 5 files changed, 22 insertions(+), 26 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java index d61747d0eb7..aa86b61767e 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java @@ -5,7 +5,6 @@ package org.opensearch.sql.executor.execution; -import org.apache.commons.lang3.NotImplementedException; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; @@ -51,6 +50,8 @@ public void execute() { @Override // TODO why can't use listener given in the constructor? public void explain(ResponseListener listener) { - throw new UnsupportedOperationException("Explain of query continuation is not supported"); + listener.onFailure(new UnsupportedOperationException( + "Explain of a paged query continuation is not supported. " + + "Use `explain` for the initial query request.")); } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 881046ec0b9..c3a84ea286e 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -76,11 +76,13 @@ public AbstractPlan create( /** * Creates a ContinuePaginatedPlan from a cursor. */ - public AbstractPlan create(String cursor, ResponseListener - queryResponseListener) { + public AbstractPlan create(String cursor, boolean isExplain, + ResponseListener queryResponseListener, + ResponseListener explainListener) { QueryId queryId = QueryId.queryId(); - return new ContinuePaginatedPlan(queryId, cursor, paginatedQueryService, paginatedPlanCache, + var cpp = new ContinuePaginatedPlan(queryId, cursor, paginatedQueryService, paginatedPlanCache, queryResponseListener); + return isExplain ? new ExplainPlan(queryId, cpp, explainListener) : cpp; } @Override diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java index c8c95a0ae68..6822eca727f 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java @@ -11,8 +11,10 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.sql.executor.PaginatedPlanCacheTest.buildCursor; @@ -114,7 +116,8 @@ public void onFailure(Exception e) { @Test public void explain_is_not_supported() { - assertThrows(UnsupportedOperationException.class, - () -> ContinuePaginatedPlan.None.explain(null)); + var listener = mock(ResponseListener.class); + ContinuePaginatedPlan.None.explain(listener); + verify(listener).onFailure(any(UnsupportedOperationException.class)); } } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index 72225f0884f..08718cb4d61 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -8,6 +8,7 @@ package org.opensearch.sql.executor.execution; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -78,8 +79,12 @@ public void createFromExplainShouldSuccess() { @Test public void createFromCursorShouldSuccess() { - AbstractPlan queryExecution = factory.create("", queryListener); - assertTrue(queryExecution instanceof ContinuePaginatedPlan); + AbstractPlan queryExecution = factory.create("", false, queryListener, explainListener); + AbstractPlan explainExecution = factory.create("", true, queryListener, explainListener); + assertAll( + () -> assertTrue(queryExecution instanceof ContinuePaginatedPlan), + () -> assertTrue(explainExecution instanceof ExplainPlan) + ); } @Test 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 2acc44de63e..7013842066c 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -68,23 +68,8 @@ private AbstractPlan plan( Optional> explainListener) { if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. - if (queryListener.isEmpty() && explainListener.isPresent()) { // explain request - // TODO explain should be processed inside the plan - explainListener.get().onFailure(new UnsupportedOperationException( - "`explain` request for cursor requests is not supported. " - + "Use `explain` for the initial query request.")); - return new AbstractPlan(QueryId.queryId()) { - @Override - public void execute() { - } - - @Override - public void explain(ResponseListener listener) { - } - }; - } - // non-explain request - return queryExecutionFactory.create(request.getCursor().get(), queryListener.get()); + return queryExecutionFactory.create(request.getCursor().get(), request.isExplainRequest(), + queryListener.orElse(null), explainListener.orElse(null)); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) ParseTree cst = parser.parse(request.getQuery()); From 1b5ab7e0a37e17402243c7c7dc12bdbf7fc490e1 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Mar 2023 13:01:14 -0800 Subject: [PATCH 23/46] Fix scroll cleaning. Signed-off-by: Yury-Fridlyand --- .../client/OpenSearchNodeClient.java | 11 +++- .../client/OpenSearchRestClient.java | 7 ++- .../request/ContinueScrollRequest.java | 12 +++-- .../request/OpenSearchScrollRequest.java | 5 +- .../scan/OpenSearchPagedIndexScan.java | 8 ++- .../client/OpenSearchNodeClientTest.java | 48 ++++++++++------- .../client/OpenSearchRestClientTest.java | 47 ++++++++-------- .../scan/OpenSearchPagedIndexScanTest.java | 53 ++++++++++++++----- 8 files changed, 119 insertions(+), 72 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java index 30584581e31..e4f25dabbdc 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java @@ -43,7 +43,7 @@ public class OpenSearchNodeClient implements OpenSearchClient { private final NodeClient client; /** - * Constructor of ElasticsearchNodeClient. + * Constructor of OpenSearchNodeClient. */ public OpenSearchNodeClient(NodeClient client) { this.client = client; @@ -172,7 +172,14 @@ public Map meta() { @Override public void cleanup(OpenSearchRequest request) { - request.clean(scrollId -> {}/* client.prepareClearScroll().addScrollId(scrollId).get() */); + request.clean(scrollId -> { + try { + client.prepareClearScroll().addScrollId(scrollId).get(); + } catch (Exception e) { + throw new IllegalStateException( + "Failed to clean up resources for search request " + request, e); + } + }); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java index 096899516f9..41efc4e5efb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java @@ -174,17 +174,16 @@ public Map meta() { @Override public void cleanup(OpenSearchRequest request) { - request.clean(scrollId -> {/* + request.clean(scrollId -> { try { ClearScrollRequest clearRequest = new ClearScrollRequest(); clearRequest.addScrollId(scrollId); client.clearScroll(clearRequest, RequestOptions.DEFAULT); - } catch (IOException e) { + } catch (Exception e) { throw new IllegalStateException( "Failed to clean up resources for search request " + request, e); - }*/ + } }); - } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java index 1ec5960b21d..cf51680b7e2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java @@ -31,6 +31,9 @@ public class ContinueScrollRequest implements OpenSearchRequest { @Getter private final OpenSearchExprValueFactory exprValueFactory; + @EqualsAndHashCode.Exclude + private boolean scrollFinished = false; + public ContinueScrollRequest(String scrollId, OpenSearchExprValueFactory exprValueFactory) { this.initialScrollId = scrollId; this.exprValueFactory = exprValueFactory; @@ -44,9 +47,9 @@ public OpenSearchResponse search(Function searchA // TODO if terminated_early - something went wrong, e.g. no scroll returned. var response = new OpenSearchResponse(openSearchResponse, exprValueFactory); - if (!response.isEmpty()) { - responseScrollId = openSearchResponse.getScrollId(); - } // else - last empty page, we should ignore the scroll even if it is returned + // on the last empty page, we should close the scroll + scrollFinished = response.isEmpty(); + responseScrollId = openSearchResponse.getScrollId(); return response; } @@ -63,6 +66,7 @@ public SearchSourceBuilder getSourceBuilder() { @Override public String toCursor() { - return responseScrollId; + // on the last page, we shouldn't return the scroll to user, it is kept for closing (clean) + return scrollFinished ? null : responseScrollId; } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 8eef94c8375..555d520d7f6 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -88,9 +88,7 @@ public OpenSearchResponse search(Function searchA } var response = new OpenSearchResponse(openSearchResponse, exprValueFactory); - if (!response.isEmpty()) { - setScrollId(openSearchResponse.getScrollId()); - } // else - last empty page, we should ignore the scroll even if it is returned + setScrollId(openSearchResponse.getScrollId()); return response; } @@ -99,6 +97,7 @@ public void clean(Consumer cleanAction) { try { if (isScrollStarted()) { cleanAction.accept(getScrollId()); + setScrollId(null); } } finally { reset(); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index 5ab2fca3935..1dc455cfd2d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -26,6 +26,7 @@ public class OpenSearchPagedIndexScan extends TableScanOperator { @ToString.Include private OpenSearchRequest request; private Iterator iterator; + private boolean needClean = false; public OpenSearchPagedIndexScan(OpenSearchClient client, PagedRequestBuilder requestBuilder) { @@ -56,6 +57,7 @@ public void open() { if (!response.isEmpty()) { iterator = response.iterator(); } else { + needClean = true; iterator = Collections.emptyIterator(); } } @@ -63,8 +65,10 @@ public void open() { @Override public void close() { super.close(); - - client.cleanup(request); + if (needClean) { + // clean on the last page only, to prevent closing the scroll/cursor in the middle of paging. + client.cleanup(request); + } } @Override diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index ab4171ad22d..206e1748f3a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -34,7 +34,8 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InOrder; @@ -73,6 +74,7 @@ import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchNodeClientTest { private static final String TEST_MAPPING_FILE = "mappings/accounts.json"; @@ -104,7 +106,7 @@ void setUp() { } @Test - void isIndexExist() { + void is_index_exist() { when(nodeClient.admin().indices() .exists(any(IndicesExistsRequest.class)).actionGet()) .thenReturn(new IndicesExistsResponse(true)); @@ -113,7 +115,7 @@ void isIndexExist() { } @Test - void isIndexNotExist() { + void is_index_not_exist() { String indexName = "test"; when(nodeClient.admin().indices() .exists(any(IndicesExistsRequest.class)).actionGet()) @@ -123,14 +125,14 @@ void isIndexNotExist() { } @Test - void isIndexExistWithException() { + void is_index_exist_with_exception() { when(nodeClient.admin().indices().exists(any())).thenThrow(RuntimeException.class); assertThrows(IllegalStateException.class, () -> client.exists("test")); } @Test - void createIndex() { + void create_index() { String indexName = "test"; Map mappings = ImmutableMap.of( "properties", @@ -143,7 +145,7 @@ void createIndex() { } @Test - void createIndexWithException() { + void create_index_with_exception() { when(nodeClient.admin().indices().create(any())).thenThrow(RuntimeException.class); assertThrows(IllegalStateException.class, @@ -151,7 +153,7 @@ void createIndexWithException() { } @Test - void getIndexMappings() throws IOException { + void get_index_mappings() throws IOException { URL url = Resources.getResource(TEST_MAPPING_FILE); String mappings = Resources.toString(url, Charsets.UTF_8); String indexName = "test"; @@ -183,7 +185,7 @@ void getIndexMappings() throws IOException { } @Test - void getIndexMappingsWithEmptyMapping() { + void get_index_mappings_with_empty_mapping() { String indexName = "test"; mockNodeClientIndicesMappings(indexName, ""); Map indexMappings = client.getIndexMappings(indexName); @@ -194,7 +196,7 @@ void getIndexMappingsWithEmptyMapping() { } @Test - void getIndexMappingsWithIOException() { + void get_index_mappings_with_IOException() { String indexName = "test"; when(nodeClient.admin().indices()).thenThrow(RuntimeException.class); @@ -202,7 +204,7 @@ void getIndexMappingsWithIOException() { } @Test - void getIndexMappingsWithNonExistIndex() { + void get_index_mappings_with_non_exist_index() { when(nodeClient.admin().indices() .prepareGetMappings(any()) .setLocal(anyBoolean()) @@ -213,7 +215,7 @@ void getIndexMappingsWithNonExistIndex() { } @Test - void getIndexMaxResultWindows() throws IOException { + void get_index_max_result_windows() throws IOException { URL url = Resources.getResource(TEST_MAPPING_SETTINGS_FILE); String indexMetadata = Resources.toString(url, Charsets.UTF_8); String indexName = "accounts"; @@ -227,7 +229,7 @@ void getIndexMaxResultWindows() throws IOException { } @Test - void getIndexMaxResultWindowsWithDefaultSettings() throws IOException { + void get_index_max_result_windows_with_default_settings() throws IOException { URL url = Resources.getResource(TEST_MAPPING_FILE); String indexMetadata = Resources.toString(url, Charsets.UTF_8); String indexName = "accounts"; @@ -241,7 +243,7 @@ void getIndexMaxResultWindowsWithDefaultSettings() throws IOException { } @Test - void getIndexMaxResultWindowsWithIOException() { + void get_index_max_result_windows_with_IOException() { String indexName = "test"; when(nodeClient.admin().indices()).thenThrow(RuntimeException.class); @@ -250,7 +252,7 @@ void getIndexMaxResultWindowsWithIOException() { /** Jacoco enforce this constant lambda be tested. */ @Test - void testAllFieldsPredicate() { + void test_all_fields_predicate() { assertTrue(OpenSearchNodeClient.ALL_FIELDS.apply("any_index").test("any_field")); } @@ -272,8 +274,7 @@ void search() { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(nodeClient.searchScroll(any()).actionGet()).thenReturn(scrollResponse); - // TODO commented out because scroll clean-up is disabled - //when(scrollResponse.getScrollId()).thenReturn("scroll456"); + when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -314,23 +315,30 @@ void cleanup() { client.cleanup(request); assertFalse(request.isScrollStarted()); - /* TODO: Scroll cleaning is temporary disabled InOrder inOrder = Mockito.inOrder(nodeClient, requestBuilder); inOrder.verify(nodeClient).prepareClearScroll(); inOrder.verify(requestBuilder).addScrollId("scroll123"); inOrder.verify(requestBuilder).get(); - */ } @Test - void cleanupWithoutScrollId() { + void cleanup_without_scrollId() { OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); client.cleanup(request); verify(nodeClient, never()).prepareClearScroll(); } @Test - void getIndices() { + void cleanup_rethrows_exception() { + when(nodeClient.prepareClearScroll()).thenThrow(new RuntimeException()); + + OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); + request.setScrollId("scroll123"); + assertThrows(IllegalStateException.class, () -> client.cleanup(request)); + } + + @Test + void get_indices() { AliasMetadata aliasMetadata = mock(AliasMetadata.class); ImmutableOpenMap.Builder> builder = ImmutableOpenMap.builder(); builder.fPut("index",Arrays.asList(aliasMetadata)); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index 67c6f51c982..9fcf7a00796 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -29,7 +29,8 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -64,6 +65,7 @@ import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchRestClientTest { private static final String TEST_MAPPING_FILE = "mappings/accounts.json"; @@ -91,7 +93,7 @@ void setUp() { } @Test - void isIndexExist() throws IOException { + void is_index_exist() throws IOException { when(restClient.indices() .exists(any(), any())) // use any() because missing equals() in GetIndexRequest .thenReturn(true); @@ -100,7 +102,7 @@ void isIndexExist() throws IOException { } @Test - void isIndexNotExist() throws IOException { + void is_index_not_exist() throws IOException { when(restClient.indices() .exists(any(), any())) // use any() because missing equals() in GetIndexRequest .thenReturn(false); @@ -109,14 +111,14 @@ void isIndexNotExist() throws IOException { } @Test - void isIndexExistWithException() throws IOException { + void is_index_exist_with_exception() throws IOException { when(restClient.indices().exists(any(), any())).thenThrow(IOException.class); assertThrows(IllegalStateException.class, () -> client.exists("test")); } @Test - void createIndex() throws IOException { + void create_index() throws IOException { String indexName = "test"; Map mappings = ImmutableMap.of( "properties", @@ -129,7 +131,7 @@ void createIndex() throws IOException { } @Test - void createIndexWithIOException() throws IOException { + void create_index_with_IOException() throws IOException { when(restClient.indices().create(any(), any())).thenThrow(IOException.class); assertThrows(IllegalStateException.class, @@ -137,7 +139,7 @@ void createIndexWithIOException() throws IOException { } @Test - void getIndexMappings() throws IOException { + void get_index_mappings() throws IOException { URL url = Resources.getResource(TEST_MAPPING_FILE); String mappings = Resources.toString(url, Charsets.UTF_8); String indexName = "test"; @@ -173,14 +175,14 @@ void getIndexMappings() throws IOException { } @Test - void getIndexMappingsWithIOException() throws IOException { + void get_index_mappings_with_IOException() throws IOException { when(restClient.indices().getMapping(any(GetMappingsRequest.class), any())) .thenThrow(new IOException()); assertThrows(IllegalStateException.class, () -> client.getIndexMappings("test")); } @Test - void getIndexMaxResultWindowsSettings() throws IOException { + void get_index_max_result_windows_settings() throws IOException { String indexName = "test"; Integer maxResultWindow = 1000; @@ -204,7 +206,7 @@ void getIndexMaxResultWindowsSettings() throws IOException { } @Test - void getIndexMaxResultWindowsDefaultSettings() throws IOException { + void get_index_max_result_windows_default_settings() throws IOException { String indexName = "test"; Integer maxResultWindow = 10000; @@ -228,7 +230,7 @@ void getIndexMaxResultWindowsDefaultSettings() throws IOException { } @Test - void getIndexMaxResultWindowsWithIOException() throws IOException { + void get_index_max_result_windows_with_IOException() throws IOException { when(restClient.indices().getSettings(any(GetSettingsRequest.class), any())) .thenThrow(new IOException()); assertThrows(IllegalStateException.class, () -> client.getIndexMaxResultWindows("test")); @@ -252,8 +254,7 @@ void search() throws IOException { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(restClient.scroll(any(), any())).thenReturn(scrollResponse); - // TODO commented out because scroll clean-up is disabled - //when(scrollResponse.getScrollId()).thenReturn("scroll456"); + when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -273,7 +274,7 @@ void search() throws IOException { } @Test - void searchWithIOException() throws IOException { + void search_with_IOException() throws IOException { when(restClient.search(any(), any())).thenThrow(new IOException()); assertThrows( IllegalStateException.class, @@ -281,7 +282,7 @@ void searchWithIOException() throws IOException { } @Test - void scrollWithIOException() throws IOException { + void scroll_with_IOException() throws IOException { // Mock first scroll request SearchResponse searchResponse = mock(SearchResponse.class); when(restClient.search(any(), any())).thenReturn(searchResponse); @@ -318,21 +319,19 @@ void cleanup() throws IOException { OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); request.setScrollId("scroll123"); client.cleanup(request); - // TODO: Scroll cleaning is temporary disabled - //verify(restClient).clearScroll(any(), any()); + verify(restClient).clearScroll(any(), any()); assertFalse(request.isScrollStarted()); } @Test - void cleanupWithoutScrollId() throws IOException { + void cleanup_without_scrollId() throws IOException { OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); client.cleanup(request); verify(restClient, never()).clearScroll(any(), any()); } - @Disabled("TODO: Scroll cleaning is temporary disabled") @Test - void cleanupWithIOException() throws IOException { + void cleanup_with_IOException() throws IOException { when(restClient.clearScroll(any(), any())).thenThrow(new IOException()); OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); @@ -341,7 +340,7 @@ void cleanupWithIOException() throws IOException { } @Test - void getIndices() throws IOException { + void get_indices() throws IOException { when(restClient.indices().get(any(GetIndexRequest.class), any(RequestOptions.class))) .thenReturn(getIndexResponse); when(getIndexResponse.getIndices()).thenReturn(new String[] {"index"}); @@ -351,7 +350,7 @@ void getIndices() throws IOException { } @Test - void getIndicesWithIOException() throws IOException { + void get_indices_with_IOException() throws IOException { when(restClient.indices().get(any(GetIndexRequest.class), any(RequestOptions.class))) .thenThrow(new IOException()); assertThrows(IllegalStateException.class, () -> client.indices()); @@ -370,7 +369,7 @@ void meta() throws IOException { } @Test - void metaWithIOException() throws IOException { + void meta_with_IOException() throws IOException { when(restClient.cluster().getSettings(any(), any(RequestOptions.class))) .thenThrow(new IOException()); @@ -378,7 +377,7 @@ void metaWithIOException() throws IOException { } @Test - void mlWithException() { + void ml_with_exception() { assertThrows(UnsupportedOperationException.class, () -> client.getNodeClient()); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java index 9006a0573d9..c13e63f01ac 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -5,6 +5,7 @@ package org.opensearch.sql.opensearch.storage.scan; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -12,6 +13,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.withSettings; import static org.opensearch.sql.data.type.ExprCoreType.STRING; @@ -31,6 +33,7 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.PagedRequestBuilder; import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; @ExtendWith(MockitoExtension.class) @@ -64,19 +67,31 @@ void query_all_results_initial_scroll_request() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - InitialPageRequestBuilder builder = new InitialPageRequestBuilder( + PagedRequestBuilder builder = new InitialPageRequestBuilder( new OpenSearchRequest.IndexName("test"), 3, settings, exprValueFactory); try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + + () -> assertFalse(indexScan.hasNext()) + ); + } + // cleanup should be called on empty response only + verify(client, never()).cleanup(any()); - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + builder = new SubsequentPageRequestBuilder( + new OpenSearchRequest.IndexName("test"), "scroll", exprValueFactory); + try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { + indexScan.open(); assertFalse(indexScan.hasNext()); } @@ -95,14 +110,26 @@ void query_all_results_continuation_scroll_request() { try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + + () -> assertFalse(indexScan.hasNext()) + ); + } + // cleanup should be called on empty response only + verify(client, never()).cleanup(any()); + + builder = new SubsequentPageRequestBuilder( + new OpenSearchRequest.IndexName("test"), "scroll", exprValueFactory); + try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { + indexScan.open(); assertFalse(indexScan.hasNext()); } From f4ea4ad8cc4afcba798e8bdccffb93688ad854bb Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Mar 2023 16:35:43 -0800 Subject: [PATCH 24/46] Store `TotalHits` and use it to report `total` in response. Signed-off-by: Yury-Fridlyand --- .../sql/executor/ExecutionEngine.java | 2 +- .../sql/planner/physical/PhysicalPlan.java | 4 + .../sql/executor/QueryServiceTest.java | 2 +- .../MicroBatchStreamingExecutionTest.java | 2 +- .../planner/physical/PhysicalPlanTest.java | 27 ++++- .../sql/executor/DefaultExecutionEngine.java | 2 +- .../sql/legacy/plugin/RestSQLQueryAction.java | 3 +- .../executor/OpenSearchExecutionEngine.java | 2 +- .../protector/ResourceMonitorPlan.java | 5 + .../response/OpenSearchResponse.java | 5 +- .../storage/scan/OpenSearchIndexScan.java | 6 ++ .../scan/OpenSearchPagedIndexScan.java | 7 ++ .../executor/ResourceMonitorPlanTest.java | 6 ++ .../response/OpenSearchResponseTest.java | 23 +++-- .../storage/scan/OpenSearchIndexScanTest.java | 99 ++++++++++++------- .../scan/OpenSearchPagedIndexScanTest.java | 6 +- .../transport/TransportPPLQueryAction.java | 2 +- .../sql/protocol/response/QueryResult.java | 5 +- .../format/JdbcResponseFormatter.java | 2 +- .../protocol/response/QueryResultTest.java | 10 +- 20 files changed, 158 insertions(+), 62 deletions(-) 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 1c78df704cc..2d98be6bbbd 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java @@ -54,7 +54,7 @@ void execute(PhysicalPlan plan, ExecutionContext context, class QueryResponse { private final Schema schema; private final List results; - + private final long total; private final Cursor cursor; } 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 5890b6f15f9..c52e658e41b 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 @@ -50,6 +50,10 @@ public ExecutionEngine.Schema schema() { + "ProjectOperator, instead of %s", this.getClass().getSimpleName())); } + public long getTotalHits() { + return getChild().stream().mapToLong(PhysicalPlan::getTotalHits).max().orElse(0); + } + public String toCursor() { throw new IllegalStateException(String.format("%s is not compatible with cursor feature", this.getClass().getSimpleName())); diff --git a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java b/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java index a0be4f8f2ec..69c819398c0 100644 --- a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java @@ -133,7 +133,7 @@ Helper executeSuccess(Split split) { invocation -> { ResponseListener listener = invocation.getArgument(2); listener.onResponse( - new ExecutionEngine.QueryResponse(schema, Collections.emptyList(), + new ExecutionEngine.QueryResponse(schema, Collections.emptyList(), 0, Cursor.None)); return null; }) diff --git a/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java b/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java index 75a62385308..f97f2b5f91b 100644 --- a/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java @@ -170,7 +170,7 @@ Helper executeSuccess(Long... offsets) { ResponseListener listener = invocation.getArgument(2); listener.onResponse( - new ExecutionEngine.QueryResponse(null, Collections.emptyList(), + new ExecutionEngine.QueryResponse(null, Collections.emptyList(), 0, Cursor.None)); PlanContext planContext = invocation.getArgument(1); 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..c6759b7f7c9 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,9 +5,16 @@ package org.opensearch.sql.planner.physical; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import java.util.List; + +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -16,6 +23,7 @@ import org.opensearch.sql.storage.split.Split; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class PhysicalPlanTest { @Mock Split split; @@ -46,8 +54,25 @@ public List getChild() { }; @Test - void addSplitToChildByDefault() { + void add_split_to_child_by_default() { testPlan.add(split); verify(child).add(split); } + + @Test + void get_total_hits_from_child() { + var plan = mock(PhysicalPlan.class); + when(child.getTotalHits()).thenReturn(42L); + when(plan.getChild()).thenReturn(List.of(child)); + when(plan.getTotalHits()).then(CALLS_REAL_METHODS); + assertEquals(42, plan.getTotalHits()); + verify(child).getTotalHits(); + } + + @Test + void get_total_hits_uses_default_value() { + var plan = mock(PhysicalPlan.class); + when(plan.getTotalHits()).then(CALLS_REAL_METHODS); + assertEquals(0, plan.getTotalHits()); + } } diff --git a/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java b/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java index 18058302711..00e02eb433e 100644 --- a/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java +++ b/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java @@ -34,7 +34,7 @@ public void execute( result.add(plan.next()); } QueryResponse response = new QueryResponse(new Schema(new ArrayList<>()), new ArrayList<>(), - Cursor.None); + 0, Cursor.None); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); 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 bd9387a68e0..4f9fdd9a535 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 @@ -172,7 +172,8 @@ private ResponseListener createQueryResponseListener( @Override public void onResponse(QueryResponse response) { sendResponse(channel, OK, - formatter.format(new QueryResult(response.getSchema(), response.getResults(), response.getCursor()))); + formatter.format(new QueryResult(response.getSchema(), response.getResults(), + response.getCursor(), response.getTotal()))); } @Override 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 cec2864c11f..b1b32821f25 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 @@ -53,7 +53,7 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context, Cursor qc = paginatedPlanCache.convertToCursor(plan); - QueryResponse response = new QueryResponse(physicalPlan.schema(), result, qc); + QueryResponse response = new QueryResponse(physicalPlan.schema(), result, plan.getTotalHits(), qc); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index 307b40dce79..3d880d82b9f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -83,6 +83,11 @@ public ExprValue next() { return delegate.next(); } + @Override + public long getTotalHits() { + return delegate.getTotalHits(); + } + @Override public String toCursor() { return delegate.toCursor(); 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 c2042d7c0ea..74aa07fccb6 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 @@ -71,7 +71,10 @@ public OpenSearchResponse(SearchHits hits, OpenSearchExprValueFactory exprValueF */ public boolean isEmpty() { return (hits.getHits() == null) || (hits.getHits().length == 0) && aggregations == null; - // TODO TBD ^ ^ + } + + public long getTotalHits() { + return hits.getTotalHits().value; } public boolean isAggregationResponse() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 3ae2e62cfd5..f9a420332d1 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -84,6 +84,12 @@ public ExprValue next() { return iterator.next(); } + @Override + public long getTotalHits() { + // TODO maybe store totalHits from `response` + return queryCount; + } + private void fetchNextBatch() { OpenSearchResponse response = client.search(request); if (!response.isEmpty()) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index 1dc455cfd2d..6626af6e9eb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -27,6 +27,7 @@ public class OpenSearchPagedIndexScan extends TableScanOperator { private OpenSearchRequest request; private Iterator iterator; private boolean needClean = false; + private long totalHits = 0; public OpenSearchPagedIndexScan(OpenSearchClient client, PagedRequestBuilder requestBuilder) { @@ -56,6 +57,7 @@ public void open() { OpenSearchResponse response = client.search(request); if (!response.isEmpty()) { iterator = response.iterator(); + totalHits = response.getTotalHits(); } else { needClean = true; iterator = Collections.emptyIterator(); @@ -71,6 +73,11 @@ public void close() { } } + @Override + public long getTotalHits() { + return totalHits; + } + @Override public String toCursor() { // TODO this assumes exactly one index is scanned. diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java index d4d987a7df2..b111047b6fd 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java @@ -107,4 +107,10 @@ void acceptSuccess() { monitorPlan.accept(visitor, context); verify(plan, times(1)).accept(visitor, context); } + + @Test + void getTotalHitsSuccess() { + monitorPlan.getTotalHits(); + verify(plan, times(1)).getTotalHits(); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java index 0a60503415d..2d1d6145f3f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java @@ -74,20 +74,29 @@ void isEmpty() { new TotalHits(2L, TotalHits.Relation.EQUAL_TO), 1.0F)); - assertFalse(new OpenSearchResponse(searchResponse, factory).isEmpty()); + var response = new OpenSearchResponse(searchResponse, factory); + assertFalse(response.isEmpty()); + assertEquals(2L, response.getTotalHits()); when(searchResponse.getHits()).thenReturn(SearchHits.empty()); when(searchResponse.getAggregations()).thenReturn(null); - assertTrue(new OpenSearchResponse(searchResponse, factory).isEmpty()); + + response = new OpenSearchResponse(searchResponse, factory); + assertTrue(response.isEmpty()); + assertEquals(0L, response.getTotalHits()); when(searchResponse.getHits()) .thenReturn(new SearchHits(null, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0)); - OpenSearchResponse response3 = new OpenSearchResponse(searchResponse, factory); - assertTrue(response3.isEmpty()); + response = new OpenSearchResponse(searchResponse, factory); + assertTrue(response.isEmpty()); + assertEquals(0L, response.getTotalHits()); when(searchResponse.getHits()).thenReturn(SearchHits.empty()); when(searchResponse.getAggregations()).thenReturn(new Aggregations(emptyList())); - assertFalse(new OpenSearchResponse(searchResponse, factory).isEmpty()); + + response = new OpenSearchResponse(searchResponse, factory); + assertFalse(response.isEmpty()); + assertEquals(0L, response.getTotalHits()); } @Test @@ -104,7 +113,8 @@ void iterator() { when(factory.construct(any())).thenReturn(exprTupleValue1).thenReturn(exprTupleValue2); int i = 0; - for (ExprValue hit : new OpenSearchResponse(searchResponse, factory)) { + var response = new OpenSearchResponse(searchResponse, factory); + for (ExprValue hit : response) { if (i == 0) { assertEquals(exprTupleValue1, hit); } else if (i == 1) { @@ -114,6 +124,7 @@ void iterator() { } i++; } + assertEquals(2L, response.getTotalHits()); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 90ad624135b..d93f4729b8b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -6,10 +6,12 @@ package org.opensearch.sql.opensearch.storage.scan; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -22,6 +24,8 @@ import java.util.HashMap; import java.util.Map; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -47,6 +51,7 @@ import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchIndexScanTest { @Mock @@ -64,19 +69,22 @@ void setup() { } @Test - void queryEmptyResult() { + void query_empty_result() { mockResponse(client); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("test", 3, settings, exprValueFactory))) { indexScan.open(); - assertFalse(indexScan.hasNext()); + assertAll( + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(0, indexScan.getTotalHits()) + ); } verify(client).cleanup(any()); } @Test - void queryAllResultsWithQuery() { + void query_all_results_with_query() { mockResponse(client, new ExprValue[]{ employee(1, "John", "IT"), employee(2, "Smith", "HR"), @@ -89,22 +97,25 @@ void queryAllResultsWithQuery() { new OpenSearchIndexScan(client, builder)) { indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - assertFalse(indexScan.hasNext()); + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); } verify(client).cleanup(any()); } @Test - void queryAllResultsWithScroll() { + void query_all_results_with_scroll() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); @@ -114,22 +125,25 @@ void queryAllResultsWithScroll() { exprValueFactory))) { indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - assertFalse(indexScan.hasNext()); + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); } verify(client).cleanup(any()); } @Test - void querySomeResultsWithQuery() { + void query_some_results_with_query() { mockResponse(client, new ExprValue[]{ employee(1, "John", "IT"), employee(2, "Smith", "HR"), @@ -142,22 +156,25 @@ void querySomeResultsWithQuery() { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - assertFalse(indexScan.hasNext()); + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); } verify(client).cleanup(any()); } @Test - void querySomeResultsWithScroll() { + void query_some_results_with_scroll() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); @@ -168,22 +185,25 @@ void querySomeResultsWithScroll() { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); - assertTrue(indexScan.hasNext()); - assertEquals(employee(1, "John", "IT"), indexScan.next()); + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(2, "Smith", "HR"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - assertTrue(indexScan.hasNext()); - assertEquals(employee(3, "Allen", "IT"), indexScan.next()); + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - assertFalse(indexScan.hasNext()); + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); } verify(client).cleanup(any()); } @Test - void pushDownFilters() { + void push_down_filters() { assertThat() .pushDown(QueryBuilders.termQuery("name", "John")) .shouldQuery(QueryBuilders.termQuery("name", "John")) @@ -201,7 +221,7 @@ void pushDownFilters() { } @Test - void pushDownHighlight() { + void push_down_highlight() { Map args = new HashMap<>(); assertThat() .pushDown(QueryBuilders.termQuery("name", "John")) @@ -212,7 +232,7 @@ void pushDownHighlight() { } @Test - void pushDownHighlightWithArguments() { + void push_down_highlight_with_arguments() { Map args = new HashMap<>(); args.put("pre_tags", new Literal("", DataType.STRING)); args.put("post_tags", new Literal("", DataType.STRING)); @@ -227,7 +247,7 @@ void pushDownHighlightWithArguments() { } @Test - void pushDownHighlightWithRepeatingFields() { + void push_down_highlight_with_repeating_fields() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); @@ -314,6 +334,9 @@ public OpenSearchResponse answer(InvocationOnMock invocation) { when(response.isEmpty()).thenReturn(false); ExprValue[] searchHit = searchHitBatches[batchNum]; when(response.iterator()).thenReturn(Arrays.asList(searchHit).iterator()); + // used in OpenSearchPagedIndexScanTest + lenient().when(response.getTotalHits()) + .thenReturn((long) searchHitBatches[batchNum].length); } else { when(response.isEmpty()).thenReturn(true); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java index c13e63f01ac..fc77ad4ff7c 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -82,7 +82,8 @@ void query_all_results_initial_scroll_request() { () -> assertTrue(indexScan.hasNext()), () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - () -> assertFalse(indexScan.hasNext()) + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) ); } // cleanup should be called on empty response only @@ -120,7 +121,8 @@ void query_all_results_continuation_scroll_request() { () -> assertTrue(indexScan.hasNext()), () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - () -> assertFalse(indexScan.hasNext()) + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) ); } // cleanup should be called on empty response only diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index d9dad9d5353..a67e077ecc5 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -140,7 +140,7 @@ private ResponseListener createListener( public void onResponse(ExecutionEngine.QueryResponse response) { String responseContent = formatter.format(new QueryResult(response.getSchema(), response.getResults(), - response.getCursor())); + response.getCursor(), response.getTotal())); listener.onResponse(new TransportPPLQueryResponse(responseContent)); } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java index 3ea5846b87c..90f422c81d2 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java @@ -36,9 +36,12 @@ public class QueryResult implements Iterable { @Getter private final Cursor cursor; + @Getter + private final long total; + public QueryResult(ExecutionEngine.Schema schema, Collection exprValues) { - this(schema, exprValues, Cursor.None); + this(schema, exprValues, Cursor.None, exprValues.size()); } /** diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java index f52ee222467..eb4de495c1b 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java @@ -40,7 +40,7 @@ protected Object buildJsonObject(QueryResult response) { json.datarows(fetchDataRows(response)); // Populate other fields - json.total(response.size()) + json.total(response.getTotal()) .size(response.size()) .status(200); if (!response.getCursor().equals(Cursor.None)) { diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java index 3db405339bc..2bfbe552784 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java @@ -36,7 +36,7 @@ void size() { tupleValue(ImmutableMap.of("name", "John", "age", 20)), tupleValue(ImmutableMap.of("name", "Allen", "age", 30)), tupleValue(ImmutableMap.of("name", "Smith", "age", 40)) - ), Cursor.None); + ), Cursor.None, 0); assertEquals(3, response.size()); } @@ -46,7 +46,7 @@ void columnNameTypes() { schema, Collections.singletonList( tupleValue(ImmutableMap.of("name", "John", "age", 20)) - ), Cursor.None); + ), Cursor.None, 0); assertEquals( ImmutableMap.of("name", "string", "age", "integer"), @@ -61,7 +61,7 @@ void columnNameTypesWithAlias() { QueryResult response = new QueryResult( schema, Collections.singletonList(tupleValue(ImmutableMap.of("n", "John"))), - Cursor.None); + Cursor.None, 0); assertEquals( ImmutableMap.of("n", "string"), @@ -73,7 +73,7 @@ void columnNameTypesWithAlias() { void columnNameTypesFromEmptyExprValues() { QueryResult response = new QueryResult( schema, - Collections.emptyList(), Cursor.None); + Collections.emptyList(), Cursor.None, 0); assertEquals( ImmutableMap.of("name", "string", "age", "integer"), response.columnNameTypes() @@ -102,7 +102,7 @@ void iterate() { Arrays.asList( tupleValue(ImmutableMap.of("name", "John", "age", 20)), tupleValue(ImmutableMap.of("name", "Allen", "age", 30)) - ), Cursor.None); + ), Cursor.None, 0); int i = 0; for (Object[] objects : response) { From 7f0acdd32447af5447766703ce988a1ba91c03dd Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Mar 2023 16:44:43 -0800 Subject: [PATCH 25/46] Add missing UT for `:protocol` module. Signed-off-by: Yury-Fridlyand --- .../format/JdbcResponseFormatterTest.java | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java index 96b383ac0ca..726fd7b2770 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java @@ -32,6 +32,7 @@ import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.exception.SemanticCheckException; +import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.protocol.response.QueryResult; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -81,6 +82,37 @@ void format_response() { formatter.format(response)); } + @Test + void format_response_with_cursor() { + QueryResult response = new QueryResult( + new Schema(ImmutableList.of( + new Column("name", "name", STRING), + new Column("address", "address", OPENSEARCH_TEXT), + new Column("age", "age", INTEGER))), + ImmutableList.of( + tupleValue(ImmutableMap.builder() + .put("name", "John") + .put("address", "Seattle") + .put("age", 20) + .build())), + new Cursor("test_cursor".getBytes()), 42); + + assertJsonEquals( + "{" + + "\"schema\":[" + + "{\"name\":\"name\",\"alias\":\"name\",\"type\":\"keyword\"}," + + "{\"name\":\"address\",\"alias\":\"address\",\"type\":\"text\"}," + + "{\"name\":\"age\",\"alias\":\"age\",\"type\":\"integer\"}" + + "]," + + "\"datarows\":[" + + "[\"John\",\"Seattle\",20]]," + + "\"total\":42," + + "\"size\":1," + + "\"cursor\":\"test_cursor\"," + + "\"status\":200}", + formatter.format(response)); + } + @Test void format_response_with_missing_and_null_value() { QueryResult response = From 2ce16263c4a37da95330ac18925aff49a7d42769 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Mar 2023 16:50:53 -0800 Subject: [PATCH 26/46] Fix PPL UTs damaged in f4ea4ad8c. Signed-off-by: Yury-Fridlyand --- .../test/java/org/opensearch/sql/ppl/PPLServiceTest.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 774143a3484..833de7b362f 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -75,7 +75,7 @@ public void cleanup() throws InterruptedException { public void testExecuteShouldPass() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), 0, Cursor.None)); return null; }).when(queryService).execute(any(), any()); @@ -97,7 +97,7 @@ public void onFailure(Exception e) { public void testExecuteCsvFormatShouldPass() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), 0, Cursor.None)); return null; }).when(queryService).execute(any(), any()); @@ -171,7 +171,7 @@ public void onFailure(Exception e) { public void testPrometheusQuery() { doAnswer(invocation -> { ResponseListener listener = invocation.getArgument(1); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), 0, Cursor.None)); return null; }).when(queryService).execute(any(), any()); From b2e6e56371a0d99270856dd9aa0b8898ba304bab Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Mar 2023 16:57:57 -0800 Subject: [PATCH 27/46] Minor checkstyle fixes. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/planner/physical/PhysicalPlanTest.java | 1 - .../sql/opensearch/executor/OpenSearchExecutionEngine.java | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) 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 c6759b7f7c9..2c02c9b9cae 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 @@ -12,7 +12,6 @@ import static org.mockito.Mockito.when; import java.util.List; - import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; 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 b1b32821f25..0d1b80fbcd0 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 @@ -53,7 +53,8 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context, Cursor qc = paginatedPlanCache.convertToCursor(plan); - QueryResponse response = new QueryResponse(physicalPlan.schema(), result, plan.getTotalHits(), qc); + QueryResponse response = new QueryResponse(physicalPlan.schema(), result, + plan.getTotalHits(), qc); listener.onResponse(response); } catch (Exception e) { listener.onFailure(e); From c7ad2193f22cc3bc760cc9b0e487411f9a653805 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Mon, 13 Mar 2023 12:08:28 -0700 Subject: [PATCH 28/46] Fallback to v1 engine for pagination (#245) * Pagination fallback integration tests. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/legacy/CursorIT.java | 29 +++- .../sql/sql/PaginationFallbackIT.java | 135 ++++++++++++++++++ .../org/opensearch/sql/sql/PaginationIT.java | 48 +++++++ .../org/opensearch/sql/util/TestUtils.java | 23 +++ .../RestSQLQueryActionCursorFallbackTest.java | 128 +++++++++++++++++ 5 files changed, 361 insertions(+), 2 deletions(-) create mode 100644 integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/sql/PaginationIT.java create mode 100644 legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionCursorFallbackTest.java diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java index 113a19885a7..5b9a583d04b 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java @@ -123,11 +123,16 @@ public void validNumberOfPages() throws IOException { String selectQuery = StringUtils.format("SELECT firstname, state FROM %s", TEST_INDEX_ACCOUNT); JSONObject response = new JSONObject(executeFetchQuery(selectQuery, 50, JDBC)); String cursor = response.getString(CURSOR); + verifyIsV1Cursor(cursor); + int pageCount = 1; while (!cursor.isEmpty()) { //this condition also checks that there is no cursor on last page response = executeCursorQuery(cursor); cursor = response.optString(CURSOR); + if (!cursor.isEmpty()) { + verifyIsV1Cursor(cursor); + } pageCount++; } @@ -136,12 +141,16 @@ public void validNumberOfPages() throws IOException { // using random value here, with fetch size of 28 we should get 36 pages (ceil of 1000/28) response = new JSONObject(executeFetchQuery(selectQuery, 28, JDBC)); cursor = response.getString(CURSOR); + verifyIsV1Cursor(cursor); System.out.println(response); pageCount = 1; while (!cursor.isEmpty()) { response = executeCursorQuery(cursor); cursor = response.optString(CURSOR); + if (!cursor.isEmpty()) { + verifyIsV1Cursor(cursor); + } pageCount++; } assertThat(pageCount, equalTo(36)); @@ -223,6 +232,7 @@ public void testCursorWithPreparedStatement() throws IOException { "}", TestsConstants.TEST_INDEX_ACCOUNT)); assertTrue(response.has(CURSOR)); + verifyIsV1Cursor(response.getString(CURSOR)); } @Test @@ -244,11 +254,13 @@ public void testRegressionOnDateFormatChange() throws IOException { StringUtils.format("SELECT login_time FROM %s LIMIT 500", TEST_INDEX_DATE_TIME); JSONObject response = new JSONObject(executeFetchQuery(selectQuery, 1, JDBC)); String cursor = response.getString(CURSOR); + verifyIsV1Cursor(cursor); actualDateList.add(response.getJSONArray(DATAROWS).getJSONArray(0).getString(0)); while (!cursor.isEmpty()) { response = executeCursorQuery(cursor); cursor = response.optString(CURSOR); + verifyIsV1Cursor(cursor); actualDateList.add(response.getJSONArray(DATAROWS).getJSONArray(0).getString(0)); } @@ -274,7 +286,6 @@ public void defaultBehaviorWhenCursorSettingIsDisabled() throws IOException { query = StringUtils.format("SELECT firstname, email, state FROM %s", TEST_INDEX_ACCOUNT); response = new JSONObject(executeFetchQuery(query, 100, JDBC)); assertTrue(response.has(CURSOR)); - wipeAllClusterSettings(); } @@ -305,12 +316,14 @@ public void testDefaultFetchSizeFromClusterSettings() throws IOException { JSONObject response = new JSONObject(executeFetchLessQuery(query, JDBC)); JSONArray datawRows = response.optJSONArray(DATAROWS); assertThat(datawRows.length(), equalTo(1000)); + verifyIsV1Cursor(response.getString(CURSOR)); updateClusterSettings(new ClusterSetting(TRANSIENT, "opensearch.sql.cursor.fetch_size", "786")); response = new JSONObject(executeFetchLessQuery(query, JDBC)); datawRows = response.optJSONArray(DATAROWS); assertThat(datawRows.length(), equalTo(786)); assertTrue(response.has(CURSOR)); + verifyIsV1Cursor(response.getString(CURSOR)); wipeAllClusterSettings(); } @@ -323,11 +336,12 @@ public void testCursorCloseAPI() throws IOException { "SELECT firstname, state FROM %s WHERE balance > 100 and age < 40", TEST_INDEX_ACCOUNT); JSONObject result = new JSONObject(executeFetchQuery(selectQuery, 50, JDBC)); String cursor = result.getString(CURSOR); - + verifyIsV1Cursor(cursor); // Retrieving next 10 pages out of remaining 19 pages for (int i = 0; i < 10; i++) { result = executeCursorQuery(cursor); cursor = result.optString(CURSOR); + verifyIsV1Cursor(cursor); } //Closing the cursor JSONObject closeResp = executeCursorCloseQuery(cursor); @@ -386,12 +400,14 @@ public void respectLimitPassedInSelectClause() throws IOException { StringUtils.format("SELECT age, balance FROM %s LIMIT %s", TEST_INDEX_ACCOUNT, limit); JSONObject response = new JSONObject(executeFetchQuery(selectQuery, 50, JDBC)); String cursor = response.getString(CURSOR); + verifyIsV1Cursor(cursor); int actualDataRowCount = response.getJSONArray(DATAROWS).length(); int pageCount = 1; while (!cursor.isEmpty()) { response = executeCursorQuery(cursor); cursor = response.optString(CURSOR); + verifyIsV1Cursor(cursor); actualDataRowCount += response.getJSONArray(DATAROWS).length(); pageCount++; } @@ -432,10 +448,12 @@ public void verifyWithAndWithoutPaginationResponse(String sqlQuery, String curso response.optJSONArray(DATAROWS).forEach(dataRows::put); String cursor = response.getString(CURSOR); + verifyIsV1Cursor(cursor); while (!cursor.isEmpty()) { response = executeCursorQuery(cursor); response.optJSONArray(DATAROWS).forEach(dataRows::put); cursor = response.optString(CURSOR); + verifyIsV1Cursor(cursor); } verifySchema(withoutCursorResponse.optJSONArray(SCHEMA), @@ -465,6 +483,13 @@ public String executeFetchAsStringQuery(String query, String fetchSize, String r return responseString; } + private void verifyIsV1Cursor(String cursor) { + if (cursor.isEmpty()) { + return; + } + assertTrue("The cursor '" + cursor + "' is not from v1 engine.", cursor.startsWith("d:")); + } + private String makeRequest(String query, String fetch_size) { return String.format("{" + " \"fetch_size\": \"%s\"," + diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java new file mode 100644 index 00000000000..19d6afb16d5 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java @@ -0,0 +1,135 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.sql; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ONLINE; +import static org.opensearch.sql.util.TestUtils.verifyIsV1Cursor; +import static org.opensearch.sql.util.TestUtils.verifyIsV2Cursor; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.Test; +import org.opensearch.sql.legacy.SQLIntegTestCase; +import org.opensearch.sql.util.TestUtils; + +public class PaginationFallbackIT extends SQLIntegTestCase { + @Override + public void init() throws IOException { + loadIndex(Index.PHRASE); + loadIndex(Index.ONLINE); + } + + @Test + public void testWhereClause() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s WHERE 1 = 1", TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testSelectAll() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_ONLINE); + verifyIsV2Cursor(response); + } + + @Test + public void testSelectWithOpenSearchFuncInFilter() throws IOException { + var response = executeQueryTemplate( + "SELECT * FROM %s WHERE `11` = match_phrase('96')", TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testSelectWithHighlight() throws IOException { + var response = executeQueryTemplate( + "SELECT highlight(`11`) FROM %s WHERE match_query(`11`, '96')", TEST_INDEX_ONLINE); + // As of 2023-03-08, WHERE clause sends the query to legacy engine and legacy engine + // does not support highlight as an expression. + assertTrue(response.has("error")); + } + + @Test + public void testSelectWithFullTextSearch() throws IOException { + var response = executeQueryTemplate( + "SELECT * FROM %s WHERE match_phrase(`11`, '96')", TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testSelectFromIndexWildcard() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s*", TEST_INDEX); + verifyIsV2Cursor(response); + } + + @Test + public void testSelectFromDataSource() throws IOException { + var response = executeQueryTemplate("SELECT * FROM @opensearch.%s", + TEST_INDEX_ONLINE); + verifyIsV2Cursor(response); + } + + @Test + public void testSelectColumnReference() throws IOException { + var response = executeQueryTemplate("SELECT `107` from %s", TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testSubquery() throws IOException { + var response = executeQueryTemplate("SELECT `107` from (SELECT * FROM %s)", + TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testSelectExpression() throws IOException { + var response = executeQueryTemplate("SELECT 1 + 1 - `107` from %s", + TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testGroupBy() throws IOException { + // GROUP BY is not paged by either engine. + var response = executeQueryTemplate("SELECT * FROM %s GROUP BY `107`", + TEST_INDEX_ONLINE); + TestUtils.verifyNoCursor(response); + } + + @Test + public void testGroupByHaving() throws IOException { + // GROUP BY is not paged by either engine. + var response = executeQueryTemplate("SELECT * FROM %s GROUP BY `107` HAVING `107` > 400", + TEST_INDEX_ONLINE); + TestUtils.verifyNoCursor(response); + } + + @Test + public void testLimit() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s LIMIT 8", TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testLimitOffset() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s LIMIT 8 OFFSET 4", + TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + @Test + public void testOrderBy() throws IOException { + var response = executeQueryTemplate("SELECT * FROM %s ORDER By `107`", + TEST_INDEX_ONLINE); + verifyIsV1Cursor(response); + } + + private JSONObject executeQueryTemplate(String queryTemplate, String index) throws IOException { + var query = String.format(queryTemplate, index); + return new JSONObject(executeFetchQuery(query, 4, "jdbc")); + } + +} diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationIT.java new file mode 100644 index 00000000000..b9e32cb1cdd --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationIT.java @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.sql; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_CALCS; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ONLINE; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.Test; +import org.opensearch.sql.legacy.SQLIntegTestCase; +import org.opensearch.sql.util.TestUtils; + +public class PaginationIT extends SQLIntegTestCase { + @Override + public void init() throws IOException { + loadIndex(Index.CALCS); + loadIndex(Index.ONLINE); + } + + @Test + public void testSmallDataSet() throws IOException { + var query = "SELECT * from " + TEST_INDEX_CALCS; + var response = new JSONObject(executeFetchQuery(query, 4, "jdbc")); + assertTrue(response.has("cursor")); + assertEquals(4, response.getInt("size")); + TestUtils.verifyIsV2Cursor(response); + } + + @Test + public void testLargeDataSetV1() throws IOException { + var v1query = "SELECT * from " + TEST_INDEX_ONLINE + " WHERE 1 = 1"; + var v1response = new JSONObject(executeFetchQuery(v1query, 4, "jdbc")); + assertEquals(4, v1response.getInt("size")); + TestUtils.verifyIsV1Cursor(v1response); + } + + @Test + public void testLargeDataSetV2() throws IOException { + var query = "SELECT * from " + TEST_INDEX_ONLINE; + var response = new JSONObject(executeFetchQuery(query, 4, "jdbc")); + assertEquals(4, response.getInt("size")); + TestUtils.verifyIsV2Cursor(response); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/util/TestUtils.java b/integ-test/src/test/java/org/opensearch/sql/util/TestUtils.java index bd75ead43b6..34ab1729d46 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/TestUtils.java +++ b/integ-test/src/test/java/org/opensearch/sql/util/TestUtils.java @@ -7,6 +7,7 @@ package org.opensearch.sql.util; import static com.google.common.base.Strings.isNullOrEmpty; +import static org.junit.Assert.assertTrue; import java.io.BufferedReader; import java.io.File; @@ -839,4 +840,26 @@ public static List> getPermutations(final List items) { return result; } + + public static void verifyIsV1Cursor(JSONObject response) { + verifyCursor(response, List.of("d:", "j:", "a:"), "v1"); + } + + + public static void verifyIsV2Cursor(JSONObject response) { + verifyCursor(response, List.of("n:"), "v2"); + } + + private static void verifyCursor(JSONObject response, List validCursorPrefix, String engineName) { + assertTrue("'cursor' property does not exist", response.has("cursor")); + + var cursor = response.getString("cursor"); + assertTrue("'cursor' property is empty", !cursor.isEmpty()); + assertTrue("The cursor '" + cursor + "' is not from " + engineName + " engine.", + validCursorPrefix.stream().anyMatch(cursor::startsWith)); + } + + public static void verifyNoCursor(JSONObject response) { + assertTrue(!response.has("cursor")); + } } diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionCursorFallbackTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionCursorFallbackTest.java new file mode 100644 index 00000000000..20d1354825f --- /dev/null +++ b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionCursorFallbackTest.java @@ -0,0 +1,128 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.legacy.plugin; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.opensearch.sql.legacy.plugin.RestSqlAction.QUERY_API_ENDPOINT; + +import java.io.IOException; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import org.json.JSONObject; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; +import org.opensearch.client.node.NodeClient; +import org.opensearch.common.Strings; +import org.opensearch.common.inject.Injector; +import org.opensearch.common.inject.ModulesBuilder; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.sql.common.antlr.SyntaxCheckException; +import org.opensearch.sql.executor.QueryManager; +import org.opensearch.sql.executor.execution.QueryPlanFactory; +import org.opensearch.sql.sql.SQLService; +import org.opensearch.sql.sql.antlr.SQLSyntaxParser; +import org.opensearch.sql.sql.domain.SQLQueryRequest; +import org.opensearch.threadpool.ThreadPool; + +/** + * A test suite that verifies fallback behaviour of cursor queries. + */ +@RunWith(MockitoJUnitRunner.class) +public class RestSQLQueryActionCursorFallbackTest extends BaseRestHandler { + + private NodeClient nodeClient; + + @Mock + private ThreadPool threadPool; + + @Mock + private QueryManager queryManager; + + @Mock + private QueryPlanFactory factory; + + @Mock + private RestChannel restChannel; + + private Injector injector; + + @Before + public void setup() { + nodeClient = new NodeClient(org.opensearch.common.settings.Settings.EMPTY, threadPool); + ModulesBuilder modules = new ModulesBuilder(); + modules.add(b -> { + b.bind(SQLService.class).toInstance(new SQLService(new SQLSyntaxParser(), queryManager, factory)); + }); + injector = modules.createInjector(); + Mockito.lenient().when(threadPool.getThreadContext()) + .thenReturn(new ThreadContext(org.opensearch.common.settings.Settings.EMPTY)); + } + + // Initial page request test cases + + @Test + public void no_fallback_with_column_reference() throws Exception { + String query = "SELECT name FROM test1"; + SQLQueryRequest request = createSqlQueryRequest(query, Optional.empty(), + Optional.of(5)); + + assertFalse(doesQueryFallback(request)); + } + + private static SQLQueryRequest createSqlQueryRequest(String query, Optional cursorId, + Optional fetchSize) throws IOException { + var builder = XContentFactory.jsonBuilder() + .startObject() + .field("query").value(query); + if (cursorId.isPresent()) { + builder.field("cursor").value(cursorId.get()); + } + + if (fetchSize.isPresent()) { + builder.field("fetch_size").value(fetchSize.get()); + } + builder.endObject(); + JSONObject jsonContent = new JSONObject(Strings.toString(builder)); + + return new SQLQueryRequest(jsonContent, query, QUERY_API_ENDPOINT, + Map.of("format", "jdbc"), cursorId.orElse("")); + } + + boolean doesQueryFallback(SQLQueryRequest request) throws Exception { + AtomicBoolean fallback = new AtomicBoolean(false); + RestSQLQueryAction queryAction = new RestSQLQueryAction(injector); + queryAction.prepareRequest(request, (channel, exception) -> { + fallback.set(true); + assertTrue(exception instanceof SyntaxCheckException); + }, (channel, exception) -> { + }).accept(restChannel); + return fallback.get(); + } + + @Override + public String getName() { + // do nothing, RestChannelConsumer is protected which required to extend BaseRestHandler + return null; + } + + @Override + protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient nodeClient) + { + // do nothing, RestChannelConsumer is protected which required to extend BaseRestHandler + return null; + } +} From 981bc258f6ce39e9ed9f8aa83ff50d55eee48863 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Mon, 13 Mar 2023 12:48:20 -0700 Subject: [PATCH 29/46] Add UT with coverage for `toCursor` serialization. Signed-off-by: Yury-Fridlyand --- .../sql/planner/physical/PhysicalPlan.java | 11 -------- .../physical/PaginateOperatorTest.java | 14 ++++++++++ .../planner/physical/PhysicalPlanTest.java | 19 +++++++++++++ .../planner/physical/ProjectOperatorTest.java | 27 ++++++++++++++----- .../executor/ResourceMonitorPlanTest.java | 6 +++++ .../request/OpenSearchRequestTest.java | 23 ++++++++++++++++ .../request/OpenSearchScrollRequestTest.java | 9 +++++++ .../scan/OpenSearchPagedIndexScanTest.java | 21 +++++++++++++++ 8 files changed, 112 insertions(+), 18 deletions(-) create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java 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 c52e658e41b..d24e3203dc7 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 @@ -19,9 +19,6 @@ public abstract class PhysicalPlan implements PlanNode, Iterator, AutoCloseable { - - public static final List FORBIDDEN_CHARS = List.of("(", ")", ","); - /** * Accept the {@link PhysicalPlanNodeVisitor}. * @@ -66,14 +63,6 @@ public String toCursor() { * @return A string that represents the plan called with those parameters. */ protected String createSection(String plan, String... params) { - if (FORBIDDEN_CHARS.stream().anyMatch(plan::contains)) { - var error = String.format("plan key '%s' contains forbidden character", - plan); - throw new RuntimeException(error); - } - - // TODO: check that each param is either a valid s-expression or - // does not contain forbidden characters. return "(" + plan + "," + String.join(",", params) + ")"; diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java index e91766b1a26..01d525978f0 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java @@ -6,6 +6,7 @@ package org.opensearch.sql.planner.physical; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; @@ -83,4 +84,17 @@ public void schema_assert() { assertThrows(Throwable.class, () -> new PaginateOperator(mock(PhysicalPlan.class), 42).schema()); } + + @Test + public void toCursor() { + var plan = mock(PhysicalPlan.class); + when(plan.toCursor()).thenReturn("Great plan, Walter, reliable as a swiss watch!", "", null); + var po = new PaginateOperator(plan, 2); + assertAll( + () -> assertEquals("(Paginate,1,2,Great plan, Walter, reliable as a swiss watch!)", + po.toCursor()), + () -> assertNull(po.toCursor()), + () -> assertNull(po.toCursor()) + ); + } } 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 2c02c9b9cae..5e70f2b9d01 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 @@ -6,6 +6,10 @@ package org.opensearch.sql.planner.physical; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -74,4 +78,19 @@ void get_total_hits_uses_default_value() { when(plan.getTotalHits()).then(CALLS_REAL_METHODS); assertEquals(0, plan.getTotalHits()); } + + @Test + void toCursor() { + var plan = mock(PhysicalPlan.class); + when(plan.toCursor()).then(CALLS_REAL_METHODS); + assertTrue(assertThrows(IllegalStateException.class, plan::toCursor) + .getMessage().contains("is not compatible with cursor feature")); + } + + @Test + void createSection() { + var plan = mock(PhysicalPlan.class); + when(plan.createSection(anyString(), any())).then(CALLS_REAL_METHODS); + assertEquals("(plan,one,two)", plan.createSection("plan", "one", "two")); + } } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index 4ab8253c040..6042eba6dcc 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -10,8 +10,10 @@ import static org.hamcrest.Matchers.allOf; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.hasItems; -import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.iterableWithSize; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.mockito.Mockito.when; import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_MISSING; import static org.opensearch.sql.data.model.ExprValueUtils.stringValue; @@ -21,13 +23,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.ObjectInputStream; -import java.io.ObjectOutputStream; import java.util.List; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -37,6 +33,7 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.DSL; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; @ExtendWith(MockitoExtension.class) class ProjectOperatorTest extends PhysicalPlanTestBase { @@ -213,4 +210,20 @@ public void project_parse_missing_will_fallback() { ExprValueUtils.tupleValue(ImmutableMap.of("action", "GET", "response", "200")), ExprValueUtils.tupleValue(ImmutableMap.of("action", "POST"))))); } + + @Test + public void toCursor() { + when(inputPlan.toCursor()).thenReturn("inputPlan", "", null); + var project = DSL.named("response", DSL.ref("response", INTEGER)); + var npe = DSL.named("action", DSL.ref("action", STRING)); + var po = project(inputPlan, List.of(project), List.of(npe)); + var serializer = new DefaultExpressionSerializer(); + var expected = String.format("(Project,(namedParseExpressions,%s),(projectList,%s),%s)", + serializer.serialize(npe), serializer.serialize(project), "inputPlan"); + assertAll( + () -> assertEquals(expected, po.toCursor()), + () -> assertNull(po.toCursor()), + () -> assertNull(po.toCursor()) + ); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java index b111047b6fd..7b1353f4a97 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java @@ -113,4 +113,10 @@ void getTotalHitsSuccess() { monitorPlan.getTotalHits(); verify(plan, times(1)).getTotalHits(); } + + @Test + void toCursorSuccess() { + monitorPlan.toCursor(); + verify(plan, times(1)).toCursor(); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java new file mode 100644 index 00000000000..d0a274ce2a2 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java @@ -0,0 +1,23 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + + +package org.opensearch.sql.opensearch.request; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.withSettings; + +import org.junit.jupiter.api.Test; + +public class OpenSearchRequestTest { + + @Test + void toCursor() { + var request = mock(OpenSearchRequest.class, withSettings().defaultAnswer(CALLS_REAL_METHODS)); + assertEquals("", request.toCursor()); + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 8bb5c1eebe5..f2a95b777d9 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -61,4 +61,13 @@ void scrollRequest() { .scrollId("scroll123"), request.scrollRequest()); } + + @Test + void toCursor() { + request.setScrollId("scroll123"); + assertEquals("scroll123", request.toCursor()); + + request.reset(); + assertEquals("", request.toCursor()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java index fc77ad4ff7c..67839ffedb9 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -15,6 +15,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import static org.opensearch.sql.data.type.ExprCoreType.STRING; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.employee; @@ -35,6 +36,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.PagedRequestBuilder; import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -143,4 +145,23 @@ void explain_not_implemented() { assertThrows(Throwable.class, () -> mock(OpenSearchPagedIndexScan.class, withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); } + + @Test + void toCursor() { + PagedRequestBuilder builder = mock(); + OpenSearchRequest request = mock(); + OpenSearchResponse response = mock(); + when(builder.build()).thenReturn(request); + when(builder.getIndexName()).thenReturn(new OpenSearchRequest.IndexName("index")); + when(client.search(request)).thenReturn(response); + when(response.isEmpty()).thenReturn(true); + when(request.toCursor()).thenReturn("cu-cursor", "", null); + OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder); + indexScan.open(); + assertAll( + () -> assertEquals("(OpenSearchPagedIndexScan,index,cu-cursor)", indexScan.toCursor()), + () -> assertEquals("", indexScan.toCursor()), + () -> assertEquals("", indexScan.toCursor()) + ); + } } From 960c0392509656572f4c223a6adbbe9be7409521 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Mon, 13 Mar 2023 16:15:56 -0700 Subject: [PATCH 30/46] Fix broken tests in `legacy`. Signed-off-by: Yury-Fridlyand --- .../sql/legacy/plugin/RestSQLQueryActionTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionTest.java index 1bc34edf50a..be572f3dfbe 100644 --- a/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionTest.java +++ b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/RestSQLQueryActionTest.java @@ -74,7 +74,7 @@ public void handleQueryThatCanSupport() throws Exception { new JSONObject("{\"query\": \"SELECT -123\"}"), "SELECT -123", QUERY_API_ENDPOINT, - ""); + "jdbc"); RestSQLQueryAction queryAction = new RestSQLQueryAction(injector); queryAction.prepareRequest(request, (channel, exception) -> { @@ -90,7 +90,7 @@ public void handleExplainThatCanSupport() throws Exception { new JSONObject("{\"query\": \"SELECT -123\"}"), "SELECT -123", EXPLAIN_API_ENDPOINT, - ""); + "jdbc"); RestSQLQueryAction queryAction = new RestSQLQueryAction(injector); queryAction.prepareRequest(request, (channel, exception) -> { @@ -107,7 +107,7 @@ public void queryThatNotSupportIsHandledByFallbackHandler() throws Exception { "{\"query\": \"SELECT name FROM test1 JOIN test2 ON test1.name = test2.name\"}"), "SELECT name FROM test1 JOIN test2 ON test1.name = test2.name", QUERY_API_ENDPOINT, - ""); + "jdbc"); AtomicBoolean fallback = new AtomicBoolean(false); RestSQLQueryAction queryAction = new RestSQLQueryAction(injector); @@ -128,7 +128,7 @@ public void queryExecutionFailedIsHandledByExecutionErrorHandler() throws Except "{\"query\": \"SELECT -123\"}"), "SELECT -123", QUERY_API_ENDPOINT, - ""); + "jdbc"); doThrow(new IllegalStateException("execution exception")) .when(queryManager) From 4f0c176f20a0b27e84d6829b9b07241afc712b68 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 14 Mar 2023 15:34:43 -0700 Subject: [PATCH 31/46] Fix getting `total` from non-paged requests and from queries without `FROM` clause. Signed-off-by: Yury-Fridlyand --- .../sql/planner/physical/PhysicalPlan.java | 7 +++ .../sql/planner/physical/ValuesOperator.java | 10 +++- .../planner/physical/ValuesOperatorTest.java | 2 + .../request/OpenSearchRequestBuilder.java | 1 - .../storage/scan/OpenSearchIndexScan.java | 5 +- .../request/OpenSearchRequestBuilderTest.java | 23 ++++---- .../storage/scan/OpenSearchIndexScanTest.java | 54 +++++++++++++------ 7 files changed, 72 insertions(+), 30 deletions(-) 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 d24e3203dc7..312e4bfff9a 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 @@ -47,6 +47,13 @@ public ExecutionEngine.Schema schema() { + "ProjectOperator, instead of %s", this.getClass().getSimpleName())); } + /** + * Returns Total hits matched the search criteria. Note: query may return less if limited. + * {@see Settings#QUERY_SIZE_LIMIT}. + * Any plan which adds/removes rows to the response should overwrite it to provide valid values. + * + * @return Total hits matched the search criteria. + */ public long getTotalHits() { return getChild().stream().mapToLong(PhysicalPlan::getTotalHits).max().orElse(0); } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java index 51d2850df72..45884830e10 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java @@ -15,6 +15,7 @@ import lombok.ToString; import org.opensearch.sql.data.model.ExprCollectionValue; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.LiteralExpression; /** @@ -55,10 +56,17 @@ public boolean hasNext() { return valuesIterator.hasNext(); } + @Override + public long getTotalHits() { + // ValuesOperator used for queries without `FROM` clause, e.g. `select 1`. + // Such query always returns 1 row. + return 1; + } + @Override public ExprValue next() { List values = valuesIterator.next().stream() - .map(expr -> expr.valueOf()) + .map(Expression::valueOf) .collect(Collectors.toList()); return new ExprCollectionValue(values); } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ValuesOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ValuesOperatorTest.java index 9acab03d2bd..bf6d28a23c6 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ValuesOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ValuesOperatorTest.java @@ -9,6 +9,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; import static org.opensearch.sql.data.model.ExprValueUtils.collectionValue; import static org.opensearch.sql.expression.DSL.literal; @@ -44,6 +45,7 @@ public void iterateSingleRow() { results, contains(collectionValue(Arrays.asList(1, "abc"))) ); + assertThat(values.getTotalHits(), equalTo(1L)); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 7239ea7c0b9..4c253a31520 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -105,7 +105,6 @@ public OpenSearchRequestBuilder(OpenSearchRequest.IndexName indexName, /** * Build DSL request. * - * @return query request or scroll request */ public OpenSearchRequest build() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index f9a420332d1..83dfb30282f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -49,6 +49,7 @@ public class OpenSearchIndexScan extends TableScanOperator { /** Number of rows returned. */ private Integer queryCount; + private long totalHits = 0; /** Search response for current batch. */ private transient Iterator iterator; @@ -86,12 +87,12 @@ public ExprValue next() { @Override public long getTotalHits() { - // TODO maybe store totalHits from `response` - return queryCount; + return totalHits; } private void fetchNextBatch() { OpenSearchResponse response = client.search(request); + totalHits += response.getTotalHits(); if (!response.isEmpty()) { iterator = response.iterator(); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 33376ece836..e4f557d865b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -19,6 +19,8 @@ import java.util.Set; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -43,6 +45,7 @@ import org.opensearch.sql.opensearch.response.agg.SingleValueParser; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class OpenSearchRequestBuilderTest { private static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); @@ -67,7 +70,7 @@ void setup() { } @Test - void buildQueryRequest() { + void build_query_request() { Integer limit = 200; Integer offset = 0; requestBuilder.pushDownLimit(limit, offset); @@ -84,7 +87,7 @@ void buildQueryRequest() { } @Test - void buildScrollRequestWithCorrectSize() { + void build_scroll_request_with_correct_size() { Integer limit = 800; Integer offset = 10; requestBuilder.pushDownLimit(limit, offset); @@ -101,7 +104,7 @@ void buildScrollRequestWithCorrectSize() { } @Test - void testPushDownQuery() { + void test_push_down_query() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); requestBuilder.pushDown(query); @@ -117,7 +120,7 @@ void testPushDownQuery() { } @Test - void testPushDownAggregation() { + void test_push_down_aggregation() { AggregationBuilder aggBuilder = AggregationBuilders.composite( "composite_buckets", Collections.singletonList(new TermsValuesSourceBuilder("longA"))); @@ -138,7 +141,7 @@ void testPushDownAggregation() { } @Test - void testPushDownQueryAndSort() { + void test_push_down_query_and_sort() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); requestBuilder.pushDown(query); @@ -156,7 +159,7 @@ void testPushDownQueryAndSort() { } @Test - void testPushDownSort() { + void test_push_down_sort() { FieldSortBuilder sortBuilder = SortBuilders.fieldSort("intA"); requestBuilder.pushDownSort(List.of(sortBuilder)); @@ -170,7 +173,7 @@ void testPushDownSort() { } @Test - void testPushDownNonFieldSort() { + void test_push_down_non_field_sort() { ScoreSortBuilder sortBuilder = SortBuilders.scoreSort(); requestBuilder.pushDownSort(List.of(sortBuilder)); @@ -184,7 +187,7 @@ void testPushDownNonFieldSort() { } @Test - void testPushDownMultipleSort() { + void test_push_down_multiple_sort() { requestBuilder.pushDownSort(List.of( SortBuilders.fieldSort("intA"), SortBuilders.fieldSort("intB"))); @@ -200,7 +203,7 @@ void testPushDownMultipleSort() { } @Test - void testPushDownProject() { + void test_push_down_project() { Set references = Set.of(DSL.ref("intA", INTEGER)); requestBuilder.pushDownProjects(references); @@ -214,7 +217,7 @@ void testPushDownProject() { } @Test - void testPushTypeMapping() { + void test_push_type_mapping() { Map typeMapping = Map.of("intA", INTEGER); requestBuilder.pushTypeMapping(typeMapping); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index d93f4729b8b..2875c9b7803 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -71,9 +71,8 @@ void setup() { @Test void query_empty_result() { mockResponse(client); - try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("test", 3, settings, - exprValueFactory))) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("test", 3, settings, exprValueFactory))) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -90,11 +89,8 @@ void query_all_results_with_query() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - OpenSearchRequestBuilder - builder = new OpenSearchRequestBuilder("employees", 10, settings, - exprValueFactory); - try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, builder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("employees", 10, settings, exprValueFactory))) { indexScan.open(); assertAll( @@ -119,10 +115,10 @@ void query_all_results_with_scroll() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); + //when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(2); - try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("employees", 2, settings, - exprValueFactory))) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("employees", 10, settings, exprValueFactory))) { indexScan.open(); assertAll( @@ -150,9 +146,8 @@ void query_some_results_with_query() { employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - try (OpenSearchIndexScan indexScan = - new OpenSearchIndexScan(client, new OpenSearchRequestBuilder("employees", 10, settings, - exprValueFactory))) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("employees", 10, settings, exprValueFactory))) { indexScan.getRequestBuilder().pushDownLimit(3, 0); indexScan.open(); @@ -167,7 +162,7 @@ void query_some_results_with_query() { () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) + () -> assertEquals(4, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); @@ -196,7 +191,34 @@ void query_some_results_with_scroll() { () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) + () -> assertEquals(4, indexScan.getTotalHits()) + ); + } + verify(client).cleanup(any()); + } + + @Test + void query_results_limited_by_query_size() { + mockResponse(client, new ExprValue[]{ + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT"), + employee(4, "Bob", "HR")}); + when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(2); + + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + new OpenSearchRequestBuilder("employees", 10, settings, exprValueFactory))) { + indexScan.open(); + + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), + + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(4, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); From bdd52a029aa3952a1c57c72213a4e913f1b01c49 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 13:55:44 -0700 Subject: [PATCH 32/46] Fix scroll cleaning. Signed-off-by: Yury-Fridlyand --- .../request/ContinueScrollRequest.java | 4 +- .../request/OpenSearchScrollRequest.java | 23 ++++--- .../response/OpenSearchResponse.java | 6 +- .../scan/OpenSearchPagedIndexScan.java | 7 +- .../client/OpenSearchNodeClientTest.java | 16 +++-- .../client/OpenSearchRestClientTest.java | 15 +++-- .../request/ContinueScrollRequestTest.java | 7 ++ .../request/OpenSearchScrollRequestTest.java | 67 +++++++++++++++++-- .../scan/OpenSearchPagedIndexScanTest.java | 12 ++-- 9 files changed, 116 insertions(+), 41 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java index cf51680b7e2..5bfe35b015b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java @@ -55,7 +55,9 @@ public OpenSearchResponse search(Function searchA @Override public void clean(Consumer cleanAction) { - cleanAction.accept(responseScrollId); + if (scrollFinished) { + cleanAction.accept(responseScrollId); + } } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 555d520d7f6..c1bc17c8009 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -53,6 +53,8 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { @Getter private String scrollId; + private boolean needClean = false; + /** Search request source builder. */ private final SearchSourceBuilder sourceBuilder; @@ -81,21 +83,24 @@ public OpenSearchScrollRequest(IndexName indexName, public OpenSearchResponse search(Function searchAction, Function scrollAction) { SearchResponse openSearchResponse; - if (isScrollStarted()) { + if (isScroll()) { openSearchResponse = scrollAction.apply(scrollRequest()); } else { openSearchResponse = searchAction.apply(searchRequest()); } var response = new OpenSearchResponse(openSearchResponse, exprValueFactory); - setScrollId(openSearchResponse.getScrollId()); + if (!(needClean = response.isEmpty())) { + setScrollId(openSearchResponse.getScrollId()); + } return response; } @Override public void clean(Consumer cleanAction) { try { - if (isScrollStarted()) { + // clean on the last page only, to prevent closing the scroll/cursor in the middle of paging. + if (needClean && isScroll()) { cleanAction.accept(getScrollId()); setScrollId(null); } @@ -121,8 +126,8 @@ public SearchRequest searchRequest() { * * @return true if scroll started */ - public boolean isScrollStarted() { - return (scrollId != null); + public boolean isScroll() { + return scrollId != null; } /** @@ -149,11 +154,7 @@ public void reset() { */ @Override public String toCursor() { - if (isScrollStarted()) { - // TODO: probably should serialize exprValueFactory here as well. - return scrollId; - } else { - return ""; - } + // TODO: probably should serialize exprValueFactory here as well. + return scrollId; } } 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 74aa07fccb6..61d4459a862 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 @@ -39,13 +39,13 @@ public class OpenSearchResponse implements Iterable { private final Aggregations aggregations; /** - * ElasticsearchExprValueFactory used to build ExprValue from search result. + * OpenSearchExprValueFactory used to build ExprValue from search result. */ @EqualsAndHashCode.Exclude private final OpenSearchExprValueFactory exprValueFactory; /** - * Constructor of ElasticsearchResponse. + * Constructor of OpenSearchResponse. */ public OpenSearchResponse(SearchResponse searchResponse, OpenSearchExprValueFactory exprValueFactory) { @@ -55,7 +55,7 @@ public OpenSearchResponse(SearchResponse searchResponse, } /** - * Constructor of ElasticsearchResponse with SearchHits. + * Constructor of OpenSearchResponse with SearchHits. */ public OpenSearchResponse(SearchHits hits, OpenSearchExprValueFactory exprValueFactory) { this.hits = hits; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index 6626af6e9eb..e9d3fd52d39 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -26,7 +26,6 @@ public class OpenSearchPagedIndexScan extends TableScanOperator { @ToString.Include private OpenSearchRequest request; private Iterator iterator; - private boolean needClean = false; private long totalHits = 0; public OpenSearchPagedIndexScan(OpenSearchClient client, @@ -59,7 +58,6 @@ public void open() { iterator = response.iterator(); totalHits = response.getTotalHits(); } else { - needClean = true; iterator = Collections.emptyIterator(); } } @@ -67,10 +65,7 @@ public void open() { @Override public void close() { super.close(); - if (needClean) { - // clean on the last page only, to prevent closing the scroll/cursor in the middle of paging. - client.cleanup(request); - } + client.cleanup(request); } @Override diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index 206e1748f3a..5e6039eeafd 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -15,7 +15,6 @@ import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; -import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -32,6 +31,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import lombok.SneakyThrows; +import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -274,7 +275,6 @@ void search() { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(nodeClient.searchScroll(any()).actionGet()).thenReturn(scrollResponse); - when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -304,16 +304,19 @@ void schedule() { } @Test + @SneakyThrows void cleanup() { ClearScrollRequestBuilder requestBuilder = mock(ClearScrollRequestBuilder.class); when(nodeClient.prepareClearScroll()).thenReturn(requestBuilder); - lenient().when(requestBuilder.addScrollId(any())).thenReturn(requestBuilder); - lenient().when(requestBuilder.get()).thenReturn(null); + when(requestBuilder.addScrollId(any())).thenReturn(requestBuilder); + when(requestBuilder.get()).thenReturn(null); OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); request.setScrollId("scroll123"); + // Enforce cleaning by setting a private field. + FieldUtils.writeField(request, "needClean", true, true); client.cleanup(request); - assertFalse(request.isScrollStarted()); + assertFalse(request.isScroll()); InOrder inOrder = Mockito.inOrder(nodeClient, requestBuilder); inOrder.verify(nodeClient).prepareClearScroll(); @@ -329,11 +332,14 @@ void cleanup_without_scrollId() { } @Test + @SneakyThrows void cleanup_rethrows_exception() { when(nodeClient.prepareClearScroll()).thenThrow(new RuntimeException()); OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); request.setScrollId("scroll123"); + // Enforce cleaning by setting a private field. + FieldUtils.writeField(request, "needClean", true, true); assertThrows(IllegalStateException.class, () -> client.cleanup(request)); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index 9fcf7a00796..ca1996a6ac9 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -27,6 +27,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import lombok.SneakyThrows; +import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -254,7 +256,6 @@ void search() throws IOException { // Mock second scroll request followed SearchResponse scrollResponse = mock(SearchResponse.class); when(restClient.scroll(any(), any())).thenReturn(scrollResponse); - when(scrollResponse.getScrollId()).thenReturn("scroll456"); when(scrollResponse.getHits()).thenReturn(SearchHits.empty()); // Verify response for first scroll request @@ -315,12 +316,15 @@ void schedule() { } @Test - void cleanup() throws IOException { + @SneakyThrows + void cleanup() { OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); + // Enforce cleaning by setting a private field. + FieldUtils.writeField(request, "needClean", true, true); request.setScrollId("scroll123"); client.cleanup(request); verify(restClient).clearScroll(any(), any()); - assertFalse(request.isScrollStarted()); + assertFalse(request.isScroll()); } @Test @@ -331,10 +335,13 @@ void cleanup_without_scrollId() throws IOException { } @Test - void cleanup_with_IOException() throws IOException { + @SneakyThrows + void cleanup_with_IOException() { when(restClient.clearScroll(any(), any())).thenThrow(new IOException()); OpenSearchScrollRequest request = new OpenSearchScrollRequest("test", factory); + // Enforce cleaning by setting a private field. + FieldUtils.writeField(request, "needClean", true, true); request.setScrollId("scroll123"); assertThrows(IllegalStateException.class, () -> client.cleanup(request)); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java index f25553a55f1..048c3e36de7 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinueScrollRequestTest.java @@ -22,6 +22,8 @@ import java.util.function.Consumer; import java.util.function.Function; +import lombok.SneakyThrows; +import org.apache.commons.lang3.reflect.FieldUtils; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; @@ -100,7 +102,12 @@ public void search_with_empty_response() { } @Test + @SneakyThrows public void clean() { + request.clean(cleanAction); + verify(cleanAction, never()).accept(any()); + // Enforce cleaning by setting a private field. + FieldUtils.writeField(request, "scrollFinished", true, true); request.clean(cleanAction); verify(cleanAction, times(1)).accept(any()); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index f2a95b777d9..6e45476306e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -8,19 +8,31 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.lucene.search.TotalHits; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchScrollRequestTest { @Mock @@ -43,13 +55,13 @@ void searchRequest() { @Test void isScrollStarted() { - assertFalse(request.isScrollStarted()); + assertFalse(request.isScroll()); request.setScrollId("scroll123"); - assertTrue(request.isScrollStarted()); + assertTrue(request.isScroll()); request.reset(); - assertFalse(request.isScrollStarted()); + assertFalse(request.isScroll()); } @Test @@ -68,6 +80,53 @@ void toCursor() { assertEquals("scroll123", request.toCursor()); request.reset(); - assertEquals("", request.toCursor()); + assertNull(request.toCursor()); + } + + @Test + void clean_on_empty_response() { + // This could happen on sequential search calls + SearchResponse searchResponse = mock(); + when(searchResponse.getScrollId()).thenReturn("scroll1", "scroll2"); + when(searchResponse.getHits()).thenReturn( + new SearchHits(new SearchHit[1], new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1F), + new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1F)); + + request.search((x) -> searchResponse, (x) -> searchResponse); + assertEquals("scroll1", request.getScrollId()); + request.search((x) -> searchResponse, (x) -> searchResponse); + assertEquals("scroll1", request.getScrollId()); + + AtomicBoolean cleanCalled = new AtomicBoolean(false); + request.clean((s) -> cleanCalled.set(true)); + + assertNull(request.getScrollId()); + assertTrue(cleanCalled.get()); + } + + @Test + void no_clean_on_non_empty_response() { + SearchResponse searchResponse = mock(); + when(searchResponse.getScrollId()).thenReturn("scroll"); + when(searchResponse.getHits()).thenReturn( + new SearchHits(new SearchHit[1], new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1F)); + + request.search((x) -> searchResponse, (x) -> searchResponse); + assertEquals("scroll", request.getScrollId()); + + request.clean((s) -> fail()); + assertNull(request.getScrollId()); + } + + @Test + void no_clean_if_no_scroll_in_response() { + SearchResponse searchResponse = mock(); + when(searchResponse.getHits()).thenReturn( + new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1F)); + + request.search((x) -> searchResponse, (x) -> searchResponse); + assertNull(request.getScrollId()); + + request.clean((s) -> fail()); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java index 67839ffedb9..eab71e3f4f1 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -13,7 +13,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; @@ -88,8 +88,7 @@ void query_all_results_initial_scroll_request() { () -> assertEquals(3, indexScan.getTotalHits()) ); } - // cleanup should be called on empty response only - verify(client, never()).cleanup(any()); + verify(client).cleanup(any()); builder = new SubsequentPageRequestBuilder( new OpenSearchRequest.IndexName("test"), "scroll", exprValueFactory); @@ -98,7 +97,7 @@ void query_all_results_initial_scroll_request() { assertFalse(indexScan.hasNext()); } - verify(client).cleanup(any()); + verify(client, times(2)).cleanup(any()); } @Test @@ -127,8 +126,7 @@ void query_all_results_continuation_scroll_request() { () -> assertEquals(3, indexScan.getTotalHits()) ); } - // cleanup should be called on empty response only - verify(client, never()).cleanup(any()); + verify(client).cleanup(any()); builder = new SubsequentPageRequestBuilder( new OpenSearchRequest.IndexName("test"), "scroll", exprValueFactory); @@ -137,7 +135,7 @@ void query_all_results_continuation_scroll_request() { assertFalse(indexScan.hasNext()); } - verify(client).cleanup(any()); + verify(client, times(2)).cleanup(any()); } @Test From a16332f7a9d4e8970fdca90096347fe05e76b297 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 13:56:36 -0700 Subject: [PATCH 33/46] Fix cursor request processing. Signed-off-by: Yury-Fridlyand --- .../opensearch/sql/sql/domain/SQLQueryRequest.java | 6 ++++-- .../sql/sql/domain/SQLQueryRequestTest.java | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) 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 0d15abf54a5..7545f4cc19f 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 @@ -89,10 +89,12 @@ public SQLQueryRequest(JSONObject jsonContent, String query, String path, public boolean isSupported() { var noCursor = !isCursor(); var noQuery = query == null; - var noParams = params.isEmpty(); + var noUnsupportedParams = params.isEmpty() + || (params.size() == 1 && params.containsKey(QUERY_PARAMS_FORMAT)); var noContent = jsonContent == null || jsonContent.isEmpty(); - return ((!noCursor && noQuery && noParams && noContent) // if cursor is given, but other things + return ((!noCursor && noQuery + && noUnsupportedParams && noContent) // if cursor is given, but other things || (noCursor && !noQuery)) // or if cursor is not given, but query && isOnlySupportedFieldInPayload() // and request has supported fields only && isSupportedFormat(); // and request is in supported format 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 3a4e4ae09ff..62bb665537c 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 @@ -160,6 +160,16 @@ public void should_not_support_request_with_cursor_and_something_else() { .cursor("n:12356") .params(Map.of("one", "two")) .build(); + SQLQueryRequest requestWithParamsWithFormat = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .params(Map.of("format", "jdbc")) + .build(); + SQLQueryRequest requestWithParamsWithFormatAnd = + SQLQueryRequestBuilder.request(null) + .cursor("n:12356") + .params(Map.of("format", "jdbc", "something", "else")) + .build(); SQLQueryRequest requestWithFetchSize = SQLQueryRequestBuilder.request(null) .cursor("n:12356") @@ -180,6 +190,8 @@ public void should_not_support_request_with_cursor_and_something_else() { () -> assertFalse(requestWithParams.isSupported()), () -> assertFalse(requestWithFetchSize.isSupported()), () -> assertTrue(requestWithNoParams.isSupported()), + () -> assertTrue(requestWithParamsWithFormat.isSupported()), + () -> assertFalse(requestWithParamsWithFormatAnd.isSupported()), () -> assertTrue(requestWithNoContent.isSupported()) ); } From 9f9e8737af91e8c74ae27ed89c45b5d7695d1451 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 13:57:35 -0700 Subject: [PATCH 34/46] Update ITs. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/sql/HighlightFunctionIT.java | 2 +- .../opensearch/sql/sql/StandalonePaginationIT.java | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.java index 809e2dc7c51..0ab6d5c70fb 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.java @@ -64,7 +64,7 @@ public void highlight_multiple_optional_arguments_test() { schema("highlight(Body, pre_tags='', " + "post_tags='')", null, "nested")); - assertEquals(1, response.getInt("total")); + assertEquals(1, response.getInt("size")); verifyDataRows(response, rows(new JSONArray(List.of("What are the differences between an IPA" + " and its variants?")), diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 46613539a4a..913433eb08c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -7,6 +7,7 @@ import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata; import static org.opensearch.sql.legacy.TestUtils.getResponseBody; +import static org.opensearch.sql.legacy.TestUtils.isIndexExist; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -66,7 +67,9 @@ public void init() { loadIndex(Index.ONLINE); loadIndex(Index.BEER); loadIndex(Index.BANK); - executeRequest(new Request("PUT", "/empty")); + if (!isIndexExist(client(), "empty")) { + executeRequest(new Request("PUT", "/empty")); + } RestHighLevelClient restClient = new InternalRestHighLevelClient(client()); client = new OpenSearchRestClient(restClient); @@ -135,6 +138,7 @@ public void onFailure(Exception e) { } // Test takes 3+ min due to a big amount of requests issued + // Skip 'online' index and/or page_size = 1 to get a significant speed-up @Test @SneakyThrows public void test_pagination_blackbox() { @@ -154,13 +158,14 @@ public void test_pagination_blackbox() { var cursor = response.getString("cursor"); assertTrue(testReportPrefix + "Cursor returned from legacy engine", cursor.startsWith("n:")); - rowsReturned += response.getInt("total"); + rowsReturned += response.getInt("size"); var datarows = response.getJSONArray("datarows"); for (int i = 0; i < datarows.length(); i++) { rowsPaged.put(datarows.get(i)); } assertTrue("Paged response schema doesn't match to non-paged", schema.similar(response.getJSONArray("schema"))); + assertEquals(indexSize, response.getInt("total")); response = executeCursorQuery(cursor); } assertEquals(testReportPrefix + "Last page is not empty", @@ -190,7 +195,7 @@ public void test_explain_not_supported() { request.setJsonEntity("{ \"cursor\" : \"n:0000\" }"); exception = assertThrows(ResponseException.class, () -> client().performRequest(request)); response = new JSONObject(new String(exception.getResponse().getEntity().getContent().readAllBytes())); - assertEquals("`explain` request for cursor requests is not supported. Use `explain` for the initial query request.", + assertEquals("Explain of a paged query continuation is not supported. Use `explain` for the initial query request.", response.getJSONObject("error").getString("details")); } From 3340e388c4fbe22f9af77f477eab50470b42bcf4 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 21:40:18 -0700 Subject: [PATCH 35/46] Fix (again) TotalHits feature. Signed-off-by: Yury-Fridlyand --- .../sql/planner/physical/FilterOperator.java | 16 ++++++++-- .../planner/physical/FilterOperatorTest.java | 32 +++++++++++++++++-- .../storage/scan/OpenSearchIndexScan.java | 6 ++-- .../system/OpenSearchSystemIndexScan.java | 11 ++++++- .../storage/scan/OpenSearchIndexScanTest.java | 6 ++-- .../system/OpenSearchSystemIndexScanTest.java | 1 + 6 files changed, 58 insertions(+), 14 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java index 86cd411a2da..a9c7597c3e5 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java @@ -17,8 +17,9 @@ import org.opensearch.sql.storage.bindingtuple.BindingTuple; /** - * The Filter operator use the conditions to evaluate the input {@link BindingTuple}. - * The Filter operator only return the results that evaluated to true. + * The Filter operator represents WHERE clause and + * uses the conditions to evaluate the input {@link BindingTuple}. + * The Filter operator only returns the results that evaluated to true. * The NULL and MISSING are handled by the logic defined in {@link BinaryPredicateOperator}. */ @EqualsAndHashCode(callSuper = false) @@ -29,7 +30,9 @@ public class FilterOperator extends PhysicalPlan { private final PhysicalPlan input; @Getter private final Expression conditions; - @ToString.Exclude private ExprValue next = null; + @ToString.Exclude + private ExprValue next = null; + private long totalHits = 0; @Override public R accept(PhysicalPlanNodeVisitor visitor, C context) { @@ -48,6 +51,7 @@ public boolean hasNext() { ExprValue exprValue = conditions.valueOf(inputValue.bindingTuples()); if (!(exprValue.isNull() || exprValue.isMissing()) && (exprValue.booleanValue())) { next = inputValue; + totalHits++; return true; } } @@ -58,4 +62,10 @@ public boolean hasNext() { public ExprValue next() { return next; } + + @Override + public long getTotalHits() { + // ignore `input.getTotalHits()`, because it returns wrong (unfiltered) value + return totalHits; + } } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/FilterOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/FilterOperatorTest.java index 288b4bf661f..f541f6a15f1 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/FilterOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/FilterOperatorTest.java @@ -17,22 +17,30 @@ import com.google.common.collect.ImmutableMap; import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.AdditionalAnswers; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.data.model.ExprIntegerValue; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.expression.DSL; @ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class FilterOperatorTest extends PhysicalPlanTestBase { @Mock private PhysicalPlan inputPlan; @Test - public void filterTest() { + public void filter_test() { FilterOperator plan = new FilterOperator(new TestScan(), DSL.equal(DSL.ref("response", INTEGER), DSL.literal(404))); List result = execute(plan); @@ -41,10 +49,11 @@ public void filterTest() { .tupleValue(ImmutableMap .of("ip", "209.160.24.63", "action", "GET", "response", 404, "referer", "www.amazon.com")))); + assertEquals(1, plan.getTotalHits()); } @Test - public void nullValueShouldBeenIgnored() { + public void null_value_should_been_ignored() { LinkedHashMap value = new LinkedHashMap<>(); value.put("response", LITERAL_NULL); when(inputPlan.hasNext()).thenReturn(true, false); @@ -54,10 +63,11 @@ public void nullValueShouldBeenIgnored() { DSL.equal(DSL.ref("response", INTEGER), DSL.literal(404))); List result = execute(plan); assertEquals(0, result.size()); + assertEquals(0, plan.getTotalHits()); } @Test - public void missingValueShouldBeenIgnored() { + public void missing_value_should_been_ignored() { LinkedHashMap value = new LinkedHashMap<>(); value.put("response", LITERAL_MISSING); when(inputPlan.hasNext()).thenReturn(true, false); @@ -67,5 +77,21 @@ public void missingValueShouldBeenIgnored() { DSL.equal(DSL.ref("response", INTEGER), DSL.literal(404))); List result = execute(plan); assertEquals(0, result.size()); + assertEquals(0, plan.getTotalHits()); + } + + @Test + public void totalHits() { + when(inputPlan.hasNext()).thenReturn(true, true, true, true, true, false); + var answers = Stream.of(200, 240, 300, 403, 404).map(c -> + new ExprTupleValue(new LinkedHashMap<>(Map.of("response", new ExprIntegerValue(c))))) + .collect(Collectors.toList()); + when(inputPlan.next()).thenAnswer(AdditionalAnswers.returnsElementsOf(answers)); + + FilterOperator plan = new FilterOperator(inputPlan, + DSL.less(DSL.ref("response", INTEGER), DSL.literal(400))); + List result = execute(plan); + assertEquals(3, result.size()); + assertEquals(3, plan.getTotalHits()); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 83dfb30282f..bc846c38d7c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -49,8 +49,6 @@ public class OpenSearchIndexScan extends TableScanOperator { /** Number of rows returned. */ private Integer queryCount; - private long totalHits = 0; - /** Search response for current batch. */ private transient Iterator iterator; @@ -87,12 +85,12 @@ public ExprValue next() { @Override public long getTotalHits() { - return totalHits; + // ignore response.getTotalHits(), because response returns entire index, regardless of LIMIT + return queryCount; } private void fetchNextBatch() { OpenSearchResponse response = client.search(request); - totalHits += response.getTotalHits(); if (!response.isEmpty()) { iterator = response.iterator(); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScan.java index eb4cb865e22..eba5eb126de 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScan.java @@ -31,9 +31,13 @@ public class OpenSearchSystemIndexScan extends TableScanOperator { */ private Iterator iterator; + private long totalHits = 0; + @Override public void open() { - iterator = request.search().iterator(); + var response = request.search(); + totalHits = response.size(); + iterator = response.iterator(); } @Override @@ -46,6 +50,11 @@ public ExprValue next() { return iterator.next(); } + @Override + public long getTotalHits() { + return totalHits; + } + @Override public String explain() { return request.toString(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 2875c9b7803..5b90451cf8f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -162,7 +162,7 @@ void query_some_results_with_query() { () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(4, indexScan.getTotalHits()) + () -> assertEquals(3, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); @@ -191,7 +191,7 @@ void query_some_results_with_scroll() { () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(4, indexScan.getTotalHits()) + () -> assertEquals(3, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); @@ -218,7 +218,7 @@ void query_results_limited_by_query_size() { () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(4, indexScan.getTotalHits()) + () -> assertEquals(2, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScanTest.java index 494f3ff2d0e..c04ef25611e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/system/OpenSearchSystemIndexScanTest.java @@ -32,6 +32,7 @@ public void queryData() { systemIndexScan.open(); assertTrue(systemIndexScan.hasNext()); assertEquals(stringValue("text"), systemIndexScan.next()); + assertEquals(1, systemIndexScan.getTotalHits()); } @Test From 524f2201017ea395c2bf072d0f76f9c81bcadd96 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 21:40:59 -0700 Subject: [PATCH 36/46] Fix typo in prometheus config. Signed-off-by: Yury-Fridlyand --- integ-test/src/test/resources/datasource/datasources.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integ-test/src/test/resources/datasource/datasources.json b/integ-test/src/test/resources/datasource/datasources.json index fbf18cfc462..487b4d16d44 100644 --- a/integ-test/src/test/resources/datasource/datasources.json +++ b/integ-test/src/test/resources/datasource/datasources.json @@ -3,7 +3,7 @@ "name" : "my_prometheus", "connector": "prometheus", "properties" : { - "prometheus.uri" : "http://localhost:9091" + "prometheus.uri" : "http://localhost:9090" } } ] From 281f3cd121317eec615a12ca9261a743eb4aef15 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 15 Mar 2023 21:41:50 -0700 Subject: [PATCH 37/46] Recover commented logging. Signed-off-by: Yury-Fridlyand --- .../java/org/opensearch/sql/legacy/plugin/RestSqlAction.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 67f5f3bb43e..e1c72f0f1ef 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -142,9 +142,8 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli } } - -// LOG.info("[{}] Incoming request {}: {}", QueryContext.getRequestId(), request.uri(), -// QueryDataAnonymizer.anonymizeData(sqlRequest.getSql())); + LOG.info("[{}] Incoming request {}: {}", QueryContext.getRequestId(), request.uri(), + QueryDataAnonymizer.anonymizeData(sqlRequest.getSql())); Format format = SqlRequestParam.getFormat(request.params()); From ca76e1b9248fc87b752bd4a4d12cd0fea566ac7e Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 16 Mar 2023 09:50:42 -0700 Subject: [PATCH 38/46] Move `test_pagination_blackbox` to a separate class and add logging. Signed-off-by: Yury-Fridlyand --- .../sql/sql/PaginationBlackboxIT.java | 114 ++++++++++++++++++ .../sql/sql/StandalonePaginationIT.java | 51 -------- 2 files changed, 114 insertions(+), 51 deletions(-) create mode 100644 integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java new file mode 100644 index 00000000000..6e19c4d64ac --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java @@ -0,0 +1,114 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + + +package org.opensearch.sql.sql; + +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; +import static org.opensearch.sql.legacy.TestUtils.isIndexExist; + +import java.util.ArrayList; +import java.util.List; +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import lombok.SneakyThrows; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.Test; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.opensearch.client.Request; +import org.opensearch.sql.legacy.SQLIntegTestCase; + +// This class has only one test case, because it is parametrized and takes significant time +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class PaginationBlackboxIT extends SQLIntegTestCase { + + private final String index; + private final Integer pageSize; + + public PaginationBlackboxIT(@Name("index") String index, + @Name("pageSize") Integer pageSize) { + this.index = index; + this.pageSize = pageSize; + } + + @ParametersFactory(argumentFormatting = "index = %1$s, page_size = %2$d") + public static Iterable compareTwoDates() { + var indices = new PaginationBlackboxHelper().getIndices(); + var pageSizes = List.of(1, 5, 10, 100, 1000); + var testData = new ArrayList(); + for (var index : indices) { + for (var pageSize : pageSizes) { + testData.add(new Object[] { index, pageSize }); + } + } + return testData; + } + + // Test takes 3+ min due to a big amount of requests issued + // Skip 'online' index and/or page_size = 1 to get a significant speed-up + @Test + @SneakyThrows + public void test_pagination_blackbox() { + var response = executeJdbcRequest(String.format("select * from %s", index)); + var indexSize = response.getInt("total"); + var rows = response.getJSONArray("datarows"); + var schema = response.getJSONArray("schema"); + var testReportPrefix = String.format("index: %s, page size: %d || ", index, pageSize); + var rowsPaged = new JSONArray(); + var rowsReturned = 0; + response = new JSONObject(executeFetchQuery( + String.format("select * from %s", index), pageSize, "jdbc")); + var responseCounter = 1; + this.logger.info(testReportPrefix + "first response"); + while (response.has("cursor")) { + assertEquals(indexSize, response.getInt("total")); + assertTrue("Paged response schema doesn't match to non-paged", + schema.similar(response.getJSONArray("schema"))); + var cursor = response.getString("cursor"); + assertTrue(testReportPrefix + "Cursor returned from legacy engine", + cursor.startsWith("n:")); + rowsReturned += response.getInt("size"); + var datarows = response.getJSONArray("datarows"); + for (int i = 0; i < datarows.length(); i++) { + rowsPaged.put(datarows.get(i)); + } + response = executeCursorQuery(cursor); + this.logger.info(testReportPrefix + + String.format("subsequent response %d/%d", responseCounter++, (indexSize / pageSize) + 1)); + } + assertTrue("Paged response schema doesn't match to non-paged", + schema.similar(response.getJSONArray("schema"))); + assertEquals(0, response.getInt("total")); + + assertEquals(testReportPrefix + "Last page is not empty", + 0, response.getInt("size")); + assertEquals(testReportPrefix + "Last page is not empty", + 0, response.getJSONArray("datarows").length()); + assertEquals(testReportPrefix + "Paged responses return another row count that non-paged", + indexSize, rowsReturned); + assertTrue(testReportPrefix + "Paged accumulated result has other rows than non-paged", + rows.similar(rowsPaged)); + } + + // A dummy class created, because accessing to `client()` isn't available from a static context, + // but it is needed before an instance of `PaginationBlackboxIT` is created. + private static class PaginationBlackboxHelper extends SQLIntegTestCase { + + @SneakyThrows + private String[] getIndices() { + initClient(); + loadIndex(Index.ACCOUNT); + loadIndex(Index.ONLINE); + loadIndex(Index.BEER); + loadIndex(Index.BANK); + if (!isIndexExist(client(), "empty")) { + executeRequest(new Request("PUT", "/empty")); + } + return getResponseBody(client().performRequest(new Request("GET", "_cat/indices?h=i")), true).split("\n"); + } + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 913433eb08c..e00efce0132 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -63,14 +63,6 @@ public class StandalonePaginationIT extends SQLIntegTestCase { @Override @SneakyThrows public void init() { - loadIndex(Index.ACCOUNT); - loadIndex(Index.ONLINE); - loadIndex(Index.BEER); - loadIndex(Index.BANK); - if (!isIndexExist(client(), "empty")) { - executeRequest(new Request("PUT", "/empty")); - } - RestHighLevelClient restClient = new InternalRestHighLevelClient(client()); client = new OpenSearchRestClient(restClient); DataSourceService dataSourceService = new DataSourceServiceImpl( @@ -137,49 +129,6 @@ public void onFailure(Exception e) { // act 3: confirm that there's no cursor. } - // Test takes 3+ min due to a big amount of requests issued - // Skip 'online' index and/or page_size = 1 to get a significant speed-up - @Test - @SneakyThrows - public void test_pagination_blackbox() { - var indices = getResponseBody(client().performRequest(new Request("GET", "_cat/indices?h=i")), true).split("\n"); - for (var index : indices) { - var response = executeJdbcRequest(String.format("select * from %s", index)); - var indexSize = response.getInt("total"); - var rows = response.getJSONArray("datarows"); - var schema = response.getJSONArray("schema"); - for (var pageSize : List.of(1, 5, 10, 100, 1000)) { - var testReportPrefix = String.format("index: %s, page size: %d || ", index, pageSize); - var rowsPaged = new JSONArray(); - var rowsReturned = 0; - response = new JSONObject(executeFetchQuery( - String.format("select * from %s", index), pageSize, "jdbc")); - while (response.has("cursor")) { - var cursor = response.getString("cursor"); - assertTrue(testReportPrefix + "Cursor returned from legacy engine", - cursor.startsWith("n:")); - rowsReturned += response.getInt("size"); - var datarows = response.getJSONArray("datarows"); - for (int i = 0; i < datarows.length(); i++) { - rowsPaged.put(datarows.get(i)); - } - assertTrue("Paged response schema doesn't match to non-paged", - schema.similar(response.getJSONArray("schema"))); - assertEquals(indexSize, response.getInt("total")); - response = executeCursorQuery(cursor); - } - assertEquals(testReportPrefix + "Last page is not empty", - 0, response.getInt("total")); - assertEquals(testReportPrefix + "Last page is not empty", - 0, response.getJSONArray("datarows").length()); - assertEquals(testReportPrefix + "Paged responses return another row count that non-paged", - indexSize, rowsReturned); - assertTrue(testReportPrefix + "Paged accumulated result has other rows than non-paged", - rows.similar(rowsPaged)); - } - } - } - @Test @SneakyThrows public void test_explain_not_supported() { From c8fcd4e7eb1e10bbfe262ec9b7ad576b7b9a41e4 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 16 Mar 2023 11:39:05 -0700 Subject: [PATCH 39/46] Address some PR feedbacks: rename some classes and revert unnecessary whitespace changed. Signed-off-by: Yury-Fridlyand --- .../sql/exception/UnsupportedCursorRequestException.java | 9 +++++++++ .../org/opensearch/sql/executor/PaginatedPlanCache.java | 2 +- .../java/org/opensearch/sql/executor/QueryService.java | 1 - .../sql/executor/execution/QueryPlanFactory.java | 4 ++-- .../legacy/plugin/UnsupportCursorRequestException.java | 9 --------- .../org/opensearch/sql/planner/DefaultImplementor.java | 1 + .../sql/planner/{ => physical}/PaginateOperator.java | 2 +- .../sql/planner/physical/PhysicalPlanNodeVisitor.java | 1 - .../opensearch/sql/executor/PaginatedPlanCacheTest.java | 4 +--- .../sql/executor/execution/QueryPlanFactoryTest.java | 4 ++-- .../opensearch/sql/planner/DefaultImplementorTest.java | 1 + .../sql/planner/physical/PaginateOperatorTest.java | 1 - .../planner/physical/PhysicalPlanNodeVisitorTest.java | 1 - .../src/test/resources/datasource/datasources.json | 2 +- .../opensearch/sql/legacy/plugin/RestSQLQueryAction.java | 3 ++- .../executor/protector/OpenSearchExecutionProtector.java | 2 +- .../executor/OpenSearchExecutionEngineTest.java | 2 +- .../protector/OpenSearchExecutionProtectorTest.java | 2 +- 18 files changed, 24 insertions(+), 27 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/exception/UnsupportedCursorRequestException.java delete mode 100644 core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java rename core/src/main/java/org/opensearch/sql/planner/{ => physical}/PaginateOperator.java (98%) diff --git a/core/src/main/java/org/opensearch/sql/exception/UnsupportedCursorRequestException.java b/core/src/main/java/org/opensearch/sql/exception/UnsupportedCursorRequestException.java new file mode 100644 index 00000000000..0ae87f15440 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/exception/UnsupportedCursorRequestException.java @@ -0,0 +1,9 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.exception; + +public class UnsupportedCursorRequestException extends RuntimeException { +} diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 26294e3be15..7a876e695a3 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -19,7 +19,7 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.executor.Cursor; -import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.storage.StorageEngine; diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index ed251e2b33c..94e70819204 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -70,7 +70,6 @@ public void executePlan(LogicalPlan plan, } } - /** * Explain the query in {@link UnresolvedPlan} using {@link ResponseListener} to * get and format explain response. diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index c3a84ea286e..7834e4a03b6 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -18,11 +18,11 @@ 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.exception.UnsupportedCursorRequestException; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryService; -import org.opensearch.sql.legacy.plugin.UnsupportCursorRequestException; /** * QueryExecution Factory. @@ -102,7 +102,7 @@ public AbstractPlan visitQuery( context.getLeft().get()); } else { // This should be picked up by the legacy engine. - throw new UnsupportCursorRequestException(); + throw new UnsupportedCursorRequestException(); } } else { return new QueryPlan(QueryId.queryId(), node.getPlan(), queryService, diff --git a/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java b/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java deleted file mode 100644 index 7fcfceeebaf..00000000000 --- a/core/src/main/java/org/opensearch/sql/legacy/plugin/UnsupportCursorRequestException.java +++ /dev/null @@ -1,9 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.legacy.plugin; - -public class UnsupportCursorRequestException extends RuntimeException { -} diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index a02d908a0f1..43422d87336 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -27,6 +27,7 @@ import org.opensearch.sql.planner.physical.EvalOperator; import org.opensearch.sql.planner.physical.FilterOperator; import org.opensearch.sql.planner.physical.LimitOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.planner.physical.RareTopNOperator; diff --git a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java similarity index 98% rename from core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java rename to core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java index 68e94ac6b21..33856015f62 100644 --- a/core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.planner; +package org.opensearch.sql.planner.physical; import java.util.List; import lombok.EqualsAndHashCode; diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java index e52e5979ed3..f8b6f2243e0 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java @@ -6,7 +6,6 @@ package org.opensearch.sql.planner.physical; -import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.write.TableWriteOperator; diff --git a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java index fe798461ba6..e99612c5f3a 100644 --- a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java @@ -28,13 +28,11 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; -import org.mockito.Answers; -import org.mockito.Mock; import org.mockito.Mockito; import org.opensearch.sql.ast.dsl.AstDSL; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.opensearch.executor.Cursor; -import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.TableScanOperator; diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index 08718cb4d61..ff02d00b4ed 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -26,10 +26,10 @@ import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.exception.UnsupportedCursorRequestException; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.executor.QueryService; -import org.opensearch.sql.legacy.plugin.UnsupportCursorRequestException; @ExtendWith(MockitoExtension.class) class QueryPlanFactoryTest { @@ -139,7 +139,7 @@ public void createQueryWithFetchSizeWhichCannotBePaged() { when(paginatedPlanCache.canConvertToCursor(plan)).thenReturn(false); factory = new QueryPlanFactory(queryService, paginatedQueryService, paginatedPlanCache); Statement query = new Query(plan, 10); - assertThrows(UnsupportCursorRequestException.class, + assertThrows(UnsupportedCursorRequestException.class, () -> factory.create(query, Optional.of(queryListener), Optional.empty())); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java index a426bbb1a94..da3f5315e46 100644 --- a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java @@ -56,6 +56,7 @@ import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalRelation; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; import org.opensearch.sql.storage.Table; diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java index 01d525978f0..19df23b9e58 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java @@ -26,7 +26,6 @@ import org.junit.jupiter.api.DisplayNameGenerator; import org.junit.jupiter.api.Test; import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.planner.PaginateOperator; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class PaginateOperatorTest extends PhysicalPlanTestBase { diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java index 27d7a43bfdc..3dfe0b5c0fd 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitorTest.java @@ -26,7 +26,6 @@ import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.PaginateOperator; /** * Todo, testing purpose, delete later. diff --git a/integ-test/src/test/resources/datasource/datasources.json b/integ-test/src/test/resources/datasource/datasources.json index 487b4d16d44..5f195747ae0 100644 --- a/integ-test/src/test/resources/datasource/datasources.json +++ b/integ-test/src/test/resources/datasource/datasources.json @@ -6,4 +6,4 @@ "prometheus.uri" : "http://localhost:9090" } } -] +] \ No newline at end of file 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 4f9fdd9a535..cbbc8c7b9cb 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 @@ -24,6 +24,7 @@ import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.utils.QueryContext; +import org.opensearch.sql.exception.UnsupportedCursorRequestException; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.legacy.metrics.MetricName; import org.opensearch.sql.legacy.metrics.Metrics; @@ -126,7 +127,7 @@ public void onResponse(T response) { @Override public void onFailure(Exception e) { - if (e instanceof SyntaxCheckException || e instanceof UnsupportCursorRequestException) { + if (e instanceof SyntaxCheckException || e instanceof UnsupportedCursorRequestException) { fallBackHandler.accept(channel, e); } else { next.onFailure(e); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java index 572b5cbadec..4d6925f1aaf 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java @@ -11,12 +11,12 @@ import org.opensearch.sql.opensearch.planner.physical.ADOperator; import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator; import org.opensearch.sql.opensearch.planner.physical.MLOperator; -import org.opensearch.sql.planner.PaginateOperator; import org.opensearch.sql.planner.physical.AggregationOperator; import org.opensearch.sql.planner.physical.DedupeOperator; import org.opensearch.sql.planner.physical.EvalOperator; import org.opensearch.sql.planner.physical.FilterOperator; import org.opensearch.sql.planner.physical.LimitOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.planner.physical.RareTopNOperator; 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 913a774e1b7..c39609cce74 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 @@ -45,7 +45,7 @@ import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; -import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index 58da3f3f9aa..d0e486fae9c 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -60,7 +60,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; -import org.opensearch.sql.planner.PaginateOperator; +import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; From ec5fb40bd31864deab2f22a67c709e373eece54b Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 16 Mar 2023 13:01:56 -0700 Subject: [PATCH 40/46] Minor commenting. Signed-off-by: Yury-Fridlyand --- .../sql/executor/execution/PaginatedQueryService.java | 9 +++++++++ .../java/org/opensearch/sql/util/StandaloneModule.java | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java index 9e2ad73336b..eba8656c4c4 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/PaginatedQueryService.java @@ -13,10 +13,19 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.executor.ExecutionContext; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.planner.Planner; import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; import org.opensearch.sql.planner.physical.PhysicalPlan; +/** + * `PaginatedQueryService` does the same as `QueryService`, but it has another planner, + * configured to handle paged index scan. + * @see OpenSearchPluginModule#queryPlanFactory (:plugin module) + * @see LogicalPlanOptimizer#paginationCreate + * @see QueryService + */ @RequiredArgsConstructor public class PaginatedQueryService { private final Analyzer analyzer; diff --git a/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java index 3551c22a1d6..9a62788a02a 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java +++ b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java @@ -111,7 +111,7 @@ public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, } @Provides - public QueryService querySerivce(ExecutionEngine executionEngine) { + public QueryService queryService(ExecutionEngine executionEngine) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); From 819ddc661f948abf63001d32442fb9606c2506fa Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 21 Mar 2023 14:17:30 -0700 Subject: [PATCH 41/46] Step 1. Signed-off-by: Yury-Fridlyand --- .../sql/executor/PaginatedPlanCache.java | 162 +++++++++--------- .../sql/opensearch/executor/Cursor.java | 16 +- .../sql/planner/SerializablePlan.java | 80 +++++++++ .../planner/physical/PaginateOperator.java | 78 ++++++--- .../sql/planner/physical/PhysicalPlan.java | 30 ++-- .../sql/planner/physical/ProjectOperator.java | 63 +++++-- .../sql/executor/PaginatedPlanCacheTest.java | 4 +- .../physical/PaginateOperatorTest.java | 8 +- .../planner/physical/PhysicalPlanTest.java | 4 +- .../planner/physical/ProjectOperatorTest.java | 8 +- .../value/OpenSearchExprValueFactory.java | 3 +- .../protector/ResourceMonitorPlan.java | 28 ++- .../request/InitialPageRequestBuilder.java | 10 +- .../request/PagedRequestBuilder.java | 13 +- .../request/SubsequentPageRequestBuilder.java | 13 +- .../setting/OpenSearchSettings.java | 4 +- .../opensearch/storage/OpenSearchIndex.java | 2 +- .../scan/OpenSearchPagedIndexScan.java | 71 +++++++- .../OpenSearchExecutionEngineTest.java | 2 +- .../executor/ResourceMonitorPlanTest.java | 4 +- .../scan/OpenSearchPagedIndexScanTest.java | 6 +- 21 files changed, 415 insertions(+), 194 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 7a876e695a3..25af7c1d47f 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -8,37 +8,63 @@ import com.google.common.hash.HashCode; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.util.ArrayList; -import java.util.List; +import java.io.Externalizable; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectInputStream; +import java.io.ObjectOutput; +import java.io.ObjectOutputStream; +import java.io.Serializable; +import java.util.zip.Deflater; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; + +import lombok.AllArgsConstructor; import lombok.Data; +import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.SneakyThrows; import org.opensearch.sql.ast.tree.UnresolvedPlan; -import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.executor.Cursor; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.TableScanOperator; -@RequiredArgsConstructor public class PaginatedPlanCache { public static final String CURSOR_PREFIX = "n:"; private final StorageEngine storageEngine; public static final PaginatedPlanCache None = new PaginatedPlanCache(null); + public PaginatedPlanCache(StorageEngine storageEngine) { + this.storageEngine = storageEngine; + SerializationContext.engine = storageEngine; + } + public boolean canConvertToCursor(UnresolvedPlan plan) { return plan.accept(new CanPaginateVisitor(), null); } - @RequiredArgsConstructor @Data - static class SerializationContext { - private final PaginatedPlanCache cache; + @AllArgsConstructor + public static class SerializationContext implements Externalizable { + private PaginateOperator plan; + private static StorageEngine engine; + + public SerializationContext() { + } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + plan.writeExternal(out); + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + var loader = (SerializablePlan.PlanLoader) in.readObject(); + plan = (PaginateOperator) loader.apply(in, engine); + } } /** @@ -46,17 +72,46 @@ static class SerializationContext { */ public Cursor convertToCursor(PhysicalPlan plan) { if (plan instanceof PaginateOperator) { - var cursor = plan.toCursor(); - if (cursor == null) { - return Cursor.None; - } - var raw = CURSOR_PREFIX + compress(cursor); - return new Cursor(raw.getBytes()); + var context = new SerializationContext((PaginateOperator) plan); + //plan.prepareToSerialization(context); + //return new Cursor(CURSOR_PREFIX + serialize(new Object[] { plan, context })); + return new Cursor(CURSOR_PREFIX + serialize(context)); } else { return Cursor.None; } } + public String serialize(Serializable object) { + try { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + ObjectOutputStream objectOutput = new ObjectOutputStream(output); + objectOutput.writeObject(object); + objectOutput.flush(); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + GZIPOutputStream gzip = new GZIPOutputStream(out) { { + this.def.setLevel(Deflater.BEST_COMPRESSION); + } }; + gzip.write(output.toByteArray()); + gzip.close(); + return HashCode.fromBytes(out.toByteArray()).toString(); + } catch (IOException e) { + throw new IllegalStateException("Failed to serialize: " + object, e); + } + } + + public Serializable deserialize(String code) { + try { + GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream( + HashCode.fromString(code).asBytes())); + ByteArrayInputStream input = new ByteArrayInputStream(gzip.readAllBytes()); + ObjectInputStream objectInput = new ObjectInputStream(input); + return (Serializable) objectInput.readObject(); + } catch (Exception e) { + throw new IllegalStateException("Failed to deserialize object", e); + } + } + /** * Compress serialized query plan. * @param str string representing a query plan @@ -75,7 +130,7 @@ public static String compress(String str) { } /** - * Decompresses a query plan that was compress with {@link PaginatedPlanCache.compress}. + * Decompresses a query plan that was compress with {@link PaginatedPlanCache#compress}. * @param input compressed query plan * @return seria */ @@ -90,76 +145,19 @@ public static String decompress(String input) { } /** - * Parse `NamedExpression`s from cursor. - * @param listToFill List to fill with data. - * @param cursor Cursor to parse. - * @return Remaining part of the cursor. + * Converts a cursor to a physical plan tree. */ - private String parseNamedExpressions(List listToFill, String cursor) { - var serializer = new DefaultExpressionSerializer(); - if (cursor.startsWith(")")) { //empty list - return cursor.substring(cursor.indexOf(',') + 1); - } - while (!cursor.startsWith("(")) { - listToFill.add((NamedExpression) - serializer.deserialize(cursor.substring(0, - Math.min(cursor.indexOf(','), cursor.indexOf(')'))))); - cursor = cursor.substring(cursor.indexOf(',') + 1); - } - return cursor; - } - - /** - * Converts a cursor to a physical plan tree. - */ public PhysicalPlan convertToPlan(String cursor) { if (cursor.startsWith(CURSOR_PREFIX)) { try { - cursor = cursor.substring(CURSOR_PREFIX.length()); - cursor = decompress(cursor); - - // TODO Parse with ANTLR or serialize as JSON/XML - if (!cursor.startsWith("(Paginate,")) { - throw new UnsupportedOperationException("Unsupported cursor"); - } - // TODO add checks for > 0 - cursor = cursor.substring(cursor.indexOf(',') + 1); - final int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); - - cursor = cursor.substring(cursor.indexOf(',') + 1); - final int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10); - - cursor = cursor.substring(cursor.indexOf(',') + 1); - if (!cursor.startsWith("(Project,")) { - throw new UnsupportedOperationException("Unsupported cursor"); - } - cursor = cursor.substring(cursor.indexOf(',') + 1); - if (!cursor.startsWith("(namedParseExpressions,")) { - throw new UnsupportedOperationException("Unsupported cursor"); - } - - cursor = cursor.substring(cursor.indexOf(',') + 1); - List namedParseExpressions = new ArrayList<>(); - cursor = parseNamedExpressions(namedParseExpressions, cursor); - - List projectList = new ArrayList<>(); - if (!cursor.startsWith("(projectList,")) { - throw new UnsupportedOperationException("Unsupported cursor"); - } - cursor = cursor.substring(cursor.indexOf(',') + 1); - cursor = parseNamedExpressions(projectList, cursor); - - if (!cursor.startsWith("(OpenSearchPagedIndexScan,")) { - throw new UnsupportedOperationException("Unsupported cursor"); - } - cursor = cursor.substring(cursor.indexOf(',') + 1); - var indexName = cursor.substring(0, cursor.indexOf(',')); - cursor = cursor.substring(cursor.indexOf(',') + 1); - var scrollId = cursor.substring(0, cursor.indexOf(')')); - TableScanOperator scan = storageEngine.getTableScan(indexName, scrollId); - - return new PaginateOperator(new ProjectOperator(scan, projectList, namedParseExpressions), - pageSize, currentPageIndex); + //var data = (Object[]) deserialize(cursor.substring(CURSOR_PREFIX.length())); + //var plan = (PhysicalPlan) data[0]; + //var context = (SerializationContext) data[1]; + //TableScanOperator scan = storageEngine.getTableScan(context.getIndexName(), context.getScrollId()); + + //Class.forName("PaginateOperator").getDeclaredConstructor() + + return ((SerializationContext) deserialize(cursor.substring(CURSOR_PREFIX.length()))).getPlan(); } catch (Exception e) { throw new UnsupportedOperationException("Unsupported cursor", e); } diff --git a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java index 73289d9066c..99ceab07cd4 100644 --- a/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java +++ b/core/src/main/java/org/opensearch/sql/opensearch/executor/Cursor.java @@ -7,23 +7,17 @@ import lombok.EqualsAndHashCode; import lombok.Getter; +import lombok.RequiredArgsConstructor; @EqualsAndHashCode +@RequiredArgsConstructor public class Cursor { - public static final Cursor None = new Cursor(); + public static final Cursor None = new Cursor(null); @Getter - private final byte[] raw; - - private Cursor() { - raw = new byte[] {}; - } - - public Cursor(byte[] raw) { - this.raw = raw; - } + private final String data; public String toString() { - return new String(raw); + return data; } } diff --git a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java new file mode 100644 index 00000000000..fabe2002a38 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java @@ -0,0 +1,80 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner; + +import org.opensearch.sql.storage.StorageEngine; + +import java.io.Externalizable; +import java.io.IOException; +import java.io.NotSerializableException; +import java.io.ObjectInput; +import java.io.ObjectOutput; +import java.io.Serializable; +import java.util.function.BiFunction; + +/** + * All instances of PhysicalPlan which needs to be serialized (in cursor feature) should + * have a public no-arg constructor and override all given here methods. + */ +public abstract class SerializablePlan implements Externalizable { + +// /** +// * Prep +// * Each plan which supports serialization should override this method, to do at least nothing. +// */ +// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { +// throw new IllegalStateException(String.format("%s is not compatible with cursor feature", +// this.getClass().getSimpleName())); +// /* Non default implementation should be like reverse visitor +// context.setSomething(data); +// getChild().forEach(plan -> plan.prepareToSerialization(context)); +// */ +// } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + throw new NotSerializableException(); + /* Each plan which supports serialization should dump itself into the stream and go recursive. + out.writeSomething(data); + for (var plan : getChild()) { + plan.writeExternal(out.getPlanForSerialization()); + } + */ + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + throw new NotSerializableException(); + /* Each plan which supports serialization should load itself from the stream and go recursive. + this.data = in.readSomething(); + for (var plan : getChild()) { + plan.readExternal(in); + } + */ + } + + /** + * TODO update comment + * Override to return null, so parent plan should skip this child for serialization, but + * it should try to serialize grandchild plan. + * + * Imagine plan structure like this + * A -> false + * `- B -> true + * `- C -> false + * In that case only plans A and C should be attempted to serialize. + * It is needed to skip a `ResourceMonitorPlan` instance only, actually. + */ + public SerializablePlan getPlanForSerialization() { + return this; + } + + @FunctionalInterface + public interface PlanLoader extends Serializable { + SerializablePlan apply(ObjectInput in, StorageEngine engine) + throws IOException, ClassNotFoundException; + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java index 33856015f62..9eea6f2388d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java @@ -5,24 +5,27 @@ package org.opensearch.sql.planner.physical; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.List; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.executor.ExecutionEngine; -import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; -import org.opensearch.sql.planner.physical.ProjectOperator; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.expression.NamedExpression; +import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.storage.StorageEngine; -@RequiredArgsConstructor @EqualsAndHashCode(callSuper = false) public class PaginateOperator extends PhysicalPlan { @Getter - private final PhysicalPlan input; + private PhysicalPlan input; @Getter - private final int pageSize; + private int pageSize; /** * Which page is this? @@ -30,9 +33,14 @@ public class PaginateOperator extends PhysicalPlan { * See usage. */ @Getter - private final int pageIndex; + private int pageIndex = 0; - int numReturned = 0; + private int numReturned = 0; + + public PaginateOperator() { + int a = 5; + // TODO validate that called only from deserializer + } /** * Page given physical plan, with pageSize elements per page, starting with the first page. @@ -40,7 +48,15 @@ public class PaginateOperator extends PhysicalPlan { public PaginateOperator(PhysicalPlan input, int pageSize) { this.pageSize = pageSize; this.input = input; - this.pageIndex = 0; + } + + /** + * Page given physical plan, with pageSize elements per page, starting with the given page. + */ + public PaginateOperator(PhysicalPlan input, int pageSize, int pageIndex) { + this.pageSize = pageSize; + this.input = input; + this.pageIndex = pageIndex; } @Override @@ -79,17 +95,39 @@ public ExecutionEngine.Schema schema() { return input.schema(); } +// @Override +// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { +// pageIndex++; +// } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + PlanLoader loader = (in, engine) -> { + var pageSize = in.readInt(); + var pageIndex = in.readInt(); + var inputLoader = (PlanLoader) in.readObject(); + var input = (PhysicalPlan) inputLoader.apply(in, engine); + return new PaginateOperator(input, pageSize, pageIndex); + }; + out.writeObject(loader); + + out.writeInt(pageSize); + out.writeInt(pageIndex + 1); + input.getPlanForSerialization().writeExternal(out); + } + @Override - public String toCursor() { - // Save cursor to read the next page. - // Could process node.getChild() here with another visitor -- one that saves the - // parameters for other physical operators -- ProjectOperator, etc. - // cursor format: n:|" - String child = getChild().get(0).toCursor(); - - var nextPage = getPageIndex() + 1; - return child == null || child.isEmpty() - ? null : createSection("Paginate", Integer.toString(nextPage), - Integer.toString(getPageSize()), child); + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + // nothing, everything done by loader } +/* + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + var loader = (PlanLoader) in.readObject(); + this = loader.apply(in, engine); + + input = (PhysicalPlan) in.readObject(); + pageSize = in.readInt(); + pageIndex = in.readInt(); + }*/ } 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 312e4bfff9a..480137511ca 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 @@ -6,19 +6,26 @@ package org.opensearch.sql.planner.physical; +import java.io.Externalizable; +import java.io.IOException; +import java.io.NotSerializableException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Iterator; import java.util.List; + import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.planner.PlanNode; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.storage.split.Split; /** * Physical plan. */ -public abstract class PhysicalPlan implements PlanNode, - Iterator, - AutoCloseable { +public abstract class PhysicalPlan extends SerializablePlan + implements PlanNode, Iterator, AutoCloseable { /** * Accept the {@link PhysicalPlanNodeVisitor}. * @@ -57,21 +64,4 @@ public ExecutionEngine.Schema schema() { public long getTotalHits() { return getChild().stream().mapToLong(PhysicalPlan::getTotalHits).max().orElse(0); } - - public String toCursor() { - throw new IllegalStateException(String.format("%s is not compatible with cursor feature", - this.getClass().getSimpleName())); - } - - /** - * Creates an S-expression that represents a plan node. - * @param plan Label for the plan. - * @param params List of serialized parameters. Including the child plans. - * @return A string that represents the plan called with those parameters. - */ - protected String createSection(String plan, String... params) { - return "(" + plan + "," - + String.join(",", params) - + ")"; - } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index c61b35e0cb6..177cc2dc0bc 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -8,6 +8,10 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMap.Builder; + +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Collections; import java.util.List; import java.util.Optional; @@ -20,23 +24,36 @@ import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.storage.StorageEngine; /** * Project the fields specified in {@link ProjectOperator#projectList} from input. */ @ToString @EqualsAndHashCode(callSuper = false) -@RequiredArgsConstructor public class ProjectOperator extends PhysicalPlan { @Getter - private final PhysicalPlan input; + private PhysicalPlan input; @Getter - private final List projectList; + private List projectList; @Getter - private final List namedParseExpressions; + private List namedParseExpressions; + + public ProjectOperator() { + int a = 5; + // TODO validate that called only from deserializer + } + + public ProjectOperator(PhysicalPlan input, List projectList, + List namedParseExpressions) { + this.input = input; + this.projectList = projectList; + this.namedParseExpressions = namedParseExpressions; + } @Override public R accept(PhysicalPlanNodeVisitor visitor, C context) { @@ -96,17 +113,33 @@ public ExecutionEngine.Schema schema() { expr.getAlias(), expr.type())).collect(Collectors.toList())); } + @SuppressWarnings("unchecked") @Override - public String toCursor() { - String child = getChild().get(0).toCursor(); - if (child == null || child.isEmpty()) { - return null; - } - var serializer = new DefaultExpressionSerializer(); - String projects = createSection("projectList", - projectList.stream().map(serializer::serialize).toArray(String[]::new)); - String namedExpressions = createSection("namedParseExpressions", - namedParseExpressions.stream().map(serializer::serialize).toArray(String[]::new)); - return createSection("Project", namedExpressions, projects, child); + public void writeExternal(ObjectOutput out) throws IOException { + PlanLoader loader = (in, engine) -> { + var projectList = (List) in.readObject(); + var namedParseExpressions = (List) in.readObject(); + var inputLoader = (PlanLoader) in.readObject(); + var input = (PhysicalPlan) inputLoader.apply(in, engine); + return new ProjectOperator(input, projectList, namedParseExpressions); + }; + out.writeObject(loader); + + out.writeObject(projectList); + out.writeObject(namedParseExpressions); + input.getPlanForSerialization().writeExternal(out); + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + // nothing, everything done by loader } +/* + @SuppressWarnings("unchecked") + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + input = (PhysicalPlan) in.readObject(); + projectList = (List) in.readObject(); + namedParseExpressions = (List) in.readObject(); + }*/ } diff --git a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java index e99612c5f3a..b64b1ad8317 100644 --- a/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/PaginatedPlanCacheTest.java @@ -359,7 +359,7 @@ void convert_deconvert_cursor() { void convertToCursor_cant_convert() { var plan = mock(MockedTableScanOperator.class); assertEquals(Cursor.None, planCache.convertToCursor(plan)); - when(plan.toCursor()).thenReturn(""); + when(plan.prepareToCursorSerialization()).thenReturn(""); assertEquals(Cursor.None, planCache.convertToCursor( new PaginateOperator(plan, 1, 2))); } @@ -443,7 +443,7 @@ public String explain() { } @Override - public String toCursor() { + public void prepareToCursorSerialization() { return createSection("OpenSearchPagedIndexScan", testIndexName, testScroll); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java index 19df23b9e58..8ab76b78649 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/PaginateOperatorTest.java @@ -87,13 +87,13 @@ public void schema_assert() { @Test public void toCursor() { var plan = mock(PhysicalPlan.class); - when(plan.toCursor()).thenReturn("Great plan, Walter, reliable as a swiss watch!", "", null); + when(plan.prepareToSerialization()).thenReturn("Great plan, Walter, reliable as a swiss watch!", "", null); var po = new PaginateOperator(plan, 2); assertAll( () -> assertEquals("(Paginate,1,2,Great plan, Walter, reliable as a swiss watch!)", - po.toCursor()), - () -> assertNull(po.toCursor()), - () -> assertNull(po.toCursor()) + po.prepareToSerialization()), + () -> assertNull(po.prepareToSerialization()), + () -> assertNull(po.prepareToSerialization()) ); } } 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 5e70f2b9d01..50d0a6299d5 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 @@ -82,8 +82,8 @@ void get_total_hits_uses_default_value() { @Test void toCursor() { var plan = mock(PhysicalPlan.class); - when(plan.toCursor()).then(CALLS_REAL_METHODS); - assertTrue(assertThrows(IllegalStateException.class, plan::toCursor) + when(plan.prepareToSerialization()).then(CALLS_REAL_METHODS); + assertTrue(assertThrows(IllegalStateException.class, plan::prepareToSerialization) .getMessage().contains("is not compatible with cursor feature")); } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index 6042eba6dcc..8b8995f5f21 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -213,7 +213,7 @@ public void project_parse_missing_will_fallback() { @Test public void toCursor() { - when(inputPlan.toCursor()).thenReturn("inputPlan", "", null); + when(inputPlan.prepareToSerialization()).thenReturn("inputPlan", "", null); var project = DSL.named("response", DSL.ref("response", INTEGER)); var npe = DSL.named("action", DSL.ref("action", STRING)); var po = project(inputPlan, List.of(project), List.of(npe)); @@ -221,9 +221,9 @@ public void toCursor() { var expected = String.format("(Project,(namedParseExpressions,%s),(projectList,%s),%s)", serializer.serialize(npe), serializer.serialize(project), "inputPlan"); assertAll( - () -> assertEquals(expected, po.toCursor()), - () -> assertNull(po.toCursor()), - () -> assertNull(po.toCursor()) + () -> assertEquals(expected, po.prepareToSerialization()), + () -> assertNull(po.prepareToSerialization()), + () -> assertNull(po.prepareToSerialization()) ); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java index 2536121e914..8f7cb939db0 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java @@ -30,6 +30,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableMap; +import java.io.Serializable; import java.time.Instant; import java.time.format.DateTimeParseException; import java.util.ArrayList; @@ -66,7 +67,7 @@ /** * Construct ExprValue from OpenSearch response. */ -public class OpenSearchExprValueFactory { +public class OpenSearchExprValueFactory implements Serializable { /** * The Mapping of Field and ExprType. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index 3d880d82b9f..f697bf47c92 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -6,12 +6,17 @@ package org.opensearch.sql.opensearch.executor.protector; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.List; import lombok.EqualsAndHashCode; import lombok.RequiredArgsConstructor; import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.monitor.ResourceMonitor; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -88,8 +93,27 @@ public long getTotalHits() { return delegate.getTotalHits(); } +// @Override +// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { +// delegate.prepareToSerialization(context); +// } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + // do nothing, we shouldn't serialize ResourceMonitorPlan + // nor its delegate (instance of TableScanOperator). + delegate.writeExternal(out); + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + // do nothing, we shouldn't serialize ResourceMonitorPlan + // nor its delegate (instance of TableScanOperator). + delegate.readExternal(in); + } + @Override - public String toCursor() { - return delegate.toCursor(); + public SerializablePlan getPlanForSerialization() { + return delegate; } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java index 7ef87473816..84bde63a442 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java @@ -23,31 +23,27 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -public class InitialPageRequestBuilder implements PagedRequestBuilder { +public class InitialPageRequestBuilder extends PagedRequestBuilder { @Getter private final OpenSearchRequest.IndexName indexName; private final SearchSourceBuilder sourceBuilder; - private final OpenSearchExprValueFactory exprValueFactory; - private final int querySize; /** * Constructor. * @param indexName index being scanned - * @param settings other settings + * @param pageSize page size * @param exprValueFactory value factory */ // TODO accept indexName as string (same way as `OpenSearchRequestBuilder` does)? public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, int pageSize, - Settings settings, // TODO: settings are not used - refactor? OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; this.exprValueFactory = exprValueFactory; - this.querySize = pageSize; this.sourceBuilder = new SearchSourceBuilder() .from(0) - .size(querySize) + .size(pageSize) .timeout(DEFAULT_QUERY_TIMEOUT); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java index 365c4a60615..6afadeda5c8 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java @@ -5,8 +5,15 @@ package org.opensearch.sql.opensearch.request; -public interface PagedRequestBuilder { - OpenSearchRequest build(); +import lombok.Getter; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; - OpenSearchRequest.IndexName getIndexName(); +public abstract class PagedRequestBuilder { + + @Getter + protected OpenSearchExprValueFactory exprValueFactory; + + abstract public OpenSearchRequest build(); + + abstract public OpenSearchRequest.IndexName getIndexName(); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java index c2fb1f3431a..062032e44f9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java @@ -6,16 +6,21 @@ package org.opensearch.sql.opensearch.request; import lombok.Getter; -import lombok.RequiredArgsConstructor; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -@RequiredArgsConstructor -public class SubsequentPageRequestBuilder implements PagedRequestBuilder { +public class SubsequentPageRequestBuilder extends PagedRequestBuilder { + + public SubsequentPageRequestBuilder(OpenSearchRequest.IndexName indexName, + String scrollId, + OpenSearchExprValueFactory factory) { + this.indexName = indexName; + this.scrollId = scrollId; + this.exprValueFactory = factory; + } @Getter private final OpenSearchRequest.IndexName indexName; private final String scrollId; - private final OpenSearchExprValueFactory exprValueFactory; @Override public OpenSearchRequest build() { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java index ae5174d678f..accd3560417 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java @@ -99,8 +99,8 @@ public class OpenSearchSettings extends Settings { Setting.Property.Dynamic); /** - * Construct ElasticsearchSetting. - * The ElasticsearchSetting must be singleton. + * Construct OpenSearchSetting. + * The OpenSearchSetting must be singleton. */ @SuppressWarnings("unchecked") public OpenSearchSettings(ClusterSettings clusterSettings) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index c5ec6c364b2..09102c3947a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -134,7 +134,7 @@ public TableScanBuilder createScanBuilder() { @Override public TableScanBuilder createPagedScanBuilder(int pageSize) { var requestBuilder = new InitialPageRequestBuilder(indexName, pageSize, - settings, new OpenSearchExprValueFactory(getFieldTypes())); + new OpenSearchExprValueFactory(getFieldTypes())); var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); return new OpenSearchPagedIndexScanBuilder(indexScan); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index e9d3fd52d39..9ef7832775f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -5,29 +5,54 @@ package org.opensearch.sql.opensearch.storage.scan; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Collections; import java.util.Iterator; +import java.util.Map; + import lombok.EqualsAndHashCode; import lombok.ToString; import org.apache.commons.lang3.NotImplementedException; +import org.opensearch.client.node.NodeClient; +import org.opensearch.common.inject.ModulesBuilder; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.datasource.DataSourceService; +import org.opensearch.sql.executor.PaginatedPlanCache; +import org.opensearch.sql.expression.function.SerializableBiFunction; +import org.opensearch.sql.expression.function.SerializableFunction; import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.ContinueScrollRequest; +import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.PagedRequestBuilder; +import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.security.SecurityAccess; +import org.opensearch.sql.opensearch.setting.OpenSearchSettings; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.TableScanOperator; @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) @ToString(onlyExplicitlyIncluded = true) public class OpenSearchPagedIndexScan extends TableScanOperator { - private final OpenSearchClient client; - private final PagedRequestBuilder requestBuilder; + private OpenSearchClient client; + private PagedRequestBuilder requestBuilder; @EqualsAndHashCode.Include @ToString.Include private OpenSearchRequest request; private Iterator iterator; private long totalHits = 0; + public OpenSearchPagedIndexScan() { + int a = 5; + // TODO validate that called only from deserializer + } + public OpenSearchPagedIndexScan(OpenSearchClient client, PagedRequestBuilder requestBuilder) { this.client = client; @@ -73,12 +98,42 @@ public long getTotalHits() { return totalHits; } +// @Override +// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { +// context.setIndexName(requestBuilder.getIndexName().toString()); +// context.setScrollId(request.toCursor()); +// } + @Override - public String toCursor() { - // TODO this assumes exactly one index is scanned. - var indexName = requestBuilder.getIndexName().getIndexNames()[0]; - var cursor = request.toCursor(); - return cursor == null || cursor.isEmpty() - ? "" : createSection("OpenSearchPagedIndexScan", indexName, cursor); + public void writeExternal(ObjectOutput out) throws IOException { + //out.writeObject(requestBuilder.getExprValueFactory()); + + PlanLoader loader = (in, engine) -> { + var indexName = (String) in.readUTF(); + var scrollId = (String) in.readUTF(); + return engine.getTableScan(indexName, scrollId); + }; + out.writeObject(loader); + out.writeUTF(requestBuilder.getIndexName().toString()); + out.writeUTF(request.toCursor()); } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + // nothing, everything done by loader + } +/* + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + //var factory = (OpenSearchExprValueFactory) in.readObject(); + var indexName = (String) in.readUTF(); + var scrollId = (String) in.readUTF(); + requestBuilder = new SubsequentPageRequestBuilder(new OpenSearchRequest.IndexName(indexName), + //scrollId, factory); + scrollId, new OpenSearchExprValueFactory(Map.of())); + + ModulesBuilder modules = new ModulesBuilder(); + var injector = modules.createInjector(); + client = SecurityAccess.doPrivileged(() -> injector.getInstance(OpenSearchClient.class)); + }*/ } 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 c39609cce74..86ff09c1f32 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 @@ -287,7 +287,7 @@ public ExecutionEngine.Schema schema() { } @Override - public String toCursor() { + public void prepareToCursorSerialization() { return "FakePaginatePlan"; } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java index 7b1353f4a97..4eb049d540d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java @@ -116,7 +116,7 @@ void getTotalHitsSuccess() { @Test void toCursorSuccess() { - monitorPlan.toCursor(); - verify(plan, times(1)).toCursor(); + monitorPlan.prepareToSerialization(); + verify(plan, times(1)).prepareToSerialization(); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java index eab71e3f4f1..93612c6e56a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java @@ -157,9 +157,9 @@ void toCursor() { OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder); indexScan.open(); assertAll( - () -> assertEquals("(OpenSearchPagedIndexScan,index,cu-cursor)", indexScan.toCursor()), - () -> assertEquals("", indexScan.toCursor()), - () -> assertEquals("", indexScan.toCursor()) + () -> assertEquals("(OpenSearchPagedIndexScan,index,cu-cursor)", indexScan.prepareToSerialization()), + () -> assertEquals("", indexScan.prepareToSerialization()), + () -> assertEquals("", indexScan.prepareToSerialization()) ); } } From 7a9d285a35dc0b6c33412893c5f653759b5f9b0f Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Tue, 21 Mar 2023 15:40:42 -0700 Subject: [PATCH 42/46] Step 2. Signed-off-by: Yury-Fridlyand --- .../sql/executor/PaginatedPlanCache.java | 74 ++++++------------- .../sql/planner/SerializablePlan.java | 41 ++-------- .../planner/physical/PaginateOperator.java | 29 +------- .../sql/planner/physical/ProjectOperator.java | 21 +----- .../protector/ResourceMonitorPlan.java | 18 +---- .../scan/OpenSearchPagedIndexScan.java | 50 ++----------- 6 files changed, 40 insertions(+), 193 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java index 25af7c1d47f..a5abd186cc4 100644 --- a/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java +++ b/core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java @@ -18,27 +18,20 @@ import java.util.zip.Deflater; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; - -import lombok.AllArgsConstructor; import lombok.Data; -import lombok.Getter; -import lombok.RequiredArgsConstructor; -import lombok.SneakyThrows; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.opensearch.executor.Cursor; import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PaginateOperator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; -import org.opensearch.sql.storage.TableScanOperator; -public class PaginatedPlanCache { +public class PaginatedPlanCache implements AutoCloseable { public static final String CURSOR_PREFIX = "n:"; - private final StorageEngine storageEngine; + //TODO remove - used in tests only public static final PaginatedPlanCache None = new PaginatedPlanCache(null); public PaginatedPlanCache(StorageEngine storageEngine) { - this.storageEngine = storageEngine; SerializationContext.engine = storageEngine; } @@ -46,18 +39,33 @@ public boolean canConvertToCursor(UnresolvedPlan plan) { return plan.accept(new CanPaginateVisitor(), null); } + // Actually called only once on server lifetime - on shutdown, so it does nothing. + @Override + public void close() throws Exception { + SerializationContext.engine = null; + } + @Data - @AllArgsConstructor public static class SerializationContext implements Externalizable { private PaginateOperator plan; private static StorageEngine engine; + /** + * If exception is thrown we don't catch it, that means something really went wrong. + * But if we can't serialize the plan, we set this flag and should return an empty cursor. + * The only case when it could happen as of now - paging is finished and there is no scroll. + */ + private boolean serializedSuccessfully = false; public SerializationContext() { } + public SerializationContext(PaginateOperator plan) { + this.plan = plan; + } + @Override public void writeExternal(ObjectOutput out) throws IOException { - plan.writeExternal(out); + serializedSuccessfully = plan.writeExternal(out); } @Override @@ -73,9 +81,8 @@ public void readExternal(ObjectInput in) throws IOException, ClassNotFoundExcept public Cursor convertToCursor(PhysicalPlan plan) { if (plan instanceof PaginateOperator) { var context = new SerializationContext((PaginateOperator) plan); - //plan.prepareToSerialization(context); - //return new Cursor(CURSOR_PREFIX + serialize(new Object[] { plan, context })); - return new Cursor(CURSOR_PREFIX + serialize(context)); + var serialized = serialize(context); + return context.serializedSuccessfully ? new Cursor(CURSOR_PREFIX + serialized) : Cursor.None; } else { return Cursor.None; } @@ -112,51 +119,12 @@ public Serializable deserialize(String code) { } } - /** - * Compress serialized query plan. - * @param str string representing a query plan - * @return str compressed with gzip. - */ - @SneakyThrows - public static String compress(String str) { - if (str == null || str.length() == 0) { - return null; - } - ByteArrayOutputStream out = new ByteArrayOutputStream(); - GZIPOutputStream gzip = new GZIPOutputStream(out); - gzip.write(str.getBytes()); - gzip.close(); - return HashCode.fromBytes(out.toByteArray()).toString(); - } - - /** - * Decompresses a query plan that was compress with {@link PaginatedPlanCache#compress}. - * @param input compressed query plan - * @return seria - */ - @SneakyThrows - public static String decompress(String input) { - if (input == null || input.length() == 0) { - return null; - } - GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream( - HashCode.fromString(input).asBytes())); - return new String(gzip.readAllBytes()); - } - /** * Converts a cursor to a physical plan tree. */ public PhysicalPlan convertToPlan(String cursor) { if (cursor.startsWith(CURSOR_PREFIX)) { try { - //var data = (Object[]) deserialize(cursor.substring(CURSOR_PREFIX.length())); - //var plan = (PhysicalPlan) data[0]; - //var context = (SerializationContext) data[1]; - //TableScanOperator scan = storageEngine.getTableScan(context.getIndexName(), context.getScrollId()); - - //Class.forName("PaginateOperator").getDeclaredConstructor() - return ((SerializationContext) deserialize(cursor.substring(CURSOR_PREFIX.length()))).getPlan(); } catch (Exception e) { throw new UnsupportedOperationException("Unsupported cursor", e); diff --git a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java index fabe2002a38..d457c8b79c0 100644 --- a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java @@ -5,39 +5,24 @@ package org.opensearch.sql.planner; -import org.opensearch.sql.storage.StorageEngine; - -import java.io.Externalizable; import java.io.IOException; -import java.io.NotSerializableException; import java.io.ObjectInput; import java.io.ObjectOutput; import java.io.Serializable; -import java.util.function.BiFunction; +import org.apache.commons.lang3.NotImplementedException; +import org.opensearch.sql.storage.StorageEngine; /** * All instances of PhysicalPlan which needs to be serialized (in cursor feature) should - * have a public no-arg constructor and override all given here methods. + * override all given here methods. */ -public abstract class SerializablePlan implements Externalizable { - -// /** -// * Prep -// * Each plan which supports serialization should override this method, to do at least nothing. -// */ -// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { -// throw new IllegalStateException(String.format("%s is not compatible with cursor feature", -// this.getClass().getSimpleName())); -// /* Non default implementation should be like reverse visitor -// context.setSomething(data); -// getChild().forEach(plan -> plan.prepareToSerialization(context)); -// */ -// } +public abstract class SerializablePlan { - @Override - public void writeExternal(ObjectOutput out) throws IOException { - throw new NotSerializableException(); + // Copied from Externalizable + public boolean writeExternal(ObjectOutput out) throws IOException { + throw new NotImplementedException(); /* Each plan which supports serialization should dump itself into the stream and go recursive. + TODO update comment out.writeSomething(data); for (var plan : getChild()) { plan.writeExternal(out.getPlanForSerialization()); @@ -45,16 +30,6 @@ public void writeExternal(ObjectOutput out) throws IOException { */ } - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - throw new NotSerializableException(); - /* Each plan which supports serialization should load itself from the stream and go recursive. - this.data = in.readSomething(); - for (var plan : getChild()) { - plan.readExternal(in); - } - */ - } /** * TODO update comment diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java index 9eea6f2388d..d959a7d645a 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/PaginateOperator.java @@ -11,13 +11,8 @@ import java.util.List; import lombok.EqualsAndHashCode; import lombok.Getter; -import lombok.RequiredArgsConstructor; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.executor.ExecutionEngine; -import org.opensearch.sql.executor.PaginatedPlanCache; -import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.planner.SerializablePlan; -import org.opensearch.sql.storage.StorageEngine; @EqualsAndHashCode(callSuper = false) public class PaginateOperator extends PhysicalPlan { @@ -95,13 +90,8 @@ public ExecutionEngine.Schema schema() { return input.schema(); } -// @Override -// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { -// pageIndex++; -// } - @Override - public void writeExternal(ObjectOutput out) throws IOException { + public boolean writeExternal(ObjectOutput out) throws IOException { PlanLoader loader = (in, engine) -> { var pageSize = in.readInt(); var pageIndex = in.readInt(); @@ -113,21 +103,6 @@ public void writeExternal(ObjectOutput out) throws IOException { out.writeInt(pageSize); out.writeInt(pageIndex + 1); - input.getPlanForSerialization().writeExternal(out); - } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - // nothing, everything done by loader + return input.getPlanForSerialization().writeExternal(out); } -/* - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - var loader = (PlanLoader) in.readObject(); - this = loader.apply(in, engine); - - input = (PhysicalPlan) in.readObject(); - pageSize = in.readInt(); - pageIndex = in.readInt(); - }*/ } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 177cc2dc0bc..5e638a77efe 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -18,17 +18,13 @@ import java.util.stream.Collectors; import lombok.EqualsAndHashCode; import lombok.Getter; -import lombok.RequiredArgsConstructor; import lombok.ToString; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; -import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; -import org.opensearch.sql.storage.StorageEngine; /** * Project the fields specified in {@link ProjectOperator#projectList} from input. @@ -115,7 +111,7 @@ public ExecutionEngine.Schema schema() { @SuppressWarnings("unchecked") @Override - public void writeExternal(ObjectOutput out) throws IOException { + public boolean writeExternal(ObjectOutput out) throws IOException { PlanLoader loader = (in, engine) -> { var projectList = (List) in.readObject(); var namedParseExpressions = (List) in.readObject(); @@ -127,19 +123,6 @@ public void writeExternal(ObjectOutput out) throws IOException { out.writeObject(projectList); out.writeObject(namedParseExpressions); - input.getPlanForSerialization().writeExternal(out); + return input.getPlanForSerialization().writeExternal(out); } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - // nothing, everything done by loader - } -/* - @SuppressWarnings("unchecked") - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - input = (PhysicalPlan) in.readObject(); - projectList = (List) in.readObject(); - namedParseExpressions = (List) in.readObject(); - }*/ } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index f697bf47c92..be888df284d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -14,7 +14,6 @@ import lombok.RequiredArgsConstructor; import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; -import org.opensearch.sql.executor.PaginatedPlanCache; import org.opensearch.sql.monitor.ResourceMonitor; import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -93,23 +92,10 @@ public long getTotalHits() { return delegate.getTotalHits(); } -// @Override -// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { -// delegate.prepareToSerialization(context); -// } - - @Override - public void writeExternal(ObjectOutput out) throws IOException { - // do nothing, we shouldn't serialize ResourceMonitorPlan - // nor its delegate (instance of TableScanOperator). - delegate.writeExternal(out); - } - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + public boolean writeExternal(ObjectOutput out) throws IOException { // do nothing, we shouldn't serialize ResourceMonitorPlan - // nor its delegate (instance of TableScanOperator). - delegate.readExternal(in); + return delegate.writeExternal(out); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index 9ef7832775f..01217d9e57c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -6,35 +6,18 @@ package org.opensearch.sql.opensearch.storage.scan; import java.io.IOException; -import java.io.ObjectInput; import java.io.ObjectOutput; import java.util.Collections; import java.util.Iterator; -import java.util.Map; import lombok.EqualsAndHashCode; import lombok.ToString; import org.apache.commons.lang3.NotImplementedException; -import org.opensearch.client.node.NodeClient; -import org.opensearch.common.inject.ModulesBuilder; import org.opensearch.sql.data.model.ExprValue; -import org.opensearch.sql.datasource.DataSourceService; -import org.opensearch.sql.executor.PaginatedPlanCache; -import org.opensearch.sql.expression.function.SerializableBiFunction; -import org.opensearch.sql.expression.function.SerializableFunction; import org.opensearch.sql.opensearch.client.OpenSearchClient; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.ContinueScrollRequest; -import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.PagedRequestBuilder; -import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; -import org.opensearch.sql.opensearch.security.SecurityAccess; -import org.opensearch.sql.opensearch.setting.OpenSearchSettings; -import org.opensearch.sql.opensearch.storage.OpenSearchIndex; -import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.TableScanOperator; @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) @@ -98,16 +81,11 @@ public long getTotalHits() { return totalHits; } -// @Override -// public void prepareToSerialization(PaginatedPlanCache.SerializationContext context) { -// context.setIndexName(requestBuilder.getIndexName().toString()); -// context.setScrollId(request.toCursor()); -// } - @Override - public void writeExternal(ObjectOutput out) throws IOException { - //out.writeObject(requestBuilder.getExprValueFactory()); - + public boolean writeExternal(ObjectOutput out) throws IOException { + if (request.toCursor() == null || request.toCursor().isEmpty()) { + return false; + } PlanLoader loader = (in, engine) -> { var indexName = (String) in.readUTF(); var scrollId = (String) in.readUTF(); @@ -116,24 +94,6 @@ public void writeExternal(ObjectOutput out) throws IOException { out.writeObject(loader); out.writeUTF(requestBuilder.getIndexName().toString()); out.writeUTF(request.toCursor()); + return true; } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - // nothing, everything done by loader - } -/* - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - //var factory = (OpenSearchExprValueFactory) in.readObject(); - var indexName = (String) in.readUTF(); - var scrollId = (String) in.readUTF(); - requestBuilder = new SubsequentPageRequestBuilder(new OpenSearchRequest.IndexName(indexName), - //scrollId, factory); - scrollId, new OpenSearchExprValueFactory(Map.of())); - - ModulesBuilder modules = new ModulesBuilder(); - var injector = modules.createInjector(); - client = SecurityAccess.doPrivileged(() -> injector.getInstance(OpenSearchClient.class)); - }*/ } From 5d085f720eedd8ea5f8315ca8a23aec4ee4a6522 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 22 Mar 2023 09:33:36 -0700 Subject: [PATCH 43/46] Support queries with `LIMIT`, `WHERE` and `ORDER BY` and queries without `FROM` into pagination. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/ast/tree/Sort.java | 4 +- .../sql/executor/CanPaginateVisitor.java | 78 +++++++++++++++---- .../sql/planner/SerializablePlan.java | 3 +- .../sql/planner/physical/FilterOperator.java | 38 ++++++++- .../sql/planner/physical/LimitOperator.java | 31 +++++++- .../sql/planner/physical/SortOperator.java | 21 ++++- .../sql/planner/physical/ValuesOperator.java | 8 ++ .../OpenSearchExecutionProtector.java | 5 +- 8 files changed, 163 insertions(+), 25 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Sort.java b/core/src/main/java/org/opensearch/sql/ast/tree/Sort.java index 5fb4139bea6..4cc71646e38 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Sort.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Sort.java @@ -12,6 +12,8 @@ import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC; import com.google.common.collect.ImmutableList; + +import java.io.Serializable; import java.util.List; import lombok.AllArgsConstructor; import lombok.Data; @@ -54,7 +56,7 @@ public T accept(AbstractNodeVisitor nodeVisitor, C context) { * Sort Options. */ @Data - public static class SortOption { + public static class SortOption implements Serializable { /** * Default ascending sort option, null first. diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index d4c1c2f3005..05ba08baf50 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -8,9 +8,18 @@ import java.util.concurrent.atomic.AtomicBoolean; 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.AllFields; +import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.WindowFunction; +import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.Relation; +import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.ast.tree.Values; /** * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache. @@ -39,58 +48,93 @@ public Boolean visitRelation(Relation node, Object context) { return Boolean.TRUE; } - /* private Boolean canPaginate(Node node, Object context) { AtomicBoolean result = new AtomicBoolean(true); - node.getChild().forEach(n -> result.set(result.get() && n.accept(this, context))); + var childList = node.getChild(); + if (childList != null) { + childList.forEach(n -> result.set(result.get() && n.accept(this, context))); + } return result.get(); } - For queries without `FROM` clause. - Required to overload `toCursor` function in `ValuesOperator` and modify cursor parsing. + //For queries without `FROM` clause. + //Required to overload `toCursor` function in `ValuesOperator` and modify cursor parsing. @Override public Boolean visitValues(Values node, Object context) { return canPaginate(node, context); } - For queries with LIMIT clause: - Required to overload `toCursor` function in `LimitOperator` and modify cursor parsing. + //For queries with LIMIT clause: + //Required to overload `toCursor` function in `LimitOperator` and modify cursor parsing. @Override public Boolean visitLimit(Limit node, Object context) { return canPaginate(node, context); } - For queries with ORDER BY clause: - Required to overload `toCursor` function in `SortOperator` and modify cursor parsing. + //For queries with ORDER BY clause: + //Required to overload `toCursor` function in `SortOperator` and modify cursor parsing. @Override public Boolean visitSort(Sort node, Object context) { return canPaginate(node, context); } - For queries with WHERE clause: - Required to overload `toCursor` function in `FilterOperator` and modify cursor parsing. + //For queries with WHERE clause: + //Required to overload `toCursor` function in `FilterOperator` and modify cursor parsing. @Override public Boolean visitFilter(Filter node, Object context) { return canPaginate(node, context); } - */ + + @Override + public Boolean visitLiteral(Literal node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitField(Field node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitAlias(Alias node, Object context) { + return canPaginate(node, context) && canPaginate(node.getDelegated(), context); + } + + @Override + public Boolean visitAllFields(AllFields node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitQualifiedName(QualifiedName node, Object context) { + return canPaginate(node, context); + } @Override public Boolean visitChildren(Node node, Object context) { return Boolean.FALSE; } + @Override + public Boolean visit(Node node, Object context) { + // for all not listed (= unchecked) - false + // TODO evaluate to return true or call `callPaginate` + return Boolean.FALSE; + } + + @Override + public Boolean visitWindowFunction(WindowFunction node, Object context) { + return Boolean.FALSE; + } + @Override public Boolean visitProject(Project node, Object context) { // Allow queries with 'SELECT *' only. Those restriction could be removed, but consider // in-memory aggregation performed by window function (see WindowOperator). // SELECT max(age) OVER (PARTITION BY city) ... - var projections = node.getProjectList(); - if (projections.size() != 1) { - return Boolean.FALSE; - } - - if (!(projections.get(0) instanceof AllFields)) { + AtomicBoolean result = new AtomicBoolean(true); + node.getProjectList().forEach(n -> result.set(result.get() && n.accept(this, context))); + if (!result.get()) { return Boolean.FALSE; } diff --git a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java index d457c8b79c0..e7709485ec7 100644 --- a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java @@ -20,7 +20,8 @@ public abstract class SerializablePlan { // Copied from Externalizable public boolean writeExternal(ObjectOutput out) throws IOException { - throw new NotImplementedException(); + throw new NotImplementedException(String.format("%s is not serializable", + getClass().getSimpleName())); /* Each plan which supports serialization should dump itself into the stream and go recursive. TODO update comment out.writeSomething(data); diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java index a9c7597c3e5..d35d816f01b 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java @@ -5,6 +5,8 @@ package org.opensearch.sql.planner.physical; +import java.io.IOException; +import java.io.ObjectOutput; import java.util.Collections; import java.util.List; import lombok.EqualsAndHashCode; @@ -34,6 +36,14 @@ public class FilterOperator extends PhysicalPlan { private ExprValue next = null; private long totalHits = 0; + // A copy constructor + public FilterOperator(PhysicalPlan input, FilterOperator other) { + this.input = input; + this.conditions = other.conditions; + this.next = other.next; + this.totalHits = other.totalHits; + } + @Override public R accept(PhysicalPlanNodeVisitor visitor, C context) { return visitor.visitFilter(this, context); @@ -46,6 +56,9 @@ public List getChild() { @Override public boolean hasNext() { + if (next != null) { + return true; + } while (input.hasNext()) { ExprValue inputValue = input.next(); ExprValue exprValue = conditions.valueOf(inputValue.bindingTuples()); @@ -60,7 +73,9 @@ public boolean hasNext() { @Override public ExprValue next() { - return next; + var res = next; + next = null; + return res; } @Override @@ -68,4 +83,25 @@ public long getTotalHits() { // ignore `input.getTotalHits()`, because it returns wrong (unfiltered) value return totalHits; } + + @Override + public boolean writeExternal(ObjectOutput out) throws IOException { + PlanLoader loader = (in, engine) -> { + var conditions = (Expression) in.readObject(); + var next = (ExprValue) in.readObject(); + var totalHits = in.readLong(); + var inputLoader = (PlanLoader) in.readObject(); + var input = (PhysicalPlan) inputLoader.apply(in, engine); + var fo = new FilterOperator(input, conditions); + fo.next = next; + fo.totalHits = totalHits; + return fo; + }; + out.writeObject(loader); + + out.writeObject(conditions); + out.writeObject(next); + out.writeLong(totalHits); + return input.getPlanForSerialization().writeExternal(out); + } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java index cd84234c4b6..168a7309585 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java @@ -7,6 +7,8 @@ package org.opensearch.sql.planner.physical; import com.google.common.collect.ImmutableList; +import java.io.IOException; +import java.io.ObjectOutput; import java.util.List; import lombok.EqualsAndHashCode; import lombok.Getter; @@ -35,6 +37,14 @@ public class LimitOperator extends PhysicalPlan { private final Integer offset; private Integer count = 0; + // Consider using a copy constructor instead -- see usage + public LimitOperator(PhysicalPlan input, Integer limit, Integer offset, Integer count) { + this.input = input; + this.limit = limit; + this.offset = offset; + this.count = count; + } + @Override public void open() { super.open(); @@ -48,7 +58,9 @@ public void open() { @Override public boolean hasNext() { - return input.hasNext() && count < offset + limit; + var inNext = input.hasNext(); + var cond = count < offset + limit; + return inNext && cond; } @Override @@ -67,4 +79,21 @@ public List getChild() { return ImmutableList.of(input); } + @Override + public boolean writeExternal(ObjectOutput out) throws IOException { + PlanLoader loader = (in, engine) -> { + var limit = in.readInt(); + var offset = in.readInt(); + var count = in.readInt(); + var inputLoader = (PlanLoader) in.readObject(); + var input = (PhysicalPlan) inputLoader.apply(in, engine); + return new LimitOperator(input, limit, offset, count); + }; + out.writeObject(loader); + + out.writeInt(limit); + out.writeInt(offset); + out.writeInt(count); + return input.getPlanForSerialization().writeExternal(out); + } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java index 4463892ca54..458f7ccb0bf 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java @@ -9,6 +9,8 @@ import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; +import java.io.IOException; +import java.io.ObjectOutput; import java.util.Collections; import java.util.Comparator; import java.util.Iterator; @@ -29,7 +31,7 @@ /** * Sort Operator.The input data is sorted by the sort fields in the {@link SortOperator#sortList}. * The sort field is specified by the {@link Expression} with {@link SortOption}. - * The count indicate how many sorted result should been return. + * The count indicate how many sorted result should be return. */ @ToString @EqualsAndHashCode(callSuper = false) @@ -47,7 +49,7 @@ public class SortOperator extends PhysicalPlan { /** * Sort Operator Constructor. * @param input input {@link PhysicalPlan} - * @param sortList list of sort sort field. + * @param sortList list of sort field. * The sort field is specified by the {@link Expression} with {@link SortOption} */ public SortOperator( @@ -134,4 +136,19 @@ public ExprValue next() { } }; } + + @Override + @SuppressWarnings("unchecked") + public boolean writeExternal(ObjectOutput out) throws IOException { + PlanLoader loader = (in, engine) -> { + var sortList = (List>) in.readObject(); + var inputLoader = (PlanLoader) in.readObject(); + var input = (PhysicalPlan) inputLoader.apply(in, engine); + return new SortOperator(input, sortList); + }; + out.writeObject(loader); + + out.writeObject(sortList); + return input.getPlanForSerialization().writeExternal(out); + } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java index 45884830e10..255fcc50eea 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ValuesOperator.java @@ -7,6 +7,9 @@ package org.opensearch.sql.planner.physical; import com.google.common.collect.ImmutableList; + +import java.io.IOException; +import java.io.ObjectOutput; import java.util.Iterator; import java.util.List; import java.util.stream.Collectors; @@ -71,4 +74,9 @@ public ExprValue next() { return new ExprCollectionValue(values); } + @Override + public boolean writeExternal(ObjectOutput out) throws IOException { + // nothing to serialize, return false to get no cursor + return false; + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java index 4d6925f1aaf..425eb37a7d4 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java @@ -44,7 +44,7 @@ public PhysicalPlan protect(PhysicalPlan physicalPlan) { @Override public PhysicalPlan visitFilter(FilterOperator node, Object context) { - return new FilterOperator(visitInput(node.getInput(), context), node.getConditions()); + return new FilterOperator(visitInput(node.getInput(), context), node); } @Override @@ -133,7 +133,8 @@ public PhysicalPlan visitLimit(LimitOperator node, Object context) { return new LimitOperator( visitInput(node.getInput(), context), node.getLimit(), - node.getOffset()); + node.getOffset(), + node.getCount()); } @Override From bd24baedc0fbc83b92ba747810ca6d79506e1d1f Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 22 Mar 2023 12:59:25 -0700 Subject: [PATCH 44/46] Add pushdown. Signed-off-by: Yury-Fridlyand --- .../sql/planner/physical/FilterOperator.java | 1 + .../sql/planner/physical/SortOperator.java | 1 + .../request/InitialPageRequestBuilder.java | 64 ++++++++++++++----- .../request/OpenSearchRequestBuilder.java | 33 +++++----- .../request/PagedRequestBuilder.java | 2 +- .../request/PushDownRequestBuilder.java | 59 +++++++++++++++++ .../scan/OpenSearchIndexScanQueryBuilder.java | 2 +- .../scan/OpenSearchPagedIndexScan.java | 2 + .../scan/OpenSearchPagedIndexScanBuilder.java | 51 +++++++++++++++ .../InitialPageRequestBuilderTest.java | 2 +- .../request/OpenSearchRequestBuilderTest.java | 4 +- .../OpenSearchIndexScanOptimizationTest.java | 2 +- .../storage/scan/OpenSearchIndexScanTest.java | 2 +- 13 files changed, 185 insertions(+), 40 deletions(-) create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java index d35d816f01b..3aa76726283 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/FilterOperator.java @@ -84,6 +84,7 @@ public long getTotalHits() { return totalHits; } + // Probably we don't need this, because we do pushDownFilter in OpenSearchPagedIndexScanBuilder @Override public boolean writeExternal(ObjectOutput out) throws IOException { PlanLoader loader = (in, engine) -> { diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java index 458f7ccb0bf..aa91ce0933c 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/SortOperator.java @@ -137,6 +137,7 @@ public ExprValue next() { }; } + // Probably we don't need this, because we do pushDownSort in OpenSearchPagedIndexScanBuilder @Override @SuppressWarnings("unchecked") public boolean writeExternal(ObjectOutput out) throws IOException { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java index 84bde63a442..defc820f4ff 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java @@ -5,17 +5,23 @@ package org.opensearch.sql.opensearch.request; +import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; +import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Set; import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; +import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.sort.SortBuilder; +import org.opensearch.search.sort.SortBuilders; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprType; @@ -52,38 +58,62 @@ public OpenSearchScrollRequest build() { return new OpenSearchScrollRequest(indexName, sourceBuilder, exprValueFactory); } - public void pushDown(QueryBuilder query) { - throw new UnsupportedOperationException("pushdown of a query is not supported"); - } - - public void pushDownAggregation( - Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { - - throw new UnsupportedOperationException("pagination of aggregation requests is not supported"); + @Override + public void pushDownFilter(QueryBuilder query) { + QueryBuilder current = sourceBuilder.query(); + + if (current == null) { + sourceBuilder.query(query); + } else { + if (isBoolFilterQuery(current)) { + ((BoolQueryBuilder) current).filter(query); + } else { + sourceBuilder.query(QueryBuilders.boolQuery() + .filter(current) + .filter(query)); + } + } + + if (sourceBuilder.sorts() == null) { + sourceBuilder.sort(DOC_FIELD_NAME, ASC); // Make sure consistent order + } } + /** + * Push down sort to DSL request. + * + * @param sortBuilders sortBuilders. + */ + @Override public void pushDownSort(List> sortBuilders) { - throw new UnsupportedOperationException("sorting of paged requests is not supported"); + if (isSortByDocOnly()) { + sourceBuilder.sorts().clear(); + } - } - - public void pushDownLimit(Integer limit, Integer offset) { - throw new UnsupportedOperationException("limit of paged requests is not supported"); - } - - public void pushDownHighlight(String field, Map arguments) { - throw new UnsupportedOperationException("highlight of paged requests is not supported"); + for (SortBuilder sortBuilder : sortBuilders) { + sourceBuilder.sort(sortBuilder); + } } /** * Push down project expression to OpenSearch. */ + @Override public void pushDownProjects(Set projects) { sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr) .distinct().toArray(String[]::new), new String[0]); } + @Override public void pushTypeMapping(Map typeMapping) { exprValueFactory.setTypeMapping(typeMapping); } + + private boolean isSortByDocOnly() { + List> sorts = sourceBuilder.sorts(); + if (sorts != null) { + return sorts.equals(Arrays.asList(SortBuilders.fieldSort(DOC_FIELD_NAME))); + } + return false; + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 4c253a31520..bf4ca7de0f7 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -39,11 +39,10 @@ /** * OpenSearch search request builder. */ -@EqualsAndHashCode +@EqualsAndHashCode(callSuper = false) @Getter @ToString -// TODO make an interface which defines all pushDown functions? -public class OpenSearchRequestBuilder { +public class OpenSearchRequestBuilder extends PushDownRequestBuilder { /** * Default query timeout in minutes. @@ -94,12 +93,12 @@ public OpenSearchRequestBuilder(OpenSearchRequest.IndexName indexName, OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; this.maxResultWindow = maxResultWindow; - this.sourceBuilder = new SearchSourceBuilder(); this.exprValueFactory = exprValueFactory; this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); - sourceBuilder.from(0); - sourceBuilder.size(querySize); - sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); + this.sourceBuilder = new SearchSourceBuilder() + .from(0) + .size(querySize) + .timeout(DEFAULT_QUERY_TIMEOUT); } /** @@ -124,7 +123,8 @@ public OpenSearchRequest build() { * * @param query query request */ - public void pushDown(QueryBuilder query) { + @Override + public void pushDownFilter(QueryBuilder query) { QueryBuilder current = sourceBuilder.query(); if (current == null) { @@ -149,9 +149,10 @@ public void pushDown(QueryBuilder query) { * * @param aggregationBuilder pair of aggregation query and aggregation parser. */ + @Override public void pushDownAggregation( Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { - aggregationBuilder.getLeft().forEach(builder -> sourceBuilder.aggregation(builder)); + aggregationBuilder.getLeft().forEach(sourceBuilder::aggregation); sourceBuilder.size(0); exprValueFactory.setParser(aggregationBuilder.getRight()); } @@ -161,6 +162,7 @@ public void pushDownAggregation( * * @param sortBuilders sortBuilders. */ + @Override public void pushDownSort(List> sortBuilders) { // TODO: Sort by _doc is added when filter push down. Remove both logic once doctest fixed. if (isSortByDocOnly()) { @@ -175,6 +177,7 @@ public void pushDownSort(List> sortBuilders) { /** * Push down size (limit) and from (offset) to DSL request. */ + @Override public void pushDownLimit(Integer limit, Integer offset) { querySize = limit; sourceBuilder.from(offset).size(limit); @@ -184,6 +187,7 @@ public void pushDownLimit(Integer limit, Integer offset) { * Add highlight to DSL requests. * @param field name of the field to highlight */ + @Override public void pushDownHighlight(String field, Map arguments) { String unquotedField = StringUtils.unquoteText(field); if (sourceBuilder.highlighter() != null) { @@ -216,20 +220,17 @@ public void pushDownHighlight(String field, Map arguments) { /** * Push down project list to DSL requets. */ + @Override public void pushDownProjects(Set projects) { - final Set projectsSet = - projects.stream().map(ReferenceExpression::getAttr).collect(Collectors.toSet()); - sourceBuilder.fetchSource(projectsSet.toArray(new String[0]), new String[0]); + sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr) + .distinct().toArray(String[]::new), new String[0]); } + @Override public void pushTypeMapping(Map typeMapping) { exprValueFactory.setTypeMapping(typeMapping); } - private boolean isBoolFilterQuery(QueryBuilder current) { - return (current instanceof BoolQueryBuilder); - } - private boolean isSortByDocOnly() { List> sorts = sourceBuilder.sorts(); if (sorts != null) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java index 6afadeda5c8..0e2935934d1 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java @@ -8,7 +8,7 @@ import lombok.Getter; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -public abstract class PagedRequestBuilder { +public abstract class PagedRequestBuilder extends PushDownRequestBuilder { @Getter protected OpenSearchExprValueFactory exprValueFactory; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java new file mode 100644 index 00000000000..3775fd56bb0 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java @@ -0,0 +1,59 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.request; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.search.aggregations.AggregationBuilder; +import org.opensearch.search.sort.SortBuilder; +import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; + +public abstract class PushDownRequestBuilder { + protected boolean isBoolFilterQuery(QueryBuilder current) { + return (current instanceof BoolQueryBuilder); + } + + private void throwUnsupported(String operation) { + throw new UnsupportedOperationException(String.format("%s: %s in requests is not supported", + getClass().getSimpleName(), operation)); + } + + public void pushDownFilter(QueryBuilder query) { + throwUnsupported("filter"); + } + + public void pushDownAggregation( + Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { + throwUnsupported("aggregation"); + } + + public void pushDownSort(List> sortBuilders) { + throwUnsupported("sorting"); + } + + public void pushDownLimit(Integer limit, Integer offset) { + throwUnsupported("limit/offset"); + } + + public void pushDownHighlight(String field, Map arguments) { + throwUnsupported("highlight"); + } + + public void pushDownProjects(Set projects) { + throwUnsupported("select list"); + } + + public void pushTypeMapping(Map typeMapping) { + throwUnsupported("type mapping"); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index 5cfde4abbee..f2e5139d01d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -61,7 +61,7 @@ public boolean pushDownFilter(LogicalFilter filter) { FilterQueryBuilder queryBuilder = new FilterQueryBuilder( new DefaultExpressionSerializer()); QueryBuilder query = queryBuilder.build(filter.getCondition()); - indexScan.getRequestBuilder().pushDown(query); + indexScan.getRequestBuilder().pushDownFilter(query); return true; } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java index 01217d9e57c..999b93447d2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java @@ -11,6 +11,7 @@ import java.util.Iterator; import lombok.EqualsAndHashCode; +import lombok.Getter; import lombok.ToString; import org.apache.commons.lang3.NotImplementedException; import org.opensearch.sql.data.model.ExprValue; @@ -24,6 +25,7 @@ @ToString(onlyExplicitlyIncluded = true) public class OpenSearchPagedIndexScan extends TableScanOperator { private OpenSearchClient client; + @Getter private PagedRequestBuilder requestBuilder; @EqualsAndHashCode.Include @ToString.Include diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java index 779df4ebec9..40fa17ca8aa 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java @@ -6,9 +6,25 @@ package org.opensearch.sql.opensearch.storage.scan; import lombok.EqualsAndHashCode; +import org.apache.commons.lang3.tuple.Pair; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.expression.Expression; +import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; +import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; +import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalProject; +import org.opensearch.sql.planner.logical.LogicalSort; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; +import java.util.List; +import java.util.stream.Collectors; + +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanQueryBuilder.findReferenceExpressions; + /** * Builder for a paged OpenSearch request. * Override pushDown* methods from TableScanBuilder as more features @@ -22,6 +38,41 @@ public OpenSearchPagedIndexScanBuilder(OpenSearchPagedIndexScan indexScan) { this.indexScan = indexScan; } + @Override + public boolean pushDownFilter(LogicalFilter filter) { + FilterQueryBuilder queryBuilder = new FilterQueryBuilder( + new DefaultExpressionSerializer()); + QueryBuilder query = queryBuilder.build(filter.getCondition()); + indexScan.getRequestBuilder().pushDownFilter(query); + return true; + } + + @Override + public boolean pushDownSort(LogicalSort sort) { + List> sortList = sort.getSortList(); + final SortQueryBuilder builder = new SortQueryBuilder(); + indexScan.getRequestBuilder().pushDownSort(sortList.stream() + .map(sortItem -> builder.build(sortItem.getValue(), sortItem.getKey())) + .collect(Collectors.toList())); + return true; + } + + // How can we set limit? Regular request sets + // .size(limit) + // Paged request sets + // .size(pageSize) + @Override + public boolean pushDownLimit(LogicalLimit limit) { + return false;//super.pushDownLimit(limit); + } + + @Override + public boolean pushDownProject(LogicalProject project) { + indexScan.getRequestBuilder().pushDownProjects( + findReferenceExpressions(project.getProjectList())); + return false; + } + @Override public TableScanOperator build() { return indexScan; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java index 023a1e397a1..0ad6c27469f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java @@ -67,7 +67,7 @@ public void build() { @Test public void pushDown_not_supported() { assertAll( - () -> assertThrows(Throwable.class, () -> requestBuilder.pushDown(mock())), + () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownFilter(mock())), () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownAggregation(mock())), () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownSort(mock())), () -> assertThrows(Throwable.class, () -> requestBuilder.pushDownLimit(1, 2)), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index e4f557d865b..50fee450fb7 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -106,7 +106,7 @@ void build_scroll_request_with_correct_size() { @Test void test_push_down_query() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); - requestBuilder.pushDown(query); + requestBuilder.pushDownFilter(query); assertEquals( new SearchSourceBuilder() @@ -143,7 +143,7 @@ void test_push_down_aggregation() { @Test void test_push_down_query_and_sort() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); - requestBuilder.pushDown(query); + requestBuilder.pushDownFilter(query); FieldSortBuilder sortBuilder = SortBuilders.fieldSort("intA"); requestBuilder.pushDownSort(List.of(sortBuilder)); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index b5de6e30c5b..043049aefe2 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -524,7 +524,7 @@ private void assertEqualsAfterOptimization(LogicalPlan expected, LogicalPlan act } private Runnable withFilterPushedDown(QueryBuilder filteringCondition) { - return () -> verify(requestBuilder, times(1)).pushDown(filteringCondition); + return () -> verify(requestBuilder, times(1)).pushDownFilter(filteringCondition); } private Runnable withAggregationPushedDown( diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 5b90451cf8f..03c9d95b774 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -311,7 +311,7 @@ public PushDownAssertion(OpenSearchClient client, } PushDownAssertion pushDown(QueryBuilder query) { - indexScan.getRequestBuilder().pushDown(query); + indexScan.getRequestBuilder().pushDownFilter(query); return this; } From 669153a7c2abe5936bd674bf53791d68ce5d932a Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Wed, 22 Mar 2023 13:00:39 -0700 Subject: [PATCH 45/46] Remove some debug stuff. Signed-off-by: Yury-Fridlyand --- .../org/opensearch/sql/planner/physical/LimitOperator.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java index 168a7309585..227419e7b18 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/LimitOperator.java @@ -58,9 +58,7 @@ public void open() { @Override public boolean hasNext() { - var inNext = input.hasNext(); - var cond = count < offset + limit; - return inNext && cond; + return input.hasNext() && count < offset + limit; } @Override From 163093cb76cba25d3acb749195a242f20d2ff887 Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Thu, 23 Mar 2023 16:30:17 -0700 Subject: [PATCH 46/46] Support more features in pagination. Signed-off-by: Yury-Fridlyand --- .../sql/executor/CanPaginateVisitor.java | 96 +++++++++++++++++++ .../request/InitialPageRequestBuilder.java | 28 ++++++ .../scan/OpenSearchPagedIndexScanBuilder.java | 10 ++ 3 files changed, 134 insertions(+) diff --git a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java index 05ba08baf50..53e283aa9e9 100644 --- a/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java +++ b/core/src/main/java/org/opensearch/sql/executor/CanPaginateVisitor.java @@ -10,10 +10,26 @@ import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.Alias; import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.expression.And; +import org.opensearch.sql.ast.expression.Between; +import org.opensearch.sql.ast.expression.Case; +import org.opensearch.sql.ast.expression.Cast; +import org.opensearch.sql.ast.expression.Compare; +import org.opensearch.sql.ast.expression.EqualTo; import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.Function; +import org.opensearch.sql.ast.expression.HighlightFunction; +import org.opensearch.sql.ast.expression.In; +import org.opensearch.sql.ast.expression.Interval; import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.ast.expression.Not; +import org.opensearch.sql.ast.expression.Or; import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.RelevanceFieldList; +import org.opensearch.sql.ast.expression.UnresolvedAttribute; +import org.opensearch.sql.ast.expression.When; import org.opensearch.sql.ast.expression.WindowFunction; +import org.opensearch.sql.ast.expression.Xor; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Limit; import org.opensearch.sql.ast.tree.Project; @@ -115,6 +131,86 @@ public Boolean visitChildren(Node node, Object context) { return Boolean.FALSE; } + @Override + public Boolean visitEqualTo(EqualTo node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitRelevanceFieldList(RelevanceFieldList node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitInterval(Interval node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitCompare(Compare node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitNot(Not node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitOr(Or node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitAnd(And node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitXor(Xor node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitFunction(Function node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitIn(In node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitBetween(Between node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitCase(Case node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitWhen(When node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitCast(Cast node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitHighlightFunction(HighlightFunction node, Object context) { + return canPaginate(node, context); + } + + @Override + public Boolean visitUnresolvedAttribute(UnresolvedAttribute node, Object context) { + return canPaginate(node, context); + } + @Override public Boolean visit(Node node, Object context) { // for all not listed (= unchecked) - false diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java index defc820f4ff..d91b5f4b0d6 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java @@ -20,11 +20,14 @@ import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder; import org.opensearch.search.sort.SortBuilder; import org.opensearch.search.sort.SortBuilders; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; @@ -95,6 +98,31 @@ public void pushDownSort(List> sortBuilders) { } } + @Override + public void pushDownHighlight(String field, Map arguments) { + HighlightBuilder highlightBuilder; + String unquotedField = StringUtils.unquoteText(field); + if (sourceBuilder.highlighter() != null) { + // OS does not allow duplicates of highlight fields + if (sourceBuilder.highlighter().fields().stream() + .anyMatch(f -> f.name().equals(unquotedField))) { + throw new SemanticCheckException(String.format( + "Duplicate field %s in highlight", field)); + } + highlightBuilder = sourceBuilder.highlighter().field(unquotedField); + } else { + highlightBuilder = new HighlightBuilder().field(unquotedField); + sourceBuilder.highlighter(highlightBuilder); + } + + if (arguments.containsKey("pre_tags")) { + highlightBuilder.preTags(arguments.get("pre_tags").toString()); + } + if (arguments.containsKey("post_tags")) { + highlightBuilder.postTags(arguments.get("post_tags").toString()); + } + } + /** * Push down project expression to OpenSearch. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java index 40fa17ca8aa..d15de343436 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java @@ -9,11 +9,13 @@ import org.apache.commons.lang3.tuple.Pair; import org.opensearch.index.query.QueryBuilder; import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalSort; @@ -73,6 +75,14 @@ public boolean pushDownProject(LogicalProject project) { return false; } + @Override + public boolean pushDownHighlight(LogicalHighlight highlight) { + indexScan.getRequestBuilder().pushDownHighlight( + StringUtils.unquoteText(highlight.getHighlightField().toString()), + highlight.getArguments()); + return true; + } + @Override public TableScanOperator build() { return indexScan;