ResourceTokenFix-ReadAllItems (#19662)

* ResourceTokenFix-ReadAllItems Co-authored-by: annie-mac <annie-mac@annie-macs-MacBook-Pro.local>
Azure · Mar 6, 2021 · 0eb4f70 · 0eb4f70
1 parent 536abd6
commit 0eb4f70
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 20 deletions.
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
@@ -874,8 +874,10 @@ public <T> Mono<FeedResponse<T>> readMany(
     public <T> CosmosPagedFlux<T> readAllItems(
         PartitionKey partitionKey,
         Class<T> classType) {
+        CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions();
+        queryRequestOptions.setPartitionKey(partitionKey);
 
-        return this.readAllItems(partitionKey, new CosmosQueryRequestOptions(), classType);
+        return this.readAllItems(partitionKey, queryRequestOptions, classType);
     }
 
     /**

diff --git a/...smos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/...smos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
@@ -2306,7 +2306,7 @@ public Flux<FeedResponse<Document>> readAllDocuments(
                         return createQueryInternal(
                             resourceLink,
                             querySpec,
-                            ModelBridgeInternal.partitionKeyRangeIdInternal(effectiveOptions, range.getId()),
+                            ModelBridgeInternal.setPartitionKeyRangeIdInternal(effectiveOptions, range.getId()),
                             Document.class,
                             ResourceType.Document,
                             queryClient,

diff --git a/...main/java/com/azure/cosmos/implementation/query/DefaultDocumentQueryExecutionContext.java b/...main/java/com/azure/cosmos/implementation/query/DefaultDocumentQueryExecutionContext.java
@@ -3,27 +3,25 @@
 package com.azure.cosmos.implementation.query;
 
 import com.azure.cosmos.BridgeInternal;
-import com.azure.cosmos.implementation.DiagnosticsClientContext;
-import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
-import com.azure.cosmos.models.CosmosQueryRequestOptions;
-import com.azure.cosmos.models.FeedResponse;
-import com.azure.cosmos.models.ModelBridgeInternal;
-import com.azure.cosmos.implementation.Resource;
-import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.implementation.BackoffRetryUtility;
 import com.azure.cosmos.implementation.Constants;
+import com.azure.cosmos.implementation.DiagnosticsClientContext;
 import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.InvalidPartitionExceptionRetryPolicy;
 import com.azure.cosmos.implementation.PartitionKeyRange;
 import com.azure.cosmos.implementation.PartitionKeyRangeGoneRetryPolicy;
 import com.azure.cosmos.implementation.PathsHelper;
 import com.azure.cosmos.implementation.QueryMetrics;
+import com.azure.cosmos.implementation.Resource;
 import com.azure.cosmos.implementation.ResourceType;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.Utils.ValueHolder;
+import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
+import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
 import com.azure.cosmos.implementation.caches.IPartitionKeyRangeCache;
 import com.azure.cosmos.implementation.caches.RxCollectionCache;
+import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
 import com.azure.cosmos.implementation.query.metrics.ClientSideMetrics;
 import com.azure.cosmos.implementation.query.metrics.FetchExecutionRangeAccumulator;
 import com.azure.cosmos.implementation.query.metrics.SchedulingStopwatch;
@@ -32,8 +30,10 @@
 import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity;
 import com.azure.cosmos.implementation.routing.Range;
 import com.azure.cosmos.implementation.routing.RoutingMapProviderHelper;
-import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
-import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
+import com.azure.cosmos.models.CosmosQueryRequestOptions;
+import com.azure.cosmos.models.FeedResponse;
+import com.azure.cosmos.models.ModelBridgeInternal;
+import com.azure.cosmos.models.SqlQuerySpec;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -45,7 +45,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
-import static com.azure.cosmos.models.ModelBridgeInternal.partitionKeyRangeIdInternal;
+import static com.azure.cosmos.models.ModelBridgeInternal.getPartitionKeyRangeIdInternal;
 
 /**
  * While this class is public, but it is not part of our published public APIs.
@@ -203,8 +203,8 @@ public RxDocumentServiceRequest createRequestAsync(String continuationToken, Int
                 this.query,
                 this.getPartitionKeyInternal());
 
-        if (!StringUtils.isEmpty(partitionKeyRangeIdInternal(cosmosQueryRequestOptions))) {
-            request.routeTo(new PartitionKeyRangeIdentity(partitionKeyRangeIdInternal(cosmosQueryRequestOptions)));
+        if (!StringUtils.isEmpty(getPartitionKeyRangeIdInternal(cosmosQueryRequestOptions))) {
+            request.routeTo(new PartitionKeyRangeIdentity(getPartitionKeyRangeIdInternal(cosmosQueryRequestOptions)));
         }
 
         return request;

diff --git a/...main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/...main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java
@@ -77,12 +77,12 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
         // FeedOptions, but have the below condition
         // for handling ParallelDocumentQueryTest#partitionKeyRangeId
         if (cosmosQueryRequestOptions != null &&
-            !StringUtils.isEmpty(ModelBridgeInternal.partitionKeyRangeIdInternal(cosmosQueryRequestOptions))) {
+            !StringUtils.isEmpty(ModelBridgeInternal.getPartitionKeyRangeIdInternal(cosmosQueryRequestOptions))) {
 
             Mono<List<PartitionKeyRange>> partitionKeyRanges = queryExecutionContext
                 .getTargetPartitionKeyRangesById(
                     collection.getResourceId(),
-                    ModelBridgeInternal.partitionKeyRangeIdInternal(cosmosQueryRequestOptions));
+                    ModelBridgeInternal.getPartitionKeyRangeIdInternal(cosmosQueryRequestOptions));
 
             return partitionKeyRanges.map(pkRanges -> {
                 List<Range<String>> ranges =

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -310,7 +310,7 @@ public static DatabaseAccount toDatabaseAccount(RxDocumentServiceResponse respon
      * @return the partitionKeyRangeId.
      */
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static String partitionKeyRangeIdInternal(CosmosQueryRequestOptions options) {
+    public static String getPartitionKeyRangeIdInternal(CosmosQueryRequestOptions options) {
         return options.getPartitionKeyRangeIdInternal();
     }
 
@@ -322,7 +322,7 @@ public static String partitionKeyRangeIdInternal(CosmosQueryRequestOptions optio
      * @return the partitionKeyRangeId.
      */
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static CosmosQueryRequestOptions partitionKeyRangeIdInternal(CosmosQueryRequestOptions options, String partitionKeyRangeId) {
+    public static CosmosQueryRequestOptions setPartitionKeyRangeIdInternal(CosmosQueryRequestOptions options, String partitionKeyRangeId) {
         return options.setPartitionKeyRangeIdInternal(partitionKeyRangeId);
     }
 

diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/ResourceTokenTestForV4.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/ResourceTokenTestForV4.java
@@ -416,6 +416,36 @@ public void queryItemFromResourceToken(
         }
     }
 
+    @Test(groups = {"simple"}, timeOut = TIMEOUT)
+    public void readAllItemFromResourceToken() throws Exception {
+
+        CosmosAsyncClient asyncClientResourceToken = null;
+        try {
+            asyncClientResourceToken = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .gatewayMode()
+                .consistencyLevel(ConsistencyLevel.SESSION)
+                .resourceToken(createdContainerPermissionWithPartitionKey.getToken())
+                .buildAsyncClient();
+
+            CosmosAsyncContainer asyncContainer =
+                asyncClientResourceToken
+                    .getDatabase(createdDatabase.getId())
+                    .getContainer(createdContainerWithPartitionKey.getId());
+
+            CosmosPagedFlux<TestObject> readAllItemObservable =
+                asyncContainer.readAllItems(new PartitionKey(PARTITION_KEY_VALUE), TestObject.class);
+            FeedResponseListValidator<TestObject> validator = new FeedResponseListValidator.Builder<TestObject>()
+                .totalSize(1)
+                .numberOfPagesIsGreaterThanOrEqualTo(1)
+                .build();
+
+            validateQuerySuccess(readAllItemObservable.byPage(100), validator);
+        } finally {
+            safeClose(asyncClientResourceToken);
+        }
+    }
+
     private CosmosAsyncClient createAsyncClientWithPermission(List<CosmosPermissionProperties> permissions) {
         assertThat(permissions).isNotNull();
 

diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ParallelDocumentQueryTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ParallelDocumentQueryTest.java
@@ -50,7 +50,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
-import static com.azure.cosmos.models.ModelBridgeInternal.partitionKeyRangeIdInternal;
+import static com.azure.cosmos.models.ModelBridgeInternal.setPartitionKeyRangeIdInternal;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.tuple;
 import static org.testng.Assert.fail;
@@ -228,7 +228,7 @@ public void partitionKeyRangeId() {
                 .map(Resource::getId).collectList().single().block()) {
             String query = "SELECT * from root";
             CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
-            partitionKeyRangeIdInternal(options, partitionKeyRangeId);
+            setPartitionKeyRangeIdInternal(options, partitionKeyRangeId);
             int queryResultCount = createdCollection.queryItems(query, options, InternalObjectNode.class)
                 .byPage()
                 .flatMap(p -> Flux.fromIterable(p.getResults()))