From e866227731fd859a5e7a29bec686e57979cc2cd5 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 10 Apr 2025 10:10:13 -0700 Subject: [PATCH 1/4] Fix rntbd headers for proxy request. --- Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs | 5 ----- .../src/ThinClientTransportSerializer.cs | 11 +++++++---- 2 files changed, 7 insertions(+), 9 deletions(-) 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(); From f2fa3e51aaadf58b7cb7d75af3a4990d1f6331e8 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 10 Apr 2025 13:18:17 -0700 Subject: [PATCH 2/4] Update tests --- .../src/Resource/Settings/ContainerProperties.cs | 2 +- .../ThinClientTransportSerializerTests.cs | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs b/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs index 32f0dd7732..31d4bd32a1 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs @@ -708,7 +708,7 @@ internal ContainerProperties(string id, PartitionKeyDefinition partitionKeyDefin /// [JsonProperty(PropertyName = Constants.Properties.RId, NullValueHandling = NullValueHandling.Ignore)] - internal string ResourceId { get; private set; } + internal string ResourceId { get; set; } internal bool HasPartitionKey => this.PartitionKey != null; 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..5b72268cfc 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( @@ -177,7 +184,8 @@ private ContainerProperties GetMockContainerProperties() PartitionKey = new PartitionKeyDefinition { Paths = new Collection { "/pk" } - } + }, + ResourceId = "-Jlvm9pqHGk=", }; } } From 762d9eae4bbfdb61c4c5131f0f548c9f77cdafd9 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 10 Apr 2025 14:16:00 -0700 Subject: [PATCH 3/4] Update ResolveThinClientEndpoint check. --- Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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); } } } From fb7d1b97eb680021f9e4e0f1320caf833db7f298 Mon Sep 17 00:00:00 2001 From: "Arooshi Avasthy (from Dev Box)" Date: Mon, 14 Apr 2025 11:41:31 -0700 Subject: [PATCH 4/4] Update test. --- .../Resource/Settings/ContainerProperties.cs | 2 +- .../ThinClientTransportSerializerTests.cs | 28 +++++++++++-------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs b/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs index 31d4bd32a1..32f0dd7732 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs @@ -708,7 +708,7 @@ internal ContainerProperties(string id, PartitionKeyDefinition partitionKeyDefin /// [JsonProperty(PropertyName = Constants.Properties.RId, NullValueHandling = NullValueHandling.Ignore)] - internal string ResourceId { get; set; } + internal string ResourceId { get; private set; } internal bool HasPartitionKey => this.PartitionKey != null; 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 5b72268cfc..12670038ca 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientTransportSerializerTests.cs @@ -175,18 +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" } - }, - ResourceId = "-Jlvm9pqHGk=", - }; - } } }