Fix v2 paginationIT

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java

@@ -69,7 +69,10 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
   /** List of includes expected in the response. */
   @EqualsAndHashCode.Exclude @ToString.Exclude private final List<String> includes;
 
-  @EqualsAndHashCode.Exclude @ToString.Exclude private boolean needClean = false;
+  @EqualsAndHashCode.Exclude @ToString.Exclude private boolean needClean = true;
+
+  /** Indicate the search already done. */
+  @EqualsAndHashCode.Exclude @ToString.Exclude private boolean searchDone = false;
 
   private String pitId;
 
@@ -145,6 +148,14 @@ public static OpenSearchQueryRequest pitOf(
     this.calciteEnabled = calciteEnabled;
   }
 
+  /** true if the request is a count aggregation request. */
+  public boolean isCountAggRequest() {
+    return !searchDone
+        && sourceBuilder.size() == 0
+        && sourceBuilder.trackTotalHitsUpTo() != null // only set in v3
+        && sourceBuilder.trackTotalHitsUpTo() == Integer.MAX_VALUE;
+  }
+
   /**
    * Constructs OpenSearchQueryRequest from serialized representation.
    *
@@ -194,25 +205,17 @@ public OpenSearchResponse search(
     }
   }
 
-  /** true if the request is a count aggregation request. */
-  public boolean isCountAggRequest() {
-    return !needClean
-        && sourceBuilder.size() == 0
-        && sourceBuilder.trackTotalHitsUpTo() != null // only set in v3
-        && sourceBuilder.trackTotalHitsUpTo() == Integer.MAX_VALUE;
-  }
-
   /** Call the old search logic for v2, since we don't support paginating aggregate in v2. */
   @Deprecated
   private OpenSearchResponse searchForV2(Function<SearchRequest, SearchResponse> searchAction) {
     // When SearchRequest doesn't contain PitId, fetch single page request
-    if (needClean) {
+    if (searchDone) {
       return new OpenSearchResponse(
           SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
     } else {
-      // get the value before set needClean = true
+      // get the value before set searchDone = true
       boolean isCountAggRequest = isCountAggRequest();
-      needClean = true;
+      searchDone = true;
       return new OpenSearchResponse(
           searchAction.apply(
               new SearchRequest().indices(indexName.getIndexNames()).source(sourceBuilder)),
@@ -224,7 +227,7 @@ private OpenSearchResponse searchForV2(Function<SearchRequest, SearchResponse> s
 
   private OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction) {
     OpenSearchResponse openSearchResponse;
-    if (needClean) {
+    if (searchDone) {
       openSearchResponse =
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
@@ -248,6 +251,7 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
               this.searchResponse, exprValueFactory, includes, isCountAggRequest());
 
       needClean = openSearchResponse.isEmpty();
+      searchDone = openSearchResponse.isEmpty();
       // Get afterKey from response
       if (openSearchResponse.isAggregationResponse()) {
         openSearchResponse.getAggregations().asList().stream()
@@ -260,7 +264,7 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
 
   public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse> searchAction) {
     OpenSearchResponse openSearchResponse;
-    if (needClean) {
+    if (searchDone) {
       openSearchResponse =
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
@@ -298,6 +302,7 @@ public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse>
               this.searchResponse, exprValueFactory, includes, isCountAggRequest());
 
       needClean = openSearchResponse.isEmpty();
+      searchDone = openSearchResponse.isEmpty();
       SearchHit[] searchHits = this.searchResponse.getHits().getHits();
       if (searchHits != null && searchHits.length > 0) {
         searchAfter = searchHits[searchHits.length - 1].getSortValues();
@@ -316,6 +321,7 @@ public void clean(Consumer<String> cleanAction) {
       // clean on the last page only, to prevent deleting the PitId in the middle of paging.
       if (this.pitId != null && needClean) {
         cleanAction.accept(this.pitId);
+        searchDone = true;
       }
     } finally {
       this.pitId = null;
@@ -329,7 +335,7 @@ public void forceClean(Consumer<String> cleanAction) {
     try {
       if (this.pitId != null) {
         cleanAction.accept(this.pitId);
-        needClean = true;
+        searchDone = true;
       }
     } finally {
       this.pitId = null;

opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java

@@ -275,7 +275,7 @@ void has_another_batch() {
             new TimeValue(1000),
             "sample",
             false);
-    assertTrue(request.hasAnotherBatch());
+    assertFalse(request.hasAnotherBatch());
   }
 
   @Test
@@ -366,7 +366,7 @@ void testCleanConditionTrue() {
     request.clean(cleanAction);
 
     verify(cleanAction, times(1)).accept("samplePid");
-    assertTrue(openSearchResponse.isEmpty());
+    assertTrue(request.isSearchDone());
     assertNull(request.getPitId());
   }
 
@@ -389,7 +389,7 @@ void testCleanConditionFalse_needCleanFalse() {
 
     request.clean(cleanAction);
     verify(cleanAction, never()).accept(anyString());
-    assertFalse(openSearchResponse.isEmpty());
+    assertFalse(request.isSearchDone());
     assertNull(request.getPitId());
   }
 
@@ -407,7 +407,7 @@ void testCleanConditionFalse_pidNull() {
 
     request.clean(cleanAction);
     verify(cleanAction, never()).accept(anyString());
-    assertFalse(request.isNeedClean());
+    assertFalse(request.isSearchDone());
     assertNull(request.getPitId());
   }