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

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

@@ -2176,16 +2176,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) {

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

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.rx;
+
+import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosDiagnostics;
+import com.azure.cosmos.CosmosDiagnosticsContext;
+import com.azure.cosmos.CosmosDiagnosticsRequestInfo;
+import com.azure.cosmos.CosmosClientBuilder;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+
+import java.util.Collection;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+/**
+ * Base class for thin client E2E tests. Provides shared setup/teardown,
+ * constants, and helper methods common to all thin client test classes.
+ */
+public abstract class ThinClientTestBase extends TestSuiteBase {
+
+    protected static final String THIN_CLIENT_ENDPOINT_INDICATOR = ":10250/";
+    protected static final String ID_FIELD = "id";
+    protected static final String PARTITION_KEY_FIELD = "mypk";
+    protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    protected CosmosAsyncClient client;
+    protected CosmosAsyncContainer container;
+
+    protected ThinClientTestBase(CosmosClientBuilder clientBuilder) {
+        super(clientBuilder);
+    }
+
+    @BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT)
+    public void before_ThinClientTest() {
+        assertThat(this.client).isNull();
+        enableThinClientForTest();
+        this.client = getClientBuilder().buildAsyncClient();
+        this.container = getSharedMultiPartitionCosmosContainer(this.client);
+
+        // Clean up shared container to prevent cross-test-class pollution.
+        cleanUpContainer(this.container);
+    }
+
+    @AfterClass(groups = {"thinclient"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
+    public void afterClass() {
+        clearThinClientForTest();
+        if (this.client != null) {
+            this.client.close();
+        }
+    }
+
+    protected static void enableThinClientForTest() {
+        System.setProperty("COSMOS.THINCLIENT_ENABLED", "true");
+    }
+
+    protected static void clearThinClientForTest() {
+        System.clearProperty("COSMOS.THINCLIENT_ENABLED");
+    }
+
+    /**
+     * Creates a test document with id and mypk fields (matching shared container partition key).
+     */
+    protected ObjectNode createTestDocument(String id, String mypk) {
+        ObjectNode doc = OBJECT_MAPPER.createObjectNode();
+        doc.put(ID_FIELD, id);
+        doc.put(PARTITION_KEY_FIELD, mypk);
+        return doc;
+    }
+
+    /**
+     * Asserts that all requests in the diagnostics were routed through the thin client endpoint.
+     */
+    protected static void assertThinClientEndpointUsed(CosmosDiagnostics diagnostics) {
+        // Delegate to the shared TestSuiteBase implementation so the thin-client routing invariant
+        // (every non-QueryPlan request via the thin-client endpoint; QueryPlan may use the classic
+        // gateway) and null-endpoint handling are applied consistently across all thin-client tests.
+        TestSuiteBase.assertThinClientEndpointUsed(diagnostics);
+    }
+
+    /**
+     * Asserts that NO requests in the diagnostics were routed through the thin client endpoint,
+     * confirming the gateway client used the standard :443 path.
+     */
+    protected static void assertGatewayEndpointUsed(CosmosDiagnostics diagnostics) {
+        assertThat(diagnostics).isNotNull();
+        CosmosDiagnosticsContext ctx = diagnostics.getDiagnosticsContext();
+        assertThat(ctx).isNotNull();
+        Collection<CosmosDiagnosticsRequestInfo> requests = ctx.getRequestInfo();
+        assertThat(requests).isNotNull();
+        assertThat(requests.size()).isPositive();
+        for (CosmosDiagnosticsRequestInfo requestInfo : requests) {
+            assertThat(requestInfo.getEndpoint())
+                .as("Gateway client must not route through thin client endpoint, but found: " + requestInfo.getEndpoint())
+                .doesNotContain(THIN_CLIENT_ENDPOINT_INDICATOR);
+        }
+    }
+}

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Unified request-level consistency override behavior across transports: invalid attempts to upgrade the request consistency level above the account default are now silently ignored instead of returning `BadRequest` in some gateway paths. - See PR [49606](https://github.com/Azure/azure-sdk-for-java/pull/49606).
+* Fixed thin-client (Gateway V2) queries with a prefix (partial) hierarchical partition key returning co-located documents from other logical partitions. - See PR [49688](https://github.com/Azure/azure-sdk-for-java/pull/49688).
 
 #### Other Changes
 * Reduced memory footprint of deserialized `PartitionKeyRange` instances by stripping unused fields in the `PartitionKeyRange(ObjectNode)` constructor - See PR [49513](https://github.com/Azure/azure-sdk-for-java/pull/49513).

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java

@@ -78,6 +78,12 @@ private static CosmosItemSerializer internalDefaultSerializer() {
     private Range<String> effectiveRange;
     private int numberOfItemsInBatchRequest;
 
+    // Marks a query whose partition key is a partial (prefix) hierarchical (MULTI_HASH) key. On this
+    // path the partial partition key is intentionally nulled upstream and the narrow prefix
+    // effective-partition-key range is carried on feedRange instead. Downstream store models reuse
+    // that range as a doc-level EPK filter rather than over-spanning the whole physical partition.
+    private boolean prefixPartitionKeyQuery;
+
     private byte[] contentAsByteArray;
 
     // NOTE: TODO: these fields are copied from .Net SDK
@@ -906,6 +912,14 @@ public FeedRangeInternal getFeedRange() {
         return this.feedRange;
     }
 
+    public boolean isPrefixPartitionKeyQuery() {
+        return this.prefixPartitionKeyQuery;
+    }
+
+    public void setPrefixPartitionKeyQuery(boolean prefixPartitionKeyQuery) {
+        this.prefixPartitionKeyQuery = prefixPartitionKeyQuery;
+    }
+
     public void applyFeedRangeFilter(FeedRangeInternal feedRange) {
         this.feedRange = feedRange;
     }
@@ -1081,6 +1095,7 @@ public RxDocumentServiceRequest clone() {
         rxDocumentServiceRequest.setResourceAddress(this.resourceAddress);
         rxDocumentServiceRequest.requestContext = this.requestContext.clone();
         rxDocumentServiceRequest.feedRange = this.feedRange;
+        rxDocumentServiceRequest.prefixPartitionKeyQuery = this.prefixPartitionKeyQuery;
         rxDocumentServiceRequest.effectiveRange = this.effectiveRange;
         rxDocumentServiceRequest.isFeed = this.isFeed;
         rxDocumentServiceRequest.resourceId = this.resourceId;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java

@@ -14,8 +14,12 @@
 import com.azure.cosmos.implementation.http.HttpHeaders;
 import com.azure.cosmos.implementation.http.HttpRequest;
 import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
+import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
 import com.azure.cosmos.implementation.routing.HexConvert;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
+import com.azure.cosmos.implementation.routing.Range;
+import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.azure.cosmos.models.PartitionKind;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufUtil;
 import io.netty.buffer.Unpooled;
@@ -27,6 +31,7 @@
 import reactor.core.publisher.Mono;
 
 import java.net.URI;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -239,16 +244,64 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque
             request.properties = new HashMap<>();
         }
 
+        PartitionKeyInternal partitionKey = request.getPartitionKeyInternal();
+
+        // Detect a partial (prefix) hierarchical partition key BEFORE the RNTBD request is serialized.
+        // The prefix's effective-partition-key sub-range [hash(prefix), hash(prefix) + "FF") is applied
+        // as the RNTBD EPK headers in the consolidated block after RntbdRequest.from() below.
+        Range<String> prefixEpkRange = null;
+        if (request.getOperationType() != OperationType.QueryPlan) {
+            if (request.isPrefixPartitionKeyQuery() && request.getFeedRange() instanceof FeedRangeEpkImpl) {
+                // Parallel prefix-query path (cached): the partial hierarchical partition key was
+                // intentionally nulled upstream (ParallelDocumentQueryExecutionContextBase#initialize), but
+                // the narrow prefix effective-partition-key range was already computed and carried on the
+                // request's feedRange (DocumentQueryExecutionContextFactory#resolveFeedRangeBasedOnPrefixContainer).
+                // Reuse that cached range directly instead of recomputing it from a partition key we no
+                // longer have.
+                prefixEpkRange = ((FeedRangeEpkImpl) request.getFeedRange()).getRange();
+            } else if (partitionKey != null) {
+                // PK-bearing prefix path: recompute the prefix range directly from the partial hierarchical
+                // partition key.
+                PartitionKeyDefinition pkDefinition = request.getPartitionKeyDefinition();
+                if (pkDefinition != null
+                    && pkDefinition.getKind() == PartitionKind.MULTI_HASH
+                    && partitionKey.getComponents().size() < pkDefinition.getPaths().size()) {
+                    prefixEpkRange = partitionKey.getEPKRangeForPrefixPartitionKey(pkDefinition);
+                }
+            }
+        }
+
         RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs(request);
 
         HttpHeaders headers = this.getHttpHeaders();
         headers.set(HttpConstants.HttpHeaders.ACTIVITY_ID, request.getActivityId().toString());
 
         RntbdRequest rntbdRequest = RntbdRequest.from(rntbdRequestArgs);
 
-        PartitionKeyInternal partitionKey = request.getPartitionKeyInternal();
-
-        if (partitionKey != null) {
+        if (prefixEpkRange != null) {
+            // Partial (prefix) hierarchical partition key - either the PK-bearing path (point/scalar
+            // prefix) or the parallel path that reused request.getFeedRange(). All five EPK-specific
+            // headers are RNTBD tokens, so set them directly on the RNTBD request here in one place.
+            //
+            // ReadFeedKeyType=EffectivePartitionKeyRange + StartEpk/EndEpk carry the prefix range as the
+            // backend's doc-level EPK filter. StartEpk/EndEpk use the hex string's UTF-8 bytes (NOT the
+            // decoded hash) - exactly as the Direct/RNTBD path does (see FeedRangeEpkImpl); the proxy
+            // forwards them to the backend (see TransportSerialization AddStartAndEndKeysFromHeaders,
+            // where ReadFeedKeyType=EffectivePartitionKeyRange selects the hex-string-as-bytes encoding).
+            // StartEpkHash/EndEpkHash (the decoded hash bytes) additionally steer the proxy to the owning
+            // physical partition(s). Without this filter the proxy resolves the request to the owning
+            // physical partition and returns every co-located document (an over-span).
+            rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.ReadFeedKeyType,
+                RntbdConstants.RntbdReadFeedKeyType.EffectivePartitionKeyRange.id());
+            rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpk,
+                prefixEpkRange.getMin().getBytes(StandardCharsets.UTF_8));
+            rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpk,
+                prefixEpkRange.getMax().getBytes(StandardCharsets.UTF_8));
+            rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpkHash,
+                HexConvert.hexToBytes(prefixEpkRange.getMin()));
+            rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpkHash,
+                HexConvert.hexToBytes(prefixEpkRange.getMax()));
+        } else if (partitionKey != null) {
             byte[] epk = partitionKey.getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition());
             rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EffectivePartitionKey, epk);
         } else if (request.requestContext.resolvedPartitionKeyRange == null) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java

