[GatewayV2] Fix over-spanning issue when partial partition keys are provided.

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

@@ -6,6 +6,7 @@
 
 package com.azure.cosmos;
 
+import com.azure.cosmos.implementation.TestConfigurations;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
@@ -42,6 +43,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
+import static com.azure.cosmos.rx.ThinClientTestBase.assertGatewayEndpointUsed;
+import static com.azure.cosmos.rx.ThinClientTestBase.assertThinClientEndpointUsed;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class CosmosMultiHashTest extends TestSuiteBase {
@@ -53,6 +56,9 @@ public class CosmosMultiHashTest extends TestSuiteBase {
     private CosmosDatabase createdDatabase;
     private CosmosContainer createdMultiHashContainer;
     private CosmosContainer createdNestedPathContainer;
+    // Reflects the active @BeforeClass lifecycle: true under the thinclient group, false under emulator.
+    // Drives mode-aware endpoint-routing assertions and the continuation-token client selection below.
+    private boolean thinClientEnabled;
 
     @Factory(dataProvider = "clientBuilders")
     public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) {
@@ -61,7 +67,25 @@ public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) {
 
     @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT)
     public void before_CosmosMultiHashTest() {
+        thinClientEnabled = false;
         client = getClientBuilder().buildClient();
+        initDatabaseAndContainers();
+    }
+
+    // Enrolls the MULTI_HASH prefix over-span coverage into the thin-client (GatewayV2, proxy :10250) group.
+    // The class-level @Factory yields a plain gateway builder without HTTP/2, which cannot route to the thin
+    // client, so this lifecycle builds an explicit HTTP/2 gateway client (mirrors clientBuildersWithGatewayAndHttp2)
+    // so the same test bodies exercise the RNTBD prefix-EPK header path. COSMOS.THINCLIENT_ENABLED is supplied by
+    // the thin-client CI lane (see sdk/cosmos/tests.yml); IDE/standalone runs must pass -DCOSMOS.THINCLIENT_ENABLED=true.
+    @BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT)
+    public void before_CosmosMultiHashTest_thinClient() {
+        thinClientEnabled = true;
+        client = createGatewayRxDocumentClient(
+            TestConfigurations.HOST, null, true, null, true, true, true).buildClient();
+        initDatabaseAndContainers();
+    }
+
+    private void initDatabaseAndContainers() {
         createdDatabase = createSyncDatabase(client, preExistingDatabaseId);
         String collectionName = UUID.randomUUID().toString();
 
@@ -103,7 +127,24 @@ public void afterClass() {
         safeCloseSyncClient(client);
     }
 
-    @Test(groups = {"emulator"}, timeOut = TIMEOUT)
+    @AfterClass(groups = {"thinclient"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
+    public void afterClass_thinClient() {
+        logger.info("starting cleanup (thin client)....");
+        safeDeleteSyncDatabase(createdDatabase);
+        safeCloseSyncClient(client);
+    }
+
+    // Mode-aware endpoint routing assertion. Under the thinclient group every non-QueryPlan (data) request
+    // must route through the thin-client proxy endpoint (:10250); under the emulator group no request may.
+    private void assertEndpointRouting(CosmosDiagnostics diagnostics) {
+        if (thinClientEnabled) {
+            assertThinClientEndpointUsed(diagnostics);
+        } else {
+            assertGatewayEndpointUsed(diagnostics);
+        }
+    }
+
+    @Test(groups = {"emulator", "thinclient"}, timeOut = TIMEOUT)
     public void itemCRUD() {
         CityItem cityItem = new CityItem(UUID.randomUUID().toString(), "Redmond", "98052", 1);
 
@@ -120,6 +161,7 @@ public void itemCRUD() {
             cityItem.getId(), partitionKey, CityItem.class);
 
         assertThat(readResponse.getItem().toString()).isEqualTo(cityItem.toString());
+        assertEndpointRouting(readResponse.getDiagnostics());
         createdMultiHashContainer.deleteItem(cityItem.getId(), partitionKey, new CosmosItemRequestOptions());
     }
 
@@ -170,7 +212,7 @@ private void validateResponse(FeedResponse<ObjectNode> response,
                     .collect(Collectors.toList())
         );
     }
-    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
     public void readManySupportsNestedPartitionKeyPaths() {
         String city = "nested-readmany-" + UUID.randomUUID();
 
@@ -186,9 +228,10 @@ public void readManySupportsNestedPartitionKeyPaths() {
 
         FeedResponse<ObjectNode> documentFeedResponse = createdNestedPathContainer.readMany(itemList, ObjectNode.class);
         validateResponse(documentFeedResponse, itemList);
+        assertEndpointRouting(documentFeedResponse.getCosmosDiagnostics());
     }
 
-    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
     public void readAllItemsSupportsNestedPartitionKeyPaths() {
         String city = "nested-readall-" + UUID.randomUUID();
 
@@ -203,7 +246,13 @@ public void readAllItemsSupportsNestedPartitionKeyPaths() {
         CosmosPagedIterable<ObjectNode> readAllResults =
             createdNestedPathContainer.readAllItems(new PartitionKey(city), ObjectNode.class);
 
-        assertThat(readAllResults.stream().map(item -> item.get("id").asText()).collect(Collectors.toList()))
+        List<String> readAllIds = new ArrayList<>();
+        for (FeedResponse<ObjectNode> page : readAllResults.iterableByPage()) {
+            page.getResults().forEach(item -> readAllIds.add(item.get("id").asText()));
+            assertEndpointRouting(page.getCosmosDiagnostics());
+        }
+
+        assertThat(readAllIds)
             .containsExactlyInAnyOrder(firstItem.get("id").asText(), secondItem.get("id").asText());
     }
 
@@ -445,7 +494,7 @@ private void validateDocCRUDAndQuery() throws InterruptedException {
         deleteAllItems();
     }
 
-    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
     private void multiHashQueryTests() {
         ArrayList<CityItem> docs = createItems();
 
@@ -530,6 +579,7 @@ private void multiHashQueryTests() {
         FeedResponse<ObjectNode> feedResponse = feedResponses.iterator().next();
         assertThat(feedResponse.getResults().size()).isEqualTo(2);
         assertThat(feedResponse.getResults().get(0).get("zipcode").asInt() < feedResponse.getResults().get(1).get("zipcode").asInt()).isTrue();
+        assertEndpointRouting(feedResponse.getCosmosDiagnostics());
 
         //Using continuation token
         testPartialPKContinuationToken();
@@ -645,7 +695,13 @@ private void deleteAllItems() {
     private void testPartialPKContinuationToken() {
         String requestContinuation = null;
         List<ObjectNode> receivedDocuments = new ArrayList<>();
-        CosmosAsyncClient asyncClient = getClientBuilder().buildAsyncClient();
+        // Under the thinclient group build an explicit HTTP/2 gateway client so this prefix continuation-token
+        // pass also routes to the thin client (:10250) instead of the plain gateway; the emulator group uses the
+        // factory-provided builder. thinClientEnabled mirrors the active @BeforeClass lifecycle.
+        CosmosAsyncClient asyncClient =
+            thinClientEnabled
+                ? createGatewayRxDocumentClient(TestConfigurations.HOST, null, true, null, true, true, true).buildAsyncClient()
+                : getClientBuilder().buildAsyncClient();
         CosmosAsyncDatabase cosmosAsyncDatabase = new CosmosAsyncDatabase(createdDatabase.getId(), asyncClient);
         CosmosAsyncContainer cosmosAsyncContainer = new CosmosAsyncContainer(createdMultiHashContainer.getId(), cosmosAsyncDatabase);
         String query = "SELECT * FROM c ORDER BY c.zipcode ASC";
@@ -671,6 +727,7 @@ private void testPartialPKContinuationToken() {
             requestContinuation = firstPage.getContinuationToken();
             receivedDocuments.addAll(firstPage.getResults());
             assertThat(firstPage.getResults().size()).isEqualTo(1);
+            assertEndpointRouting(firstPage.getCosmosDiagnostics());
         } while (requestContinuation != null);
         assertThat(receivedDocuments.size()).isEqualTo(3);
         asyncClient.close();

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java

@@ -2710,16 +2710,26 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnosticsContext ctx)
         assertThat(requests).isNotNull();
         assertThat(requests.size()).isPositive();
 
+        // Validate every request rather than early-returning on the first thin-client match: a mixed
+        // scenario (some data requests via the thin-client endpoint, some via the classic gateway) must
+        // fail. Every non-QueryPlan (data) request must route through the thin-client proxy endpoint;
+        // QueryPlan calls are resolved via the classic gateway in thin-client mode, so they are the only
+        // requests allowed to target a non-thin-client endpoint.
         for (CosmosDiagnosticsRequestInfo requestInfo : requests) {
-            if (requestInfo.getEndpoint() != null
-                && requestInfo.getEndpoint().contains(THIN_CLIENT_ENDPOINT_INDICATOR)) {
-                return;
+            // requestType has the form "<ResourceType>:<OperationType>" (OperationType.QueryPlan
+            // stringifies to "QueryPlan").
+            String requestType = requestInfo.getRequestType();
+            if (requestType != null && requestType.endsWith(":QueryPlan")) {
+                continue;
             }
-        }
 
-        assertThat(false)
-            .as("No request targeting thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR + ")")
-            .isTrue();
+            String endpoint = requestInfo.getEndpoint();
+            assertThat(endpoint != null && endpoint.contains(THIN_CLIENT_ENDPOINT_INDICATOR))
+                .as("Non-QueryPlan request must target the thin client proxy endpoint ("
+                    + THIN_CLIENT_ENDPOINT_INDICATOR + "), but was: " + endpoint
+                    + " (requestType: " + requestType + ")")
+                .isTrue();
+        }
     }
 
     protected static void safeClose(AsyncDocumentClient client) {