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

[jeet1995]

Fix thin-client over-span on MULTI_HASH prefix partition-key queries

Problem

Thin-client (GatewayV2, proxy :10250) requests against a hierarchical (MULTI_HASH) container using only a prefix of the key paths (e.g. PK [tenantId, userId], supply only tenantId) over-spanned to every co-located document on the physical partition, returning other tenants' documents. Full partition keys and QueryPlan requests were unaffected.

Fix (ThinClientStoreModel.java)

On the prefix case (kind == MULTI_HASH && components < paths), compute the prefix EPK range and set these RNTBD headers on the request:

• ReadFeedKeyType = EffectivePartitionKeyRange — tells the backend to interpret the EPK tokens as a range filter; without it the pair is ignored and the read falls back to the whole partition.
• StartEpk / EndEpk — hex EPK bounds of the prefix range; the backend's doc-level filter that trims out-of-prefix documents.
• StartEpkHash / EndEpkHash — binary hash bounds that steer to the correct physical partition.

Core HPK tests (ThinClientQueryE2ETest)

Seed = 4 tenants × 6 users on [/tenantId, /userId]; every test is parity-checked against a Direct-mode baseline.

• testHierarchicalPrefixHalfOpenRange — the query guard: runs SELECT * FROM c ORDER BY c.idx with a prefix PK (tenantId only) → asserts exactly 6 documents (the tenant's users); with the full key → exactly 1.
• testHierarchicalReadAllItemsPrefixPartitionKey — readAllItems prefix → exactly the tenant's 6 documents; a per-document assertion fails if any foreign tenant leaks in.
• testHierarchicalReadManyByPartitionKeysPrefixPartitionKey — readMany prefix → stays scoped to 6.
• testHierarchicalReadAllItemsFullPartitionKey — full key → 1 document (no regression).

[Copilot]

Pull request overview

Fixes a thin-client (Gateway V2 / :10250 proxy) routing bug where prefix hierarchical (MULTI_HASH) partition keys could over-span to all co-located documents on the physical partition, potentially returning documents from other logical partitions. The change scopes these requests by sending the correct RNTBD EPK-range headers and adds E2E coverage to prevent regression.

Changes:

• Compute the prefix effective-partition-key (EPK) range for MULTI_HASH prefix keys and apply it via RNTBD headers (ReadFeedKeyType=EffectivePartitionKeyRange, Start/EndEpk, Start/EndEpkHash) in ThinClientStoreModel.
• Introduce a request flag (prefixPartitionKeyQuery) to preserve prefix-query intent through the parallel query initialization path when the partition key is intentionally nulled.
• Add thin-client E2E tests validating parity with Direct mode for hierarchical prefix queries, plus supporting test infrastructure and changelog entry.

Reviewed changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java	Detect MULTI_HASH prefix PK requests and set EPK-range RNTBD headers to prevent over-span.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java	Add prefixPartitionKeyQuery flag with accessors + clone propagation.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java	Set prefixPartitionKeyQuery when a partial hierarchical key is detected and PK header is suppressed.
sdk/cosmos/azure-cosmos/CHANGELOG.md	Document the thin-client prefix HPK over-span fix.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java	New shared thin-client test base + endpoint-routing assertions.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java	New Direct-vs-thin-client query E2E suite including hierarchical prefix regression tests.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java	Update thin-client endpoint assertion behavior (noting QueryPlan can use classic gateway).

[jeet1995]

@sdkReviewAgent

[tvaron3]

LGTM

[FabianMeiswinkel]

If getFeedrange is of different type - should we convert it to FeedRangeEpkImpl here?

[FabianMeiswinkel]

LGTM except for one possible gap - if the conversion to FeedRangeEpk has happened always before already at least add a comment along those lines - if not, then I think the conversion would have to happen here.

[xinlian12]

QQ: I think we should have some HPK related tests for gateway v1 and direct mode, why the need for creating new tests case vs also enable those for gateway v2?

[xinlian12]

QQ: how these headers are populated in gateway v1? at some point upstream, we should already populate startEpk & endEpk? why not based on those headers existence to decide which header need to populate, similar idea as how direct mode populate the headers?

and also why we have to calculate whether this is a prefix here only for thin client

[xinlian12]

🔴 Blocking · Correctness · Prefix over-span reintroduced after partition split

prefixEpkRange = ((FeedRangeEpkImpl) request.getFeedRange()).getRange();

The parallel-prefix branch here assumes request.getFeedRange() still carries the narrow prefix EPK range that DocumentQueryExecutionContextFactory.resolveFeedRangeBasedOnPrefixContainer produced. That is true on the initial dispatch, but it breaks on split:

1. Parent producer has feedRange = FeedRangeEpkImpl(prefixEpkRange) (narrow).
2. Backend returns GONE; DocumentProducer.produceOnFeedRangeGone calls getReplacementRanges(feedRange.getRange()), which returns the physical partition ranges overlapping the prefix.
3. When multiple physicals overlap, createReplacingDocumentProducersOnSplit creates each child with new FeedRangeEpkImpl(targetRange.toRange()) — the full physical partition range, not the intersection with the parent prefix (DocumentProducer.java:325).
4. On the child, createRequestFunc still evaluates isPartialPartitionKeyQuery(...) = true on the same query options, so request.setPrefixPartitionKeyQuery(true).
5. In this branch we then read feedRange.getRange() and emit StartEpk/EndEpk/StartEpkHash/EndEpkHash = the entire physical partition range as the "prefix" doc filter.

Net result: after a real split (or a stale-routing-map GONE+refresh), the exact over-span this PR is fixing returns — silently, with no exception, no diagnostic log. Because partitionKeyInternal was intentionally nulled upstream, there is no fallback tenant filter either.

Suggested action: Either

• On the split child, intersect the parent's prefix range with the child physical range before wrapping in FeedRangeEpkImpl (child gets a genuinely narrower EPK range), or
• In createChildDocumentProducerOnSplit, clear the prefix flag on split-child requests and fall back to the point path (accept a small RU hit but preserve correctness), or
• Recompute the prefix EPK range from cosmosQueryRequestOptions.getPartitionKey() + partitionKeyDefinition here when isPrefixPartitionKeyQuery() is true, and only use feedRange.getRange() if it is provably ⊆ that recomputed range.

Worth adding a fault-injection test that forces GONE on a prefix query and asserts no foreign-tenant docs come back.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

[xinlian12]

🟡 Recommendation · Correctness · Silent fallback when the invariant breaks [Partially covered by @FabianMeiswinkel]

Building on Fabian's question just above ("If getFeedrange is of different type — should we convert it to FeedRangeEpkImpl here?"): converting isn't quite the answer, because the underlying Range<String> might still be the wrong range. The deeper issue is that today, if isPrefixPartitionKeyQuery() == true but the feedRange is not a FeedRangeEpkImpl or does not actually contain the prefix range, we fall through — and on the parallel path partitionKey == null (intentionally nulled upstream), so the PK-bearing branch at line 262 also can't fire. prefixEpkRange stays null and execution lands in the else at line 313, which emits StartEpkHash/EndEpkHash from the resolved physical partition range: the exact over-span this PR is fixing, silently, with no test signal.

Concrete scenarios where this fires today:

1. After a split — DocumentProducer.createChildDocumentProducerOnSplit (DocumentProducer.java:325) wraps the physical partition range in FeedRangeEpkImpl while isPrefixPartitionKeyQuery stays true. The instanceof check passes, but the range is no longer a prefix. See the separate blocking comment on line 261.
2. Future refactor — any wrapping/decorating of feed ranges (versioned FeedRange, composed feed range) would silently break this branch.

Suggested change: fail loud if the invariant is violated, rather than silently coercing.

if (request.isPrefixPartitionKeyQuery()) {
    FeedRangeInternal fr = request.getFeedRange();
    if (!(fr instanceof FeedRangeEpkImpl)) {
        throw new IllegalStateException(
            "prefixPartitionKeyQuery=true requires FeedRangeEpkImpl on the request, got: " + fr);
    }
    prefixEpkRange = ((FeedRangeEpkImpl) fr).getRange();
}

Also worth guarding against PartitionKey.NONE on the PK-bearing branch: PartitionKeyInternal.None.getComponents() is null. The caller in ParallelDocumentQueryExecutionContextBase guards != PartitionKey.NONE, but DocumentQueryExecutionContextFactory#resolveFeedRangeBasedOnPrefixContainer does not, so a NONE HPK query would NPE inside getEPKRangeForPrefixPartitionKey. Cheap defensive check while you're here.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

[xinlian12]

🟡 Recommendation · Test Coverage · This test does not actually exercise the ThinClientStoreModel prefix-header fix

The test name and its Javadoc ("groupPartitionKeysByPhysicalPartition expands the prefix to its effective-partition-key range; the thin client must apply that as a doc-level filter") suggest this exercises the new StartEpk/EndEpk RNTBD path — but that path is never hit for readManyByPartitionKeys:

• RxDocumentClientImpl.buildSequentialFluxFromScratch (RxDocumentClientImpl.java:4967) sets batchQueryOptions.setFeedRange(new FeedRangeEpkImpl(bd.partitionScope)) — the physical partition range, not batchFilter (the narrow prefix-derived batch range).
• It does not set PartitionKey on batchQueryOptions. So downstream ParallelDocumentQueryExecutionContextBase.initialize sees cosmosQueryRequestOptions.getPartitionKey() == null, isPartialPartitionKeyQuery never runs, and request.setPrefixPartitionKeyQuery(false) is what ends up on the wire.
• In ThinClientStoreModel.wrapInHttpRequest, neither the parallel-prefix branch (line 254, needs isPrefixPartitionKeyQuery()) nor the PK-bearing prefix branch (line 262, needs partitionKey != null) fires. The request falls through to the physical-range else at line 313.

Tenant scoping in readMany is instead enforced by the generated SQL predicate WHERE c.<hpkPath> = @... in ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec. Revert the whole ThinClientStoreModel change and this test still passes — it's a routing + SQL-rewriting regression test, not a fix-guard.

Suggested action: Either rename/scope the test to be honest about what it actually validates, or add a diagnostics-level assertion that the x-ms-read-key-type=EffectivePartitionKeyRange header / StartEpk+EndEpk RNTBD tokens were emitted for the batch request (something that would flip if ThinClientStoreModel regressed). The testHierarchicalReadAllItemsPrefixPartitionKey / testHierarchicalPrefixHalfOpenRange tests do genuinely exercise the fix — this one gives false confidence.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

[xinlian12]

🟡 Recommendation · Test Coverage · Regression is not deterministically reproducible from this fixture

The three prefix regression tests (testHierarchicalPrefixHalfOpenRange, testHierarchicalReadAllItemsPrefixPartitionKey, testHierarchicalReadManyByPartitionKeysPrefixPartitionKey) all probe only hierarchicalKeys.get(0) — tenant-0 — and the per-doc tenant assertion ("returned a foreign tenant's document") only trips if some foreign tenant is co-located with tenant-0 on the same physical partition.

With this fixture:

• Tenant IDs are "tenant-" + t + "-" + UUID.randomUUID() — hash placement is effectively random per run.
• 4 tenants × 6 users = 24 docs on a 12,000 RU/s container (≈ 2 physical partitions).
• With random assignment of 4 tenants across ~2 physicals, there is a non-trivial probability (roughly (1/2)^3 = 12.5% in the worst case) that tenant-0 is the only tenant on its physical partition. In that case, an over-spanning physical-partition read still returns exactly the 6 tenant-0 docs, and every assertion (count == 6, IDs match Direct, tenant == tenant-0) passes whether or not the fix is present.

That means on some fraction of CI runs the whole regression guard is inert. This weakens the "removing the fix must break CI" contract that these tests are supposed to enforce.

Suggested action (any of):

1. Iterate all tenants in the assertion loop (for (int t = 0; t < HIER_TENANTS; ...)) — probability tenant-0 and every other tenant is alone drops to near zero.
2. Bump HIER_TENANTS to ~16 so pigeonhole guarantees co-location per physical partition.
3. After seeding, use Direct diagnostics (distinctPartitionKeyRangesContacted / activity from queryItems("SELECT c.<tenant> FROM c", crossPartition)) to bucket tenants by physical PKRangeId, pick a tenant proven to share a range with another, and run the assertions against that tenant. Fail-fast the test if no co-located pair exists (setup problem, not a product bug).
4. Add a nonexistent-tenant prefix negative test — readAllItems with a made-up tenantId prefix should return 0 docs; without the fix it would return whatever co-located tenant's docs are in the resolved physical partition.

Option 4 is arguably the cleanest: it makes co-location irrelevant and directly exposes over-span.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

[xinlian12]

💬 Observation · Cross-SDK · Peer .NET thin-client likely has the same over-span bug

Not actionable for this PR, but worth flagging so the team can decide whether to file a peer issue on Azure/azure-cosmos-dotnet-v3: the .NET thin-client path appears to have the same defect this PR is fixing here.

• ThinClientStoreClient.PrepareRequestForProxyAsync sets ProxyStartEpk/ProxyEndEpk from request.RequestContext?.ResolvedPartitionKeyRange.MinInclusive/.MaxExclusive — i.e. the resolved physical partition range.
• ThinClientTransportSerializer.SerializeProxyRequestAsync then unconditionally propagates those values into HttpConstants.HttpHeaders.StartEpk / EndEpk (and WFConstants.BackendHeaders.StartEpkHash/EndEpkHash), with no branch that narrows to a prefix EPK range for partial MULTI_HASH keys. The upstream RequestInvokerHandler.ResolveFeedRangeBasedOnPrefixContainerAsync correctly computes a prefix FeedRangeEpk, but that narrower range is discarded by the time the RNTBD body is serialized.
• .NET has no CosmosMultiHashTest variant that combines thin-client + prefix HPK (checked ThinClientTransportSerializerTests.cs and CosmosItemThinClientTests.cs).

Python and Rust are architecturally unaffected — neither serializes RNTBD tokens on the client; they send x-ms-start-epk / x-ms-end-epk as plain HTTP headers and the (correctly narrow) prefix range flows through.

If confirmed, this Java PR is a good reference implementation for the .NET-side fix.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

[xinlian12]

✅ Review complete (57:48)

Posted 5 inline comment(s).

_{Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage}