diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java index 169ef019538b0..7034b39b12069 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java @@ -90,7 +90,6 @@ public boolean test(FeedResponse frp) { private volatile int collectedItems = 0; private volatile boolean lastPage = false; - @Override public FeedResponse apply(FeedResponse t) { @@ -101,9 +100,15 @@ public FeedResponse apply(FeedResponse 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 index dd73fb4d93d6b..4e315e629f460 100644 --- 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 receivedDocuments = this.queryWithContinuationTokens(query, pageSize);