diff --git a/Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs b/Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs index 00670beafd..430fdaf1ae 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs @@ -771,11 +771,7 @@ public IList GetEffectivePreferredLocations() public Uri ResolveThinClientEndpoint(DocumentServiceRequest request) { - bool isReadRequest = request.IsReadOnlyRequest - || request.OperationType == OperationType.Query - || request.OperationType == OperationType.ReadFeed; - - return this.locationCache.ResolveThinClientEndpoint(request, isReadRequest); + return this.locationCache.ResolveThinClientEndpoint(request, request.IsReadOnlyRequest); } } } diff --git a/Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs b/Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs index 1ddce45dd4..ad49baf04e 100644 --- a/Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs +++ b/Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs @@ -91,11 +91,6 @@ private async ValueTask PrepareRequestForProxyAsync( BufferProviderWrapper bufferProviderWrapper = this.bufferProviderWrapperPool.Get(); try { - ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync( - request, - CancellationToken.None, - NoOpTrace.Singleton); - request.ResourceId = collection.ResourceId; requestMessage.Headers.TryAddWithoutValidation( ThinClientConstants.ProxyOperationType, request.OperationType.ToOperationTypeString()); diff --git a/Microsoft.Azure.Cosmos/src/ThinClientTransportSerializer.cs b/Microsoft.Azure.Cosmos/src/ThinClientTransportSerializer.cs index dd05ae9d15..f7d8ec71a3 100644 --- a/Microsoft.Azure.Cosmos/src/ThinClientTransportSerializer.cs +++ b/Microsoft.Azure.Cosmos/src/ThinClientTransportSerializer.cs @@ -63,6 +63,11 @@ public static async Task SerializeProxyRequestAsync( requestStream, AuthorizationTokenType.PrimaryMasterKey, dictionaryCollection); + ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync( + request, + CancellationToken.None, + NoOpTrace.Singleton); + if (operationType.IsPointOperation()) { string partitionKey = request.Headers.Get(HttpConstants.HttpHeaders.PartitionKey); @@ -72,10 +77,6 @@ public static async Task SerializeProxyRequestAsync( throw new InternalServerErrorException(); } - ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync( - request, - CancellationToken.None, - NoOpTrace.Singleton); string epk = GetEffectivePartitionKeyHash(partitionKey, collection.PartitionKey); request.Properties = new Dictionary @@ -96,6 +97,8 @@ public static async Task SerializeProxyRequestAsync( request.Headers.Add(HttpConstants.HttpHeaders.StartEpk, request.Headers[ThinClientConstants.ProxyStartEpk]); request.Headers.Add(HttpConstants.HttpHeaders.EndEpk, request.Headers[ThinClientConstants.ProxyEndEpk]); } + request.ResourceId = collection.ResourceId; + request.Headers.Add(WFConstants.BackendHeaders.CollectionRid, collection.ResourceId); await request.EnsureBufferedBodyAsync(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs index 7aa7d10fc9..12670038ca 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs @@ -57,6 +57,13 @@ public async Task SerializeProxyRequestAsync_ShouldThrowIfNoPartitionKeyInPointO CallBase = true }; + clientCollectionCacheMock + .Setup(c => c.ResolveCollectionAsync( + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(this.GetMockContainerProperties()); + // Act & Assert await Assert.ThrowsExceptionAsync( () => ThinClientTransportSerializer.SerializeProxyRequestAsync( @@ -168,17 +175,24 @@ public async Task ConvertProxyResponseAsync_ShouldReturnHttpResponse_WhenValid() Assert.IsTrue( converted.Headers.Any(h => h.Key == ThinClientConstants.RoutedViaProxy), "Expected 'x-ms-thinclient-route-via-proxy' header to be set in the converted response."); + } + + private ContainerProperties GetMockContainerProperties() + { + ContainerProperties containerProperties = new ContainerProperties + { + PartitionKey = new PartitionKeyDefinition + { + Paths = new Collection { "/pk" } + } + }; + + typeof(ContainerProperties) + .GetProperty("ResourceId", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance) + ?.SetValue(containerProperties, "-Jlvm9pqHGk="); + + return containerProperties; } - private ContainerProperties GetMockContainerProperties() - { - return new ContainerProperties - { - PartitionKey = new PartitionKeyDefinition - { - Paths = new Collection { "/pk" } - } - }; - } } }