[Internal] Thin Client Integration: Fixes rntbd headers for proxy request.

Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs

@@ -771,11 +771,7 @@ public IList<string> 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);
         }
     }
 }

Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs

@@ -91,11 +91,6 @@ private async ValueTask<HttpRequestMessage> 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());

Microsoft.Azure.Cosmos/src/ThinClientTransportSerializer.cs

@@ -63,6 +63,11 @@ public static async Task<Stream> 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<Stream> SerializeProxyRequestAsync(
                     throw new InternalServerErrorException();
                 }
 
-                ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(
-                     request,
-                     CancellationToken.None,
-                     NoOpTrace.Singleton);
                 string epk = GetEffectivePartitionKeyHash(partitionKey, collection.PartitionKey);
 
                 request.Properties = new Dictionary<string, object>
@@ -96,6 +97,8 @@ public static async Task<Stream> 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();
 

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<DocumentServiceRequest>(),
+                  It.IsAny<CancellationToken>(),
+                  It.IsAny<ITrace>()))
+              .ReturnsAsync(this.GetMockContainerProperties());
+
             // Act & Assert
             await Assert.ThrowsExceptionAsync<InternalServerErrorException>(
                 () => 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<string> { "/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<string> { "/pk" }
-                }
-            };
-        }
     }
 }