@@ -81,20 +81,24 @@ protected void initialize(
                 headers.put(HttpConstants.HttpHeaders.CORRELATED_ACTIVITY_ID, correlatedActivityId.toString());
 
                 PartitionKeyInternal partitionKeyInternal = null;
+                boolean isPrefixPartitionKeyQuery = false;
                 if (cosmosQueryRequestOptions.getPartitionKey() != null && cosmosQueryRequestOptions.getPartitionKey() != PartitionKey.NONE) {
                     // the next if statement is for the method `createDocumentServiceRequestWithFeedRange` below
                     // with this added logic we don't have to remove the header (since the header never gets set) and
                     // partitionKeyInternal gets passed as null which avoids the feedRange normalization.
-                    if (!PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey())) {
+                    isPrefixPartitionKeyQuery = PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey());
+                    if (!isPrefixPartitionKeyQuery) {
                         partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey());
                         headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
                     }
                 }
 
                 queryOptionsAccessor().setPartitionKeyDefinition(cosmosQueryRequestOptions, collection.getPartitionKey());
                 queryOptionsAccessor().setCollectionRid(cosmosQueryRequestOptions, collection.getResourceId());
-                return this.createDocumentServiceRequestWithFeedRange(headers, querySpecForInit, partitionKeyInternal, feedRange,
+                RxDocumentServiceRequest request = this.createDocumentServiceRequestWithFeedRange(headers, querySpecForInit, partitionKeyInternal, feedRange,
                                                          collection.getResourceId(), cosmosQueryRequestOptions.getThroughputControlGroupName());
+                request.setPrefixPartitionKeyQuery(isPrefixPartitionKeyQuery);
+                return request;
             };
 
             Function<RxDocumentServiceRequest, Mono<FeedResponse<T>>> executeFunc = (request) ->  this.executeRequestAsync(