[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;
@@ -62,6 +63,22 @@ public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) {
     @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT)
     public void before_CosmosMultiHashTest() {
         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)
+    // and sets COSMOS.THINCLIENT_ENABLED so the same test bodies exercise the RNTBD prefix-EPK header path.
+    @BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT)
+    public void before_CosmosMultiHashTest_thinClient() {
+        System.setProperty("COSMOS.THINCLIENT_ENABLED", "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 +120,15 @@ 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);
+        System.clearProperty("COSMOS.THINCLIENT_ENABLED");
+    }
+
+    @Test(groups = {"emulator", "thinclient"}, timeOut = TIMEOUT)
     public void itemCRUD() {
         CityItem cityItem = new CityItem(UUID.randomUUID().toString(), "Redmond", "98052", 1);
 
@@ -170,7 +195,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();
 
@@ -188,7 +213,7 @@ public void readManySupportsNestedPartitionKeyPaths() {
         validateResponse(documentFeedResponse, itemList);
     }
 
-    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
     public void readAllItemsSupportsNestedPartitionKeyPaths() {
         String city = "nested-readall-" + UUID.randomUUID();
 
@@ -445,7 +470,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();
 
@@ -645,7 +670,12 @@ private void deleteAllItems() {
     private void testPartialPKContinuationToken() {
         String requestContinuation = null;
         List<ObjectNode> receivedDocuments = new ArrayList<>();
-        CosmosAsyncClient asyncClient = getClientBuilder().buildAsyncClient();
+        // Under the thinclient group the process-wide property is set; build an HTTP/2 gateway client so this
+        // prefix continuation-token pass also routes to the thin client (:10250) instead of the plain gateway.
+        CosmosAsyncClient asyncClient =
+            Boolean.parseBoolean(System.getProperty("COSMOS.THINCLIENT_ENABLED"))
+                ? 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";

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) {