diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/HttpConstants.java index 6afcbdf58cba..3afcf3769bf2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/HttpConstants.java @@ -125,6 +125,7 @@ public static class HttpHeaders { // Address related headers. public static final String FORCE_REFRESH = "x-ms-force-refresh"; + public static final String FORCE_COLLECTION_ROUTING_MAP_REFRESH = "x-ms-collectionroutingmap-refresh"; public static final String ITEM_COUNT = "x-ms-item-count"; public static final String NEW_RESOURCE_ID = "x-ms-new-resource-id"; public static final String USE_MASTER_COLLECTION_RESOLVER = "x-ms-use-master-collection-resolver"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCache.java index 2639ce723979..75278b464062 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCache.java @@ -155,7 +155,7 @@ public Mono tryGetAddresses(RxDocumentServiceRequest reque if (suboptimalServerPartitionTimestamp != null) { logger.debug("suboptimalServerPartitionTimestamp is {}", suboptimalServerPartitionTimestamp); - + boolean forceRefreshDueToSuboptimalPartitionReplicaSet = Duration.between(suboptimalServerPartitionTimestamp, Instant.now()).getSeconds() > this.suboptimalPartitionForceRefreshIntervalInSeconds; @@ -235,7 +235,7 @@ public Mono tryGetAddresses(RxDocumentServiceRequest reque } return Mono.error(unwrappedException); } - + }); } @@ -258,6 +258,11 @@ public Mono> getServerAddressesViaGatewayAsync( headers.put(HttpConstants.HttpHeaders.FORCE_REFRESH, Boolean.TRUE.toString()); } + if(request.forceCollectionRoutingMapRefresh) + { + headers.put(HttpConstants.HttpHeaders.FORCE_COLLECTION_ROUTING_MAP_REFRESH, Boolean.TRUE.toString()); + } + addressQuery.put(HttpConstants.QueryStrings.FILTER, HttpUtils.urlEncode(this.protocolFilter)); addressQuery.put(HttpConstants.QueryStrings.PARTITION_KEY_RANGE_IDS, String.join(",", partitionKeyRangeIds)); @@ -452,6 +457,11 @@ public Mono> getMasterAddressesViaGatewayAsync( headers.put(HttpConstants.HttpHeaders.USE_MASTER_COLLECTION_RESOLVER, Boolean.TRUE.toString()); } + if(request.forceCollectionRoutingMapRefresh) + { + headers.put(HttpConstants.HttpHeaders.FORCE_COLLECTION_ROUTING_MAP_REFRESH, Boolean.TRUE.toString()); + } + queryParameters.put(HttpConstants.QueryStrings.FILTER, HttpUtils.urlEncode(this.protocolFilter)); headers.put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123()); String token = this.tokenProvider.getUserAuthorizationToken( diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCacheTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCacheTest.java index b84e6208a4e9..9c4af7c66937 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCacheTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/internal/directconnectivity/GatewayAddressCacheTest.java @@ -100,21 +100,24 @@ public void getServerAddressesViaGateway(List partitionKeyRangeIds, authorizationTokenProvider, null, getHttpClient(configs)); - - RxDocumentServiceRequest req = + for (int i = 0; i < 2; i++) { + RxDocumentServiceRequest req = RxDocumentServiceRequest.create(OperationType.Create, ResourceType.Document, - collectionLink + "/docs/", - getDocumentDefinition(), new HashMap<>()); - - Mono> addresses = cache.getServerAddressesViaGatewayAsync( + collectionLink + "/docs/", + getDocumentDefinition(), new HashMap<>()); + if (i == 1) { + req.forceCollectionRoutingMapRefresh = true; //testing address api with x-ms-collectionroutingmap-refresh true + } + Mono> addresses = cache.getServerAddressesViaGatewayAsync( req, createdCollection.getResourceId(), partitionKeyRangeIds, false); - PartitionReplicasAddressesValidator validator = new PartitionReplicasAddressesValidator.Builder() + PartitionReplicasAddressesValidator validator = new PartitionReplicasAddressesValidator.Builder() .withProtocol(protocol) .replicasOfPartitions(partitionKeyRangeIds) .build(); - validateSuccess(addresses, validator, TIMEOUT); + validateSuccess(addresses, validator, TIMEOUT); + } } @Test(groups = { "direct" }, dataProvider = "protocolProvider", timeOut = TIMEOUT) @@ -129,21 +132,24 @@ public void getMasterAddressesViaGatewayAsync(Protocol protocol) throws Exceptio authorizationTokenProvider, null, getHttpClient(configs)); - - RxDocumentServiceRequest req = + for (int i = 0; i < 2; i++) { + RxDocumentServiceRequest req = RxDocumentServiceRequest.create(OperationType.Create, ResourceType.Database, - "/dbs", - new Database(), new HashMap<>()); - - Mono> addresses = cache.getMasterAddressesViaGatewayAsync(req, ResourceType.Database, + "/dbs", + new Database(), new HashMap<>()); + if (i == 1) { + req.forceCollectionRoutingMapRefresh = true; //testing address api with x-ms-collectionroutingmap-refresh true + } + Mono> addresses = cache.getMasterAddressesViaGatewayAsync(req, ResourceType.Database, null, "/dbs/", false, false, null); - PartitionReplicasAddressesValidator validator = new PartitionReplicasAddressesValidator.Builder() + PartitionReplicasAddressesValidator validator = new PartitionReplicasAddressesValidator.Builder() .withProtocol(protocol) .replicasOfSamePartition() .build(); - validateSuccess(addresses, validator, TIMEOUT); + validateSuccess(addresses, validator, TIMEOUT); + } } @DataProvider(name = "targetPartitionsKeyRangeAndCollectionLinkParams") @@ -802,7 +808,7 @@ public static void validateSuccess(Mono> observable, testSubscriber.assertValueCount(1); validator.validate(testSubscriber.values().get(0)); } - + @BeforeClass(groups = { "direct" }, timeOut = SETUP_TIMEOUT) public void beforeClass() { client = clientBuilder().build();