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 d61747d0eb..aa86b61767 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 881046ec0b..c3a84ea286 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 c8c95a0ae6..6822eca727 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 72225f0884..08718cb4d6 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 2acc44de63..7013842066 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());