Fix ITs

Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

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

@@ -33,6 +33,7 @@
 import org.opensearch.sql.analysis.AnalysisContext;
 import org.opensearch.sql.analysis.Analyzer;
 import org.opensearch.sql.ast.statement.Explain;
+import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
@@ -95,7 +96,12 @@ public void execute(
     if (shouldUseCalcite(queryType)) {
       executeWithCalcite(plan, queryType, listener, pageSize, offset);
     } else {
-      executeWithLegacy(plan, queryType, listener, Optional.empty());
+      // For legacy path (SQL), wrap with Paginate when pagination is enabled
+      if (pageSize > 0) {
+        executeWithLegacy(new Paginate(pageSize, plan), queryType, listener, Optional.empty());
+      } else {
+        executeWithLegacy(plan, queryType, listener, Optional.empty());
+      }
     }
   }
 

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java

@@ -17,6 +17,7 @@
 import org.opensearch.sql.ast.tree.FetchCursor;
 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.QueryId;
 import org.opensearch.sql.executor.QueryService;
@@ -106,15 +107,31 @@ public AbstractPlan visitQuery(
           context) {
     requireNonNull(context.getLeft(), "[BUG] query listener must be not null");
     if (node.getFetchSize() > 0) {
-      // Pagination enabled - use pageSize and offset
-      return new QueryPlan(
-          QueryId.queryId(),
-          node.getQueryType(),
-          node.getPlan(),
-          node.getFetchSize(),
-          node.getPaginationOffset(),
-          queryService,
-          context.getLeft());
+      if (node.getQueryType() == QueryType.PPL) {
+        // PPL pagination - use pageSize and offset (Calcite path)
+        return new QueryPlan(
+            QueryId.queryId(),
+            node.getQueryType(),
+            node.getPlan(),
+            node.getFetchSize(),
+            node.getPaginationOffset(),
+            queryService,
+            context.getLeft());
+      } else {
+        // SQL pagination - use legacy cursor-based pagination
+        if (canConvertToCursor(node.getPlan())) {
+          return new QueryPlan(
+              QueryId.queryId(),
+              node.getQueryType(),
+              node.getPlan(),
+              node.getFetchSize(),
+              queryService,
+              context.getLeft());
+        } else {
+          // This should be picked up by the legacy engine.
+          throw new UnsupportedCursorRequestException();
+        }
+      }
     } else {
       return new QueryPlan(
           QueryId.queryId(), node.getQueryType(), node.getPlan(), queryService, context.getLeft());

core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java

@@ -99,15 +99,17 @@ public void no_consumer_response_channel() {
   @Test
   public void create_query_with_fetch_size_should_create_query_plan() {
     factory = new QueryPlanFactory(queryService);
-    Statement query = new Query(plan, 10, queryType);
+    // Use PPL query type for new pagination feature
+    Statement query = new Query(plan, 10, QueryType.PPL);
     AbstractPlan queryExecution = factory.create(query, queryListener, explainListener);
     assertTrue(queryExecution instanceof QueryPlan);
   }
 
   @Test
   public void create_query_with_fetch_size_and_offset_should_create_query_plan() {
     factory = new QueryPlanFactory(queryService);
-    Statement query = new Query(plan, 10, queryType, 50);
+    // Use PPL query type for new pagination with offset feature
+    Statement query = new Query(plan, 10, QueryType.PPL, 50);
     AbstractPlan queryExecution = factory.create(query, queryListener, explainListener);
     assertTrue(queryExecution instanceof QueryPlan);
   }