escapeNonAscIIPkValueForQueryPlanAndQuery

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDiagnosticsTest.java

@@ -555,6 +555,7 @@ public void requestSessionTokenDiagnostics() {
     @Test(groups = {"fast"})
     public void databaseAccountToClients() {
         CosmosClient testClient = null;
+        CosmosClient testClient2 = null;
         try {
             testClient = new CosmosClientBuilder()
                 .endpoint(TestConfigurations.HOST)
@@ -588,8 +589,7 @@ public void databaseAccountToClients() {
             String intString = substrings[substrings.length-1];
             int intValue = Integer.parseInt(intString);
 
-
-            CosmosClient testClient2 = new CosmosClientBuilder()
+            testClient2 = new CosmosClientBuilder()
                 .endpoint(TestConfigurations.HOST)
                 .key(TestConfigurations.MASTER_KEY)
                 .contentResponseOnWriteEnabled(true)
@@ -615,13 +615,9 @@ public void databaseAccountToClients() {
             intString = substrings[substrings.length-1];
             assertThat(Integer.parseInt(intString)).isEqualTo(intValue+1);
 
-            //close second client
-            testClient2.close();
-
         } finally {
-            if (testClient != null) {
-                testClient.close();
-            }
+            safeCloseSyncClient(testClient);
+            safeCloseSyncClient(testClient2);
         }
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -504,7 +504,14 @@ private RxDocumentClientImpl(URI serviceEndpoint,
         this.clientId = clientIdGenerator.incrementAndGet();
         this.clientCorrelationId = Strings.isNullOrWhiteSpace(clientCorrelationId) ?
             String.format("%05d",this.clientId): clientCorrelationId;
-        clientMap.put(serviceEndpoint.toString(), clientMap.getOrDefault(serviceEndpoint.toString(), 0) + 1);
+        clientMap.compute(serviceEndpoint.toString(), (key, value) -> {
+            if (value == null) {
+                return 1;
+            }
+
+            return value++;
+        });
+
         this.diagnosticsClientConfig = new DiagnosticsClientConfig();
         this.diagnosticsClientConfig.withClientId(this.clientId);
         this.diagnosticsClientConfig.withActiveClientCounter(activeClientsCnt);

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java

@@ -742,7 +742,7 @@ public static CosmosChangeFeedRequestOptions getEffectiveCosmosChangeFeedRequest
                 cosmosChangeFeedRequestRequestOptions, pagedFluxOptions);
     }
 
-    static String escapeNonAscii(String partitionKeyJson) {
+    public static String escapeNonAscii(String partitionKeyJson) {
         // if all are ascii original string will be returned, and avoids copying data.
         StringBuilder sb = null;
         for (int i = 0; i < partitionKeyJson.length(); i++) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/feedranges/FeedRangePartitionKeyImpl.java

@@ -156,7 +156,7 @@ public Mono<RxDocumentServiceRequest> populateFeedRangeFilteringHeaders(
 
         request.getHeaders().put(
             HttpConstants.HttpHeaders.PARTITION_KEY,
-            this.partitionKey.toJson());
+            Utils.escapeNonAscii(this.partitionKey.toJson()));
         request.setPartitionKeyInternal(this.partitionKey);
 
         MetadataDiagnosticsContext metadataDiagnosticsCtx =

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java

@@ -330,7 +330,7 @@ private void populatePartitionKeyInfo(RxDocumentServiceRequest request, Partitio
             if (partitionKey != null) {
                 request.setPartitionKeyInternal(partitionKey);
                 request.setPartitionKeyDefinition(partitionKeyDefinition);
-                request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKey.toJson());
+                request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKey.toJson()));
             }
         }
     }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java

@@ -13,6 +13,7 @@
 import com.azure.cosmos.implementation.ResourceType;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.Strings;
+import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -83,7 +84,7 @@ protected void initialize(
                     // partitionKeyInternal gets passed as null which avoids the feedRange normalization.
                     if (!PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey())) {
                         partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey());
-                        headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
+                        headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
                     }
                 }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java

@@ -96,7 +96,7 @@ static Mono<PartitionedQueryExecutionInfo> getQueryPlanThroughGatewayAsync(Diagn
 
         if (partitionKey != null && partitionKey != PartitionKey.NONE) {
             PartitionKeyInternal partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(partitionKey);
-            requestHeaders.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
+            requestHeaders.put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
         }
 
         final RxDocumentServiceRequest queryPlanRequest = RxDocumentServiceRequest.create(diagnosticsClientContext,