Fixing the top query continuation issue (#27377)

Co-authored-by: Aayush Kataria <[email protected]>
Azure · Mar 4, 2022 · 81e6f0f · 81e6f0f
1 parent 236dd53
commit 81e6f0f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/...src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java b/...src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java
@@ -90,7 +90,6 @@ public boolean test(FeedResponse<T> frp) {
 
             private volatile int collectedItems = 0;
             private volatile boolean lastPage = false;
-
             @Override
             public FeedResponse<T> apply(FeedResponse<T> t) {
 
@@ -101,9 +100,15 @@ public FeedResponse<T> apply(FeedResponse<T> t) {
                     if (top != collectedItems) {
                         // Add Take Continuation Token
                         String sourceContinuationToken = t.getContinuationToken();
-                        TakeContinuationToken takeContinuationToken = new TakeContinuationToken(top - collectedItems,
+                        if (sourceContinuationToken != null) {
+                            TakeContinuationToken takeContinuationToken = new TakeContinuationToken(top - collectedItems,
                                 sourceContinuationToken);
-                        headers.put(HttpConstants.HttpHeaders.CONTINUATION, takeContinuationToken.toJson());
+                            headers.put(HttpConstants.HttpHeaders.CONTINUATION, takeContinuationToken.toJson());
+                        } else {
+                            // Null out the continuation token. The sourceContinuationToken being null means
+                            // that this is the last page and there are no more elements left to fetch.
+                            headers.put(HttpConstants.HttpHeaders.CONTINUATION, null);
+                        }
                     } else {
                         // Null out the continuation token
                         headers.put(HttpConstants.HttpHeaders.CONTINUATION, null);

diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java
@@ -127,6 +127,12 @@ public void queryDocumentsWithTopContinuationTokens() throws Exception {
         this.queryWithContinuationTokensAndPageSizes(query, new int[] { 1, 5, 10 }, 8);
     }
 
+    @Test(groups = { "simple" }, timeOut = TIMEOUT * 1000, retryAnalyzer = RetryAnalyzer.class)
+    public void queryDocumentsWithTopGreaterThanItemsContinuationTokens() throws Exception {
+        String query = "SELECT TOP 2147483647 * FROM c";
+        this.queryWithContinuationTokensAndPageSizes(query, new int[] {1}, 20);
+    }
+
     private void queryWithContinuationTokensAndPageSizes(String query, int[] pageSizes, int topCount) {
         for (int pageSize : pageSizes) {
             List<InternalObjectNode> receivedDocuments = this.queryWithContinuationTokens(query, pageSize);