From 499f928f1973953e1b5b172304dfaf0ff0bdea15 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 09:48:53 -0400 Subject: [PATCH 01/15] Fix thin-client MULTI_HASH prefix partition key EPK over-span For a partial (prefix) hierarchical (MULTI_HASH) partition key, the thin-client (Gateway V2) store model previously sent only StartEpkHash/EndEpkHash routing headers and no doc-level EPK filter. The proxy therefore resolved the request to the owning physical partition and returned every co-located document (an over-span) instead of only those matching the prefix. ThinClientStoreModel.wrapInHttpRequest now detects a prefix MULTI_HASH key and sets READ_FEED_KEY_TYPE=EffectivePartitionKeyRange, START_EPK and END_EPK on the request headers before RntbdRequest.from(), so RntbdRequestHeaders serializes the prefix EPK sub-range [hash(prefix), hash(prefix) + "FF") as the backend doc-level filter (hex string as UTF-8 bytes), mirroring the Direct/RNTBD FeedRangeEpkImpl path. StartEpkHash/EndEpkHash routing headers are retained. QueryPlan requests and full (non-prefix) keys are unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../implementation/ThinClientStoreModel.java | 47 +++++++++++++++++-- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index 3a849b4c5baa..6fe5267ab28a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -16,6 +16,9 @@ import com.azure.cosmos.implementation.caches.RxClientCollectionCache; 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; @@ -239,6 +242,35 @@ 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. + // For such a query the prefix's effective-partition-key sub-range + // [hash(prefix), hash(prefix) + "FF") must reach the thin-client proxy as a doc-level EPK + // filter. We convey it by setting StartEpk/EndEpk + ReadFeedKeyType=EffectivePartitionKeyRange + // on the request headers here, so that RntbdRequest.from() -> + // RntbdRequestHeaders.addStartAndEndKeys serializes them into the RNTBD request with the + // correct encoding (the hex string taken as UTF-8 bytes, not the decoded hash) - exactly as + // the Direct/RNTBD path does for the same query (see FeedRangeEpkImpl). The proxy forwards + // these to the backend as the doc-level EPK filter (see TransportSerialization + // AddStartAndEndKeysFromHeaders, where ReadFeedKeyType=EffectivePartitionKeyRange selects the + // hex-string-as-bytes encoding). Without this filter the proxy resolves the request to the + // owning physical partition and returns every co-located document (an over-span). + Range prefixEpkRange = null; + if (request.getOperationType() != OperationType.QueryPlan && partitionKey != null) { + PartitionKeyDefinition pkDefinition = request.getPartitionKeyDefinition(); + if (pkDefinition != null + && pkDefinition.getKind() == PartitionKind.MULTI_HASH + && partitionKey.getComponents().size() < pkDefinition.getPaths().size()) { + prefixEpkRange = partitionKey.getEPKRangeForPrefixPartitionKey(pkDefinition); + request.getHeaders().put( + HttpConstants.HttpHeaders.READ_FEED_KEY_TYPE, + ReadFeedKeyType.EffectivePartitionKeyRange.name()); + request.getHeaders().put(HttpConstants.HttpHeaders.START_EPK, prefixEpkRange.getMin()); + request.getHeaders().put(HttpConstants.HttpHeaders.END_EPK, prefixEpkRange.getMax()); + } + } + RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs(request); HttpHeaders headers = this.getHttpHeaders(); @@ -246,11 +278,18 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque RntbdRequest rntbdRequest = RntbdRequest.from(rntbdRequestArgs); - PartitionKeyInternal partitionKey = request.getPartitionKeyInternal(); - if (partitionKey != null) { - byte[] epk = partitionKey.getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition()); - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EffectivePartitionKey, epk); + if (prefixEpkRange != null) { + // Partial (prefix) hierarchical partition key. StartEpk/EndEpk + ReadFeedKeyType were + // already set on the request headers above and serialized into the RNTBD request as the + // backend's doc-level EPK filter. StartEpkHash/EndEpkHash additionally steer the proxy to + // the owning physical partition(s), mirroring the Direct/RNTBD path. + rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMin())); + rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMax())); + } else { + byte[] epk = partitionKey.getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition()); + rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EffectivePartitionKey, epk); + } } else if (request.requestContext.resolvedPartitionKeyRange == null) { throw new IllegalStateException( "Resolved partition key range should not be null at this point. ResourceType: " From 8df2e85d7e97aa3d4c373400ba0200851073cd38 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 10:05:27 -0400 Subject: [PATCH 02/15] Add thin-client query E2E parity regression tests Ports ThinClientQueryE2ETest and its ThinClientTestBase helper. The suite self-baselines every query against a Direct (RNTBD) client and the thin-client (:10250 HTTP/2) path, asserting identical results. This directly guards the MULTI_HASH prefix partition-key EPK over-span fix: prefix (single-component) hierarchical-partition-key queries must return only the narrow matching set, matching Direct, rather than over-spanning to all co-located documents. Both files are runtime self-enabling (enableThinClientForTest() + builder-driven HTTP/2 via setEnabled(true)); no module property or TestSuiteBase changes needed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../cosmos/rx/ThinClientQueryE2ETest.java | 2031 +++++++++++++++++ .../azure/cosmos/rx/ThinClientTestBase.java | 111 + 2 files changed, 2142 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java new file mode 100644 index 000000000000..4fcdcde14da7 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -0,0 +1,2031 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.rx; + +import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnostics; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosDiagnosticsRequestInfo; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.models.CompositePath; +import com.azure.cosmos.models.CompositePathSortOrder; +import com.azure.cosmos.models.CosmosBulkOperations; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosFullTextIndex; +import com.azure.cosmos.models.CosmosFullTextPath; +import com.azure.cosmos.models.CosmosFullTextPolicy; +import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosVectorDataType; +import com.azure.cosmos.models.CosmosVectorDistanceFunction; +import com.azure.cosmos.models.CosmosVectorEmbedding; +import com.azure.cosmos.models.CosmosVectorEmbeddingPolicy; +import com.azure.cosmos.models.CosmosVectorIndexSpec; +import com.azure.cosmos.models.CosmosVectorIndexType; +import com.azure.cosmos.models.ExcludedPath; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.IncludedPath; +import com.azure.cosmos.models.IndexingMode; +import com.azure.cosmos.models.IndexingPolicy; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyBuilder; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKeyDefinitionVersion; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.implementation.TestConfigurations; +import com.azure.cosmos.implementation.directconnectivity.Protocol; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; + +import static com.azure.cosmos.rx.ThinClientTestBase.assertThinClientEndpointUsed; +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.Fail.fail; + +/** + * Thin client query E2E tests comparing Direct TCP (baseline) vs Gateway V2 (thin client). + *

+ * Each query is executed through both connection modes and results are compared: + *

    + *
  • Direct TCP — baseline, runs against backend partition replicas.
    • + *
  • Gateway V2 (thin client) — system under test, routes through proxy (:10250), + * proxy returns raw PartitionKeyInternal arrays, SDK converts to EPK client-side.
    • + *
+ * Assertions: + *
    + *
  1. Gateway V2 requests routed through the :10250 thin client endpoint.
    2. + *
  2. Result set sizes match between Direct and Gateway V2.
    3. + *
  3. Result set contents match (document IDs in order for ordered queries, set equality for unordered).
    4. + *
+ */ +public class ThinClientQueryE2ETest extends TestSuiteBase { + + private CosmosAsyncClient directClient; // Baseline: Direct TCP + private CosmosAsyncClient thinClient; // System under test: Gateway V2 (thin client) + private CosmosAsyncContainer directContainer; + private CosmosAsyncContainer thinClientContainer; + + // Dedicated container whose documents reside in multiple physical partitions (distinct partition + // keys over a 12,000 RU/s container). Used by the cross-partition tests so they genuinely fan out. + private CosmosAsyncContainer directCrossPartitionContainer; + private CosmosAsyncContainer thinClientCrossPartitionContainer; + private String crossPartitionContainerId; + + // Dedicated container with a HIERARCHICAL (MULTI_HASH) partition key (/tenantId, /userId) over a + // 12,000 RU/s container, so its rows span multiple physical partitions. The thin-client (Gateway + // V2 proxy) emits this container's QueryPlan with MULTI-COMPONENT PartitionKeyInternal min/max + // arrays (one element per hierarchical path), which the client converts to EPK ranges through the + // MULTI_HASH branch of convertToSortedEpkRanges. Single-path containers never reach that branch, + // so this fixture is required for the hierarchical-PK coverage and half-open prefix-range + // coverage. + private CosmosAsyncContainer directHierarchicalContainer; + private CosmosAsyncContainer thinClientHierarchicalContainer; + private String hierarchicalContainerId; + private final List hierarchicalKeys = new ArrayList<>(); // {tenantId, userId} per seeded doc + private static final String TENANT_FIELD = "tenantId"; + private static final String USER_FIELD = "userId"; + private static final int HIER_TENANTS = 4; + private static final int HIER_USERS_PER_TENANT = 6; + + private final List seededDocs = new ArrayList<>(); + private final String commonPk = "tc-query-" + UUID.randomUUID().toString().substring(0, 8); + + // Use constants and helpers from ThinClientTestBase to avoid duplication. + private static final String ID_FIELD = ThinClientTestBase.ID_FIELD; + private static final String PK_FIELD = ThinClientTestBase.PARTITION_KEY_FIELD; + private static final ObjectMapper OBJECT_MAPPER = ThinClientTestBase.OBJECT_MAPPER; + + @BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT * 2) + public void before_ThinClientQueryE2ETest() { + try { + // 1. Direct TCP client (baseline) + CosmosClientBuilder directClientBuilder = createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.TCP, false, null, true, true); + this.directClient = directClientBuilder.buildAsyncClient(); + this.directContainer = getSharedMultiPartitionCosmosContainer(this.directClient); + + // 2. Gateway V2 thin client (system under test) + ThinClientTestBase.enableThinClientForTest(); + CosmosClientBuilder thinClientBuilder = createGatewayRxDocumentClient( + TestConfigurations.HOST, null, true, null, true, true, true); + this.thinClient = thinClientBuilder.buildAsyncClient(); + this.thinClientContainer = this.thinClient.getDatabase( + directContainer.getDatabase().getId()).getContainer(directContainer.getId()); + + // 3. Clean up shared container to prevent cross-test-class pollution + cleanUpContainer(this.directContainer); + + // 4. Seed diverse test data for broad query coverage + seedTestData(); + + // 5. Seed a dedicated multi-physical-partition container for genuine cross-partition coverage + seedCrossPartitionData(); + + // 6. Seed a dedicated hierarchical (MULTI_HASH) partition-key container for the + // multi-component QueryPlan range-conversion coverage. + seedHierarchicalData(); + } catch (Exception e) { + // Clean up any clients that were successfully created before the failure + if (this.thinClient != null) { this.thinClient.close(); this.thinClient = null; } + if (this.directClient != null) { this.directClient.close(); this.directClient = null; } + throw e; + } + } + + /** + * Seeds 10 documents into the shared container with the following schema: + * 
+     * {
+     *   "id":        "tcdoc-{i}-{uuid}",  // unique document ID
+     *   "mypk":      "{commonPk}",         // partition key — same for all seeded docs
+     *   "category":  string,               // one of: electronics, books, clothing, toys
+     *   "status":    string,               // "active" or "inactive"
+     *   "age":       int,                  // range: 8–61
+     *   "price":     double,               // range: 7.50–549.99
+     *   "idx":       int,                  // sequential index 0–9
+     *   "isActive":  boolean,              // derived from status == "active"
+     *   "address":   { "city": string, "zip": int },  // nested object
+     *   "scores":    [int, int],           // two-element int array: [i*10, i*10+5]
+     *   "tags":      [string, ...]         // variable-length string array
+     * }
+     *
+ */ + private void seedTestData() { + String[] categories = {"electronics", "books", "clothing", "electronics", "books", + "clothing", "electronics", "toys", "toys", "books"}; + String[] statuses = {"active", "inactive", "active", "active", "inactive", + "active", "inactive", "active", "active", "active"}; + int[] ages = {25, 30, 17, 42, 55, 19, 38, 12, 8, 61}; + double[] prices = {99.99, 14.50, 45.00, 299.99, 9.99, 25.00, 549.99, 19.99, 7.50, 22.00}; + + for (int i = 0; i < 10; i++) { + String docId = "tcdoc-" + i + "-" + UUID.randomUUID().toString().substring(0, 8); + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, docId); + doc.put(PK_FIELD, commonPk); + doc.put("category", categories[i]); + doc.put("status", statuses[i]); + doc.put("age", ages[i]); + doc.put("price", prices[i]); + doc.put("idx", i); + doc.put("isActive", statuses[i].equals("active")); + + ObjectNode address = OBJECT_MAPPER.createObjectNode(); + address.put("city", i % 2 == 0 ? "Seattle" : "Portland"); + address.put("zip", 98100 + i); + doc.set("address", address); + + doc.putArray("scores").add(i * 10).add(i * 10 + 5); + + // Tags array for JOIN/EXISTS tests — varies per doc + ArrayNode tags = doc.putArray("tags"); + tags.add(categories[i]); // first tag matches category + if (i % 2 == 0) tags.add("on-sale"); + if (i % 3 == 0) tags.add("featured"); + + seededDocs.add(doc); + } + + bulkInsert(directContainer, seededDocs).blockLast(); + } + + /** + * Seeds a dedicated container whose documents reside in multiple physical partitions. The + * container is provisioned at 12,000 RU/s - above the 10,000 RU/s single-partition limit, so the + * service splits it into multiple physical partitions - and each document is given a distinct + * partition key, spreading the rows across those partitions. Cross-partition queries against this + * container therefore genuinely fan out, unlike the single-partition shared fixture. + */ + private void seedCrossPartitionData() { + crossPartitionContainerId = "thinXp_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + CosmosContainerProperties props = new CosmosContainerProperties(crossPartitionContainerId, pkDef); + db.createContainer(props, ThroughputProperties.createManualThroughput(12000)).block(); + + this.directCrossPartitionContainer = db.getContainer(crossPartitionContainerId); + this.thinClientCrossPartitionContainer = + thinClient.getDatabase(db.getId()).getContainer(crossPartitionContainerId); + + String[] categories = {"electronics", "books", "clothing", "toys"}; + for (int i = 0; i < 30; i++) { + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "xp-" + i + "-" + UUID.randomUUID().toString().substring(0, 8)); + // Distinct partition key per document so the rows spread across physical partitions. + doc.put(PK_FIELD, "xppk-" + i + "-" + UUID.randomUUID()); + doc.put("category", categories[i % categories.length]); + doc.put("idx", i); // distinct -> ORDER BY c.idx is a total order + directCrossPartitionContainer.createItem(doc, new PartitionKey(doc.get(PK_FIELD).asText()), null).block(); + } + } + + /** + * Seeds a dedicated container with a HIERARCHICAL (MULTI_HASH) partition key (/tenantId, /userId), + * provisioned at 12,000 RU/s so the rows span multiple physical partitions. Exactly + * {@code HIER_TENANTS * HIER_USERS_PER_TENANT} documents are seeded, one per distinct + * (tenantId, userId) pair, with a distinct {@code idx} (so {@code ORDER BY c.idx} is a total + * order) and a cycled {@code category}. The thin-client proxy emits this container's QueryPlan + * with MULTI-COMPONENT PartitionKeyInternal min/max arrays, exercising the MULTI_HASH branch of + * convertToSortedEpkRanges that single-path containers never reach. + */ + private void seedHierarchicalData() { + hierarchicalContainerId = "thinHier_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setKind(PartitionKind.MULTI_HASH); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setPaths(Arrays.asList("/" + TENANT_FIELD, "/" + USER_FIELD)); + CosmosContainerProperties props = new CosmosContainerProperties(hierarchicalContainerId, pkDef); + db.createContainer(props, ThroughputProperties.createManualThroughput(12000)).block(); + + this.directHierarchicalContainer = db.getContainer(hierarchicalContainerId); + this.thinClientHierarchicalContainer = + thinClient.getDatabase(db.getId()).getContainer(hierarchicalContainerId); + + String[] categories = {"electronics", "books", "clothing", "toys"}; + int idx = 0; + for (int t = 0; t < HIER_TENANTS; t++) { + String tenantId = "tenant-" + t + "-" + UUID.randomUUID().toString().substring(0, 8); + for (int u = 0; u < HIER_USERS_PER_TENANT; u++) { + String userId = "user-" + u + "-" + UUID.randomUUID().toString().substring(0, 8); + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "hier-" + idx + "-" + UUID.randomUUID().toString().substring(0, 8)); + doc.put(TENANT_FIELD, tenantId); + doc.put(USER_FIELD, userId); + doc.put("category", categories[idx % categories.length]); + doc.put("idx", idx); // distinct -> ORDER BY c.idx is a total order + PartitionKey pk = new PartitionKeyBuilder().add(tenantId).add(userId).build(); + directHierarchicalContainer.createItem(doc, pk, null).block(); + hierarchicalKeys.add(new String[]{tenantId, userId}); + idx++; + } + } + } + + @AfterClass(groups = {"thinclient"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + if (directContainer != null && !seededDocs.isEmpty()) { + try { + List deleteOps = seededDocs.stream() + .map(doc -> CosmosBulkOperations.getDeleteItemOperation( + doc.get(ID_FIELD).asText(), new PartitionKey(commonPk))) + .collect(Collectors.toList()); + directContainer.executeBulkOperations(Flux.fromIterable(deleteOps)).blockLast(); + } catch (Exception e) { + logger.warn("Bulk delete of seeded docs failed: {}", e.getMessage()); + } + } + ThinClientTestBase.clearThinClientForTest(); + if (directCrossPartitionContainer != null) { + safeDeleteContainer(directCrossPartitionContainer); + } + if (directHierarchicalContainer != null) { + safeDeleteContainer(directHierarchicalContainer); + } + if (this.thinClient != null) { this.thinClient.close(); } + if (this.directClient != null) { this.directClient.close(); } + } + + // ==================== Equality & Filter Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSelectAll() { + assertDirectAndThinClientMatch("SELECT * FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereEquality() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category = 'electronics'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereEqualityParameterized() { + SqlQuerySpec qs = new SqlQuerySpec("SELECT * FROM c WHERE c.category = @cat"); + qs.setParameters(Arrays.asList(new SqlParameter("@cat", "books"))); + assertDirectAndThinClientMatch(qs, partitionedOptions()); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereRangeGreaterThan() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.age > 30"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereRangeLessThanOrEqual() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.price <= 25.00"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereRangeBetween() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.age >= 18 AND c.age <= 40"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereIn() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category IN ('electronics', 'toys')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereCompoundAndOr() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.status = 'active' AND (c.category = 'electronics' OR c.category = 'books')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereNotEqual() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.status != 'inactive'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereBooleanField() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.isActive = true"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereIsDefined() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE IS_DEFINED(c.address)"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereStartsWith() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE STARTSWITH(c.category, 'elec')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereContains() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE CONTAINS(c.category, 'ook')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereArrayContains() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE ARRAY_CONTAINS(c.scores, 50)"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testWhereNestedProperty() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.address.city = 'Seattle'"); + } + + // ==================== Projection Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSelectSpecificFields() { + String query = "SELECT c.id, c.category, c.price FROM c"; + QueryResult gwResult = drainQuery(directContainer, query, partitionedOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientContainer, query, partitionedOptions(), ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Count mismatch: " + query).isEqualTo(gwResult.results.size()); + for (int i = 0; i < gwResult.results.size(); i++) { + assertThat(tcResult.results.get(i).get("category").asText()).isEqualTo(gwResult.results.get(i).get("category").asText()); + assertThat(tcResult.results.get(i).get("price").asDouble()).isEqualTo(gwResult.results.get(i).get("price").asDouble()); + } + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSelectComputedAlias() { + String query = "SELECT c.id, c.price * 1.1 AS taxedPrice FROM c"; + QueryResult gwResult = drainQuery(directContainer, query, partitionedOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientContainer, query, partitionedOptions(), ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Count mismatch: " + query).isEqualTo(gwResult.results.size()); + } + + // ==================== ORDER BY Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testOrderByAsc() { + assertDirectAndThinClientMatch("SELECT * FROM c ORDER BY c.age"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testOrderByDesc() { + assertDirectAndThinClientMatch("SELECT * FROM c ORDER BY c.price DESC"); + } + + /** + * Multi-property ORDER BY (composite ordering). A multi-key ORDER BY can only be served when a + * matching composite index exists, so this test provisions a dedicated container with a + * (category ASC, age DESC) composite index, seeds documents on a single logical partition, and + * asserts the thin client returns the exact same strictly-ordered sequence as Direct. This + * exercises the multi-component ORDER BY composition / merge path, which single-key ORDER BY + * does not. Mirrors the multi-ORDER-BY coverage in the .NET SDK. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 2) + public void testMultipleOrderBy() { + String containerId = "multiOrderBy_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directTestContainer = db.getContainer(containerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + CosmosContainerProperties props = new CosmosContainerProperties(containerId, pkDef); + + // Multi-property ORDER BY requires a matching composite index. + IndexingPolicy idxPolicy = new IndexingPolicy(); + List composite = new ArrayList<>(); + composite.add(new CompositePath().setPath("/category").setOrder(CompositePathSortOrder.ASCENDING)); + composite.add(new CompositePath().setPath("/age").setOrder(CompositePathSortOrder.DESCENDING)); + idxPolicy.setCompositeIndexes(Collections.singletonList(composite)); + props.setIndexingPolicy(idxPolicy); + db.createContainer(props, ThroughputProperties.createManualThroughput(400)).block(); + + CosmosAsyncContainer tcContainer = thinClient.getDatabase(db.getId()).getContainer(containerId); + + String pk = "mob-" + UUID.randomUUID().toString().substring(0, 8); + String[] categories = {"alpha", "beta", "alpha", "gamma", "beta", "alpha"}; + int[] ages = {30, 25, 20, 40, 35, 50}; + for (int i = 0; i < categories.length; i++) { + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "mob-" + i + "-" + UUID.randomUUID().toString().substring(0, 8)); + doc.put(PK_FIELD, pk); + doc.put("category", categories[i]); + doc.put("age", ages[i]); + directTestContainer.createItem(doc, new PartitionKey(pk), null).block(); + } + + String query = "SELECT * FROM c ORDER BY c.category ASC, c.age DESC"; + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions().setPartitionKey(new PartitionKey(pk)); + + QueryResult directResult = drainQuery(directTestContainer, query, options, ObjectNode.class); + QueryResult tcResult = drainQuery(tcContainer, query, options, ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Multi-ORDER-BY count mismatch: " + query).isEqualTo(directResult.results.size()); + assertThat(tcResult.results.size()).isEqualTo(categories.length); + // Strict sequence parity — the multi-key ORDER BY ordering must be identical to Direct. + assertSameDocumentIds(directResult.results, tcResult.results, true, query); + } finally { + safeDeleteContainer(directTestContainer); + } + } + + // ==================== DISTINCT Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testDistinctValue() { + assertScalarDirectAndThinClientMatch("SELECT DISTINCT VALUE c.category FROM c", String.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testDistinctValueBoolean() { + assertScalarDirectAndThinClientMatch("SELECT DISTINCT VALUE c.isActive FROM c", Boolean.class); + } + + // ==================== TOP Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testTop() { + assertDirectAndThinClientMatch("SELECT TOP 3 * FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testTopWithOrderBy() { + assertDirectAndThinClientMatch("SELECT TOP 5 * FROM c ORDER BY c.price DESC"); + } + + // ==================== Aggregate Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCount() { + assertScalarDirectAndThinClientMatch("SELECT VALUE COUNT(1) FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSum() { + assertScalarDirectAndThinClientMatch("SELECT VALUE SUM(c.price) FROM c", Double.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testAvg() { + assertScalarDirectAndThinClientMatch("SELECT VALUE AVG(c.age) FROM c", Double.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMin() { + assertScalarDirectAndThinClientMatch("SELECT VALUE MIN(c.price) FROM c", Double.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMax() { + assertScalarDirectAndThinClientMatch("SELECT VALUE MAX(c.age) FROM c", Integer.class); + } + + /** + * DCount (distinct count). Validates the thin client serves the distinct-count aggregate identically + * to Direct. Cosmos SQL does not accept the standard {@code COUNT(DISTINCT ...)} form; the canonical + * DCount idiom is {@code COUNT(1)} over a {@code SELECT DISTINCT VALUE} subquery, which the query + * engine folds into a DCount aggregate. DCount is a distinct query feature that must be advertised in + * SupportedQueryFeatures for the proxy-generated query plan; this guards that negotiation path. + * The .NET SDK has equivalent DCount coverage. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testDCount() { + assertScalarDirectAndThinClientMatch( + "SELECT VALUE COUNT(1) FROM (SELECT DISTINCT VALUE c.category FROM c) AS t", Integer.class); + } + + // ==================== GROUP BY Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testGroupByCount() { + assertGroupByDirectAndThinClientMatch("SELECT c.category, COUNT(1) as cnt FROM c GROUP BY c.category", "category"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testGroupBySumAvg() { + assertGroupByDirectAndThinClientMatch("SELECT c.category, SUM(c.price) as total, AVG(c.price) as avg FROM c GROUP BY c.category", "category"); + } + + // ==================== OFFSET / LIMIT Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testOffsetLimit() { + assertDirectAndThinClientMatch("SELECT * FROM c ORDER BY c.idx OFFSET 3 LIMIT 4"); + } + + // ==================== JOIN Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testJoinScoresArray() { + // Self-join on scores array — produces one row per array element + assertDirectAndThinClientMatch("SELECT c.id, s AS score FROM c JOIN s IN c.scores"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testJoinWithFilter() { + // Self-join with WHERE filter on the joined element + assertDirectAndThinClientMatch("SELECT c.id, s AS score FROM c JOIN s IN c.scores WHERE s >= 50"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testJoinTagsArray() { + // Self-join on tags string array + assertDirectAndThinClientMatch("SELECT c.id, t AS tag FROM c JOIN t IN c.tags"); + } + + // ==================== EXISTS Subquery Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testExistsSubquery() { + // Docs pattern: use EXISTS to check if any array element matches + assertDirectAndThinClientMatch( + "SELECT * FROM c WHERE EXISTS (SELECT VALUE s FROM s IN c.scores WHERE s > 60)"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testExistsSubqueryWithStringMatch() { + // EXISTS on tags array with string match + assertDirectAndThinClientMatch( + "SELECT * FROM c WHERE EXISTS (SELECT VALUE t FROM t IN c.tags WHERE t = 'on-sale')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testExistsAliasInProjection() { + // EXISTS aliased in SELECT — returns boolean column + assertDirectAndThinClientMatch( + "SELECT c.id, EXISTS (SELECT VALUE s FROM s IN c.scores WHERE s > 60) AS hasHighScore FROM c"); + } + + // ==================== LIKE Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikePrefix() { + // LIKE with prefix pattern + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE 'elec%'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikeSuffix() { + // LIKE with suffix pattern + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE '%ing'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikeContains() { + // LIKE with contains pattern (substring match via wildcards) + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE '%ook%'"); + } + + // ==================== BETWEEN Keyword ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testBetween() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.age BETWEEN 18 AND 40"); + } + + // ==================== String Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringConcat() { + assertDirectAndThinClientMatch("SELECT CONCAT(c.category, '-', c.status) AS label FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringEndsWith() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE ENDSWITH(c.category, 'ics')"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringLower() { + assertDirectAndThinClientMatch("SELECT LOWER(c.category) AS lowerCat FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringUpper() { + assertDirectAndThinClientMatch("SELECT UPPER(c.status) AS upperStatus FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringLength() { + assertDirectAndThinClientMatch("SELECT c.category, LENGTH(c.category) AS len FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringSubstring() { + assertDirectAndThinClientMatch("SELECT SUBSTRING(c.category, 0, 4) AS prefix FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringReplace() { + assertDirectAndThinClientMatch("SELECT REPLACE(c.category, 'o', '0') AS replaced FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringIndexOf() { + assertDirectAndThinClientMatch("SELECT INDEX_OF(c.category, 'o') AS pos FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringLeft() { + assertDirectAndThinClientMatch("SELECT LEFT(c.category, 3) AS l FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringReverse() { + assertDirectAndThinClientMatch("SELECT REVERSE(c.category) AS rev FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStringTrim() { + assertDirectAndThinClientMatch("SELECT TRIM(c.status) AS trimmed FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testRegexMatch() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE RegexMatch(c.category, '^elec.*')"); + } + + // ==================== Type Checking Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsArray() { + assertDirectAndThinClientMatch("SELECT c.id, IS_ARRAY(c.scores) AS isArr FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsBool() { + assertDirectAndThinClientMatch("SELECT c.id, IS_BOOL(c.isActive) AS isBool FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsNull() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE IS_NULL(c.nonExistentField)"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsNumber() { + assertDirectAndThinClientMatch("SELECT c.id, IS_NUMBER(c.age) AS isNum FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsString() { + assertDirectAndThinClientMatch("SELECT c.id, IS_STRING(c.category) AS isStr FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIsObject() { + assertDirectAndThinClientMatch("SELECT c.id, IS_OBJECT(c.address) AS isObj FROM c"); + } + + // ==================== Math Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMathAbs() { + assertDirectAndThinClientMatch("SELECT ABS(c.age - 30) AS diff FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMathCeilingFloor() { + assertDirectAndThinClientMatch("SELECT CEILING(c.price) AS ceil, FLOOR(c.price) AS flr FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMathRound() { + assertDirectAndThinClientMatch("SELECT ROUND(c.price) AS rounded FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMathPower() { + assertDirectAndThinClientMatch("SELECT POWER(c.age, 2) AS ageSq FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testMathSqrt() { + assertDirectAndThinClientMatch("SELECT SQRT(c.price) AS sqrtPrice FROM c"); + } + + // ==================== Array Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testArrayLength() { + assertDirectAndThinClientMatch("SELECT c.id, ARRAY_LENGTH(c.scores) AS len FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testArraySlice() { + assertDirectAndThinClientMatch("SELECT c.id, ARRAY_SLICE(c.tags, 0, 1) AS firstTag FROM c"); + } + + // ==================== Conditional Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testIif() { + assertDirectAndThinClientMatch("SELECT c.id, IIF(c.age >= 18, 'adult', 'minor') AS ageGroup FROM c"); + } + + // ==================== Date/Time Function Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testGetCurrentDateTime() { + // Only assert both paths return a non-empty ISO 8601 string — exact values + // will differ because gateway and proxy execute at slightly different times. + QueryResult gwResult = drainQuery(directContainer, + "SELECT VALUE GetCurrentDateTime()", partitionedOptions(), String.class); + QueryResult tcResult = drainQuery(thinClientContainer, + "SELECT VALUE GetCurrentDateTime()", partitionedOptions(), String.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + assertThat(gwResult.results.size()).isEqualTo(1); + assertThat(tcResult.results.size()).isEqualTo(1); + assertThat(gwResult.results.get(0)).matches("\\d{4}-\\d{2}-\\d{2}T.*Z"); + assertThat(tcResult.results.get(0)).matches("\\d{4}-\\d{2}-\\d{2}T.*Z"); + } + + // ==================== SELECT VALUE / Nested Projection Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSelectValueObject() { + assertDirectAndThinClientMatch( + "SELECT VALUE { name: c.category, loc: c.address.city } FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testSelectValueScalar() { + assertScalarDirectAndThinClientMatch("SELECT VALUE c.category FROM c", String.class); + } + + // ==================== QueryOracle-derived coverage ==================== + // The CosmosDB QueryOracle (CloudTest "Release-QueryOracle") fuzzes queries from category + // configs (Like, ScalarExpression, Builtin, Aggregate, ...). These tests bring the LIKE and + // scalar-expression categories that were not yet exercised above into the Direct-vs-thin-client + // oracle. Geospatial (ST_DISTANCE/ST_WITHIN) and UDF categories are intentionally NOT covered + // here because they require a geospatial-indexed container and registered user-defined + // functions respectively, which the shared seeded fixture does not provide. + + // ---- LIKE: advanced patterns (character classes, single-char wildcard, negation) ---- + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikeSingleCharWildcard() { + // '_' matches exactly one character: 'electronic_' matches "electronics". + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE 'electronic_'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikeCharacterClassRange() { + // Character range '[a-c]' — categories starting with a..c (books, clothing). + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE '[a-c]%'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testLikeNegatedCharacterClass() { + // Negated character class '[^bc]' — categories NOT starting with b or c (electronics, toys). + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category LIKE '[^bc]%'"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testNotLike() { + assertDirectAndThinClientMatch("SELECT * FROM c WHERE c.category NOT LIKE 'elec%'"); + } + + // ---- Scalar expressions: coalesce, computed member access, array literal, unary, modulo, ternary ---- + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCoalesceOperator() { + // '??' returns the right operand when the left is undefined. + assertScalarDirectAndThinClientMatch("SELECT VALUE (c.missingField ?? c.category) FROM c", String.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testComputedMemberIndexer() { + // Quoted/computed property accessor c["category"]. + assertScalarDirectAndThinClientMatch("SELECT VALUE c[\"category\"] FROM c", String.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testArrayLiteralProjection() { + // Array-create scalar expression in the projection. + assertDirectAndThinClientMatch("SELECT c.id, [c.age, c.idx] AS pair FROM c"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testUnaryNegation() { + assertScalarDirectAndThinClientMatch("SELECT VALUE -c.age FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testModuloOperator() { + assertScalarDirectAndThinClientMatch("SELECT VALUE (c.age % 2) FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testTernaryConditional() { + // Ternary '?:' — distinct parse path from IIF(...). + assertScalarDirectAndThinClientMatch( + "SELECT VALUE (c.age >= 18 ? 'adult' : 'minor') FROM c", String.class); + } + + // ==================== Query Plan Caching Tests ==================== + // A query scoped to a single logical partition (partition key set on the options) caches the + // proxy-generated query plan on the client and reuses it on subsequent identical queries. This + // test verifies that the cached plan still executes correctly - i.e. a cached QueryPlan obtained + // from the thin-client proxy continues to produce results identical to the Direct (TCP) baseline. + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCachedQueryPlanFromProxyExecutesCorrectly() { + // GROUP BY requires a query plan. With a partition key set (assertGroupByDirectAndThinClientMatch + // uses partitionedOptions()), the proxy-generated plan is cached after the first execution. + String query = "SELECT c.category, COUNT(1) AS cachedPlanCount FROM c GROUP BY c.category"; + + // First execution fetches and caches the query plan generated by the thin-client proxy. + assertGroupByDirectAndThinClientMatch(query, "category"); + + // Second execution reuses the cached plan; assert it still matches the Direct baseline. + assertGroupByDirectAndThinClientMatch(query, "category"); + } + + // ==================== Cross-Partition Tests ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionSelectAll() { + assertCrossPartitionDirectAndThinClientMatch("SELECT * FROM c ORDER BY c.idx"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionWhereFilter() { + assertCrossPartitionDirectAndThinClientMatch( + "SELECT * FROM c WHERE c.category = 'electronics' ORDER BY c.idx"); + } + + // ============ Cross-Partition Aggregate / GROUP BY Merge Tests ============ + // These run aggregate and GROUP BY queries with NO partition key against the dedicated + // multi-physical-partition container, so the thin client must fan out to every physical partition + // and MERGE the partial results into the final value/groups. This is distinct from the + // single-partition aggregate/GROUP BY tests above (which set a partition key and never merge + // across partitions). Each helper also asserts the Direct baseline genuinely contacted more than + // one partition key range, proving the cross-partition merge path actually ran. + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionCount() { + assertCrossPartitionScalarMatch("SELECT VALUE COUNT(1) FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionSum() { + assertCrossPartitionScalarMatch("SELECT VALUE SUM(c.idx) FROM c", Double.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionAvg() { + assertCrossPartitionScalarMatch("SELECT VALUE AVG(c.idx) FROM c", Double.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionMin() { + assertCrossPartitionScalarMatch("SELECT VALUE MIN(c.idx) FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionMax() { + assertCrossPartitionScalarMatch("SELECT VALUE MAX(c.idx) FROM c", Integer.class); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionGroupByCount() { + assertCrossPartitionGroupByMatch( + "SELECT c.category, COUNT(1) AS cnt FROM c GROUP BY c.category", "category"); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testCrossPartitionGroupBySumAvg() { + assertCrossPartitionGroupByMatch( + "SELECT c.category, SUM(c.idx) AS total, AVG(c.idx) AS avg FROM c GROUP BY c.category", "category"); + } + + // ============ QueryPlan Range-Conversion Boundary Tests ============ + // These tests target the thin-client (Gateway V2 proxy) QueryPlan emission, where the + // PartitionedQueryExecutionInfo carries queryRanges as PartitionKeyInternal min/max JSON arrays + // (e.g. {"min":[],"max":["Infinity"]}) rather than the Gateway V1 EPK hex strings. The client + // converts those arrays to EPK ranges via PartitionKeyInternalHelper.convertToSortedEpkRanges: + // - an empty min array -> MinimumInclusiveEffectivePartitionKey ("") + // - an "Infinity" max -> MaximumExclusiveEffectivePartitionKey ("FF") + // - a MULTI_HASH key -> the multi-component branch of the conversion + // We assert result parity against the Direct-TCP baseline (whose plan comes from Gateway V1), + // proving both QueryPlan sources resolve to equivalent ranges. + + /** + * Full-range / Infinity boundary. {@code SELECT * FROM c} with no partition key forces the + * proxy to emit the widest possible range (empty min -> "", "Infinity" max -> "FF"). Asserts the + * cross-partition container returns ALL seeded rows on both paths (absolute count, distinct from + * {@link #testCrossPartitionSelectAll} which only checks parity) and that the fan-out genuinely + * spans more than one partition key range. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testFullRangeInfinityBoundary() { + int expected = 30; // seedCrossPartitionData seeds exactly 30 docs + CosmosQueryRequestOptions crossPartitionOptions = new CosmosQueryRequestOptions(); + QueryResult directResult = + drainQuery(directCrossPartitionContainer, "SELECT * FROM c", crossPartitionOptions, ObjectNode.class); + QueryResult tcResult = + drainQuery(thinClientCrossPartitionContainer, "SELECT * FROM c", crossPartitionOptions, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(directResult.results.size()).as("Direct full-range count").isEqualTo(expected); + assertThat(tcResult.results.size()).as("Thin-client full-range count").isEqualTo(expected); + assertSameDocumentIds(directResult.results, tcResult.results, false, "full-range SELECT *"); + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("full-range query must contact more than one partition key range") + .isGreaterThan(1); + } + + /** + * Hierarchical (MULTI_HASH) partition key cross-partition parity. Cross-partition queries + * over the /tenantId,/userId container force the proxy to emit MULTI-COMPONENT PartitionKeyInternal + * arrays, exercising the MULTI_HASH branch of convertToSortedEpkRanges. Validates plain, ORDER BY, + * and filtered-ORDER-BY shapes against the Direct baseline. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testHierarchicalPartitionKeyCrossPartition() { + assertHierarchicalCrossPartitionMatch("SELECT * FROM c"); + assertHierarchicalCrossPartitionMatch("SELECT * FROM c ORDER BY c.idx"); + assertHierarchicalCrossPartitionMatch("SELECT * FROM c WHERE c.category = 'books' ORDER BY c.idx"); + } + + /** + * Hierarchical prefix (half-open) range. Setting only the FIRST hierarchical component + * (tenantId) on the options makes the proxy emit a half-open prefix range whose max is the + * "Infinity"-suffixed sibling of the min - the prefix-range boundary case of the conversion. The + * full key (both components) instead resolves to a single point range. Asserts the prefix query + * returns exactly the tenant's users and the full key returns exactly one document, with parity + * to the Direct baseline. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testHierarchicalPrefixHalfOpenRange() { + String[] firstKey = hierarchicalKeys.get(0); + String tenant0 = firstKey[0]; + String user0 = firstKey[1]; + + CosmosQueryRequestOptions fullKeyOptions = new CosmosQueryRequestOptions(); + fullKeyOptions.setPartitionKey(new PartitionKeyBuilder().add(tenant0).add(user0).build()); + assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", fullKeyOptions, 1); + + CosmosQueryRequestOptions prefixOptions = new CosmosQueryRequestOptions(); + prefixOptions.setPartitionKey(new PartitionKeyBuilder().add(tenant0).build()); + assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", prefixOptions, HIER_USERS_PER_TENANT); + } + + /** + * Two-source QueryPlan parity. The same ORDER BY projection query is resolved through both + * QueryPlan sources: the single-path cross-partition container (proxy emits single-component + * arrays) and the hierarchical container (proxy emits multi-component arrays). Both must match + * the Gateway V1 Direct baseline, formalizing the invariant that the two emission formats are + * interchangeable from the client's perspective. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testTwoSourceQueryPlanParity() { + String q = "SELECT c.id, c.idx, c.category FROM c WHERE c.idx >= 0 ORDER BY c.idx"; + assertCrossPartitionDirectAndThinClientMatch(q); + assertHierarchicalCrossPartitionMatch(q); + } + + // ==================== Multi-EPK-Range Tests (Sort Validation) ==================== + // These tests use a dedicated higher-throughput (24,000 RU/s) container so that metadata + // exposes multiple EPK (effective partition key) ranges. Note: on the emulator / single-box + // backend these ranges are typically served by a single physical process (no true multi-box + // partitioning), so this does not guarantee documents physically land on separate partitions. + // It does, however, exercise the SDK-side routing and sort pipeline: documents with different + // partition keys map to different EPK ranges, and after PartitionKeyInternal → EPK hash + // conversion, the sort in parseQueryRangesForThinClient() ensures + // RoutingMapProviderHelper.getOverlappingRanges() doesn't throw IllegalArgumentException for + // unsorted ranges. + + /** + * Helper: creates a higher-throughput (24K RU) container exposing multiple EPK ranges, runs the + * test, deletes the container. + */ + private void runMultiRangeTest(String[] pkValues, String queryTemplate, int expectedCount, + boolean assertMultipleRangesContacted) { + String containerId = "multiRange_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directTestContainer = db.getContainer(containerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + CosmosContainerProperties props = new CosmosContainerProperties(containerId, pkDef); + db.createContainer(props, ThroughputProperties.createManualThroughput(24000)).block(); + + CosmosAsyncContainer tcContainer = thinClient.getDatabase(db.getId()).getContainer(containerId); + + for (int i = 0; i < pkValues.length; i++) { + String docId = "mr-" + i + "-" + UUID.randomUUID().toString().substring(0, 8); + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, docId); + doc.put(PK_FIELD, pkValues[i]); + doc.put("idx", i); + doc.put("val", i * 100); + directTestContainer.createItem(doc, new PartitionKey(pkValues[i]), null).block(); + } + + String query = queryTemplate; + + QueryResult directResult = drainQuery(directTestContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(tcContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Multi-range count mismatch for: " + query).isEqualTo(directResult.results.size()); + assertThat(tcResult.results.size()).isEqualTo(expectedCount); + + List directIds = directResult.results.stream().map(d -> d.get(ID_FIELD).asText()).sorted().collect(Collectors.toList()); + List tcIds = tcResult.results.stream().map(d -> d.get(ID_FIELD).asText()).sorted().collect(Collectors.toList()); + assertThat(tcIds).isEqualTo(directIds); + + if (assertMultipleRangesContacted) { + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("multi-range query must contact more than one partition key range: " + query) + .isGreaterThan(1); + } + + } finally { + safeDeleteContainer(directTestContainer); + } + } + + /** + * Test: IN clause on partition key with 3 values → 3 disjoint EPK ranges across 3 physical partitions. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 3) + public void testMultiRangePartitionKeyInClause() { + String[] pkValues = {"pk-alpha", "pk-beta", "pk-gamma", "pk-delta", "pk-epsilon"}; + runMultiRangeTest(pkValues, + "SELECT * FROM c WHERE c.mypk IN ('pk-alpha', 'pk-gamma', 'pk-epsilon')", + 3, false); + } + + /** + * Test: OR on partition key values → 2 disjoint EPK ranges. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 3) + public void testMultiRangePartitionKeyOrClause() { + String[] pkValues = {"pk-or-1", "pk-or-2", "pk-or-3"}; + runMultiRangeTest(pkValues, + "SELECT * FROM c WHERE c.mypk = 'pk-or-1' OR c.mypk = 'pk-or-3'", + 2, false); + } + + /** + * Test: IN clause with 10 PK values → 10 disjoint EPK ranges, stress test for sort correctness. + * Uses UUID-based PK values to maximize EPK hash spread. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 3) + public void testMultiRangeManyPartitionKeys() { + String[] pkValues = new String[10]; + for (int i = 0; i < 10; i++) { + pkValues[i] = "pk-many-" + UUID.randomUUID().toString(); + } + + // Build IN clause dynamically from the random PK values + StringBuilder sb = new StringBuilder("SELECT * FROM c WHERE c.mypk IN ("); + for (int i = 0; i < pkValues.length; i++) { + if (i > 0) sb.append(", "); + sb.append("'").append(pkValues[i]).append("'"); + } + sb.append(")"); + runMultiRangeTest(pkValues, sb.toString(), 10, true); + } + + // ==================== Continuation Token Draining ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testContinuationTokenDraining() { + // Drain gateway fully for expected count + QueryResult gwResult = drainQuery(directContainer, "SELECT * FROM c", partitionedOptions(), ObjectNode.class); + + // Drain thin client with small page size to force multiple continuations + List tcAll = new ArrayList<>(); + List tcDiag = new ArrayList<>(); + String continuationToken = null; + int pageCount = 0; + int maxIterations = 100; + do { + Iterable> pages = thinClientContainer + .queryItems("SELECT * FROM c", partitionedOptions(), ObjectNode.class) + .byPage(continuationToken, 3) // small page size + .toIterable(); + for (FeedResponse page : pages) { + tcAll.addAll(page.getResults()); + tcDiag.add(page.getCosmosDiagnostics()); + continuationToken = page.getContinuationToken(); + pageCount++; + } + } while (continuationToken != null && --maxIterations > 0); + + for (CosmosDiagnostics d : tcDiag) { assertThinClientEndpointUsed(d); } + assertThat(pageCount).as("Should have multiple pages with page size 3").isGreaterThan(1); + assertThat(tcAll.size()).as("Continuation draining count mismatch").isEqualTo(gwResult.results.size()); + + List tcDrainedIds = idsSorted(tcAll); + List gwDrainedIds = idsSorted(gwResult.results); + assertThat(tcDrainedIds).as("Continuation draining ID-set mismatch").isEqualTo(gwDrainedIds); + assertThat(new HashSet<>(tcDrainedIds).size()) + .as("Duplicate IDs encountered across drained continuation pages") + .isEqualTo(tcDrainedIds.size()); + } + + // ==================== Invalid Query ==================== + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testInvalidQueryReturnsBadRequest() { + try { + thinClientContainer.queryItems("SELEC * FORM c", new CosmosQueryRequestOptions(), ObjectNode.class) + .byPage().blockFirst(); + fail("Expected exception for invalid query"); + } catch (CosmosException e) { + assertThat(e.getStatusCode()) + .as("Invalid query must return 400 Bad Request (statusCode 0 indicates the thin-client " + + "decode regression), got " + e.getStatusCode()) + .isEqualTo(400); + assertThinClientEndpointUsed(e.getDiagnostics()); + logger.info("Expected error for invalid query: {} (status {})", e.getMessage(), e.getStatusCode()); + } + } + + // ==================== Vector Search ==================== + + /** + * Creates a vector-enabled container, runs VectorDistance query through both + * Direct and thin client, compares results. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 2) + public void testVectorSearchGatewayVsThinClient() { + String vectorContainerId = "vecCompare_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directVecContainer = db.getContainer(vectorContainerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + + CosmosContainerProperties props = new CosmosContainerProperties(vectorContainerId, pkDef); + + CosmosVectorEmbeddingPolicy policy = new CosmosVectorEmbeddingPolicy(); + CosmosVectorEmbedding emb = new CosmosVectorEmbedding(); + emb.setPath("/embedding"); + emb.setDataType(CosmosVectorDataType.FLOAT32); + emb.setEmbeddingDimensions(3); + emb.setDistanceFunction(CosmosVectorDistanceFunction.COSINE); + policy.setCosmosVectorEmbeddings(Collections.singletonList(emb)); + props.setVectorEmbeddingPolicy(policy); + + IndexingPolicy idxPolicy = new IndexingPolicy(); + idxPolicy.setIndexingMode(IndexingMode.CONSISTENT); + idxPolicy.setIncludedPaths(Collections.singletonList(new IncludedPath("/*"))); + idxPolicy.setExcludedPaths(Arrays.asList(new ExcludedPath("/embedding/*"), new ExcludedPath("/\"_etag\"/?"))); + CosmosVectorIndexSpec vecIdx = new CosmosVectorIndexSpec(); + vecIdx.setPath("/embedding"); + vecIdx.setType(CosmosVectorIndexType.FLAT.toString()); + idxPolicy.setVectorIndexes(Collections.singletonList(vecIdx)); + props.setIndexingPolicy(idxPolicy); + + db.createContainer(props).block(); + CosmosAsyncContainer tcVecContainer = thinClient.getDatabase(db.getId()).getContainer(vectorContainerId); + + double[][] embeddings = { + {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0}, + {1.0, 1.0, 0.0}, {0.9, 0.1, 0.0}, + }; + + String vecPk = UUID.randomUUID().toString(); + List docIds = new ArrayList<>(); + for (int i = 0; i < embeddings.length; i++) { + String docId = "vec_" + i + "_" + UUID.randomUUID().toString().substring(0, 8); + docIds.add(docId); + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, docId); + doc.put(PK_FIELD, vecPk); + doc.put("text", "document " + i); + ArrayNode arr = doc.putArray("embedding"); + for (double v : embeddings[i]) { arr.add(v); } + directVecContainer.createItem(doc, new PartitionKey(vecPk), null).block(); + } + + String query = "SELECT TOP 5 c.id, c.text, VectorDistance(c.embedding, [1.0, 0.0, 0.0]) AS score " + + "FROM c ORDER BY VectorDistance(c.embedding, [1.0, 0.0, 0.0])"; + + QueryResult directResult = drainQuery(directVecContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(tcVecContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).isEqualTo(directResult.results.size()); + assertThat(tcResult.results.size()).isEqualTo(5); + + for (int i = 0; i < directResult.results.size(); i++) { + assertThat(tcResult.results.get(i).get("id").asText()).isEqualTo(directResult.results.get(i).get("id").asText()); + } + + assertThat(tcResult.results.get(0).get("id").asText()).isEqualTo(docIds.get(0)); + assertThat(tcResult.results.get(0).get("score").asDouble()).isGreaterThan(0.99); + + // Euclidean variant — validates ORDER BY score semantics with a different distance function + String euclideanQuery = "SELECT TOP 5 c.id, VectorDistance(c.embedding, [1.0, 0.0, 0.0], false, {'distanceFunction':'euclidean'}) AS score " + + "FROM c ORDER BY VectorDistance(c.embedding, [1.0, 0.0, 0.0], false, {'distanceFunction':'euclidean'})"; + + QueryResult directEuclidean = drainQuery(directVecContainer, euclideanQuery, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcEuclidean = drainQuery(tcVecContainer, euclideanQuery, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcEuclidean.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcEuclidean.results.size()).isEqualTo(directEuclidean.results.size()); + assertThat(tcEuclidean.results.size()).isEqualTo(5); + + for (int i = 0; i < directEuclidean.results.size(); i++) { + assertThat(tcEuclidean.results.get(i).get("id").asText()) + .as("Euclidean vector search result mismatch at position " + i) + .isEqualTo(directEuclidean.results.get(i).get("id").asText()); + } + + } finally { + safeDeleteContainer(directVecContainer); + } + } + + // ==================== Full-Text Search ==================== + + /** + * Creates a container with full-text policy and index, runs FullTextContains query. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 2) + public void testFullTextSearchGatewayVsThinClient() { + String containerId = "ftsCompare_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directFtsContainer = db.getContainer(containerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + + CosmosContainerProperties props = new CosmosContainerProperties(containerId, pkDef); + + CosmosFullTextPath ftPath = new CosmosFullTextPath(); + ftPath.setPath("/text"); + ftPath.setLanguage("en-US"); + CosmosFullTextPolicy ftPolicy = new CosmosFullTextPolicy(); + ftPolicy.setDefaultLanguage("en-US"); + ftPolicy.setPaths(Collections.singletonList(ftPath)); + props.setFullTextPolicy(ftPolicy); + + IndexingPolicy idxPolicy = new IndexingPolicy(); + idxPolicy.setIndexingMode(IndexingMode.CONSISTENT); + idxPolicy.setIncludedPaths(Collections.singletonList(new IncludedPath("/*"))); + idxPolicy.setExcludedPaths(Collections.singletonList(new ExcludedPath("/\"_etag\"/?"))); + CosmosFullTextIndex ftIndex = new CosmosFullTextIndex(); + ftIndex.setPath("/text"); + idxPolicy.setCosmosFullTextIndexes(Collections.singletonList(ftIndex)); + props.setIndexingPolicy(idxPolicy); + + db.createContainer(props).block(); + CosmosAsyncContainer tcFtsContainer = thinClient.getDatabase(db.getId()).getContainer(containerId); + + String ftsPk = UUID.randomUUID().toString(); + String[] texts = { + "The quick brown fox jumps over the lazy dog", + "A red bicycle parked near the mountain trail", + "Electronic devices on sale at the downtown store", + "Mountain biking trails with scenic views", + "The lazy cat sleeps on the warm brown couch" + }; + for (int i = 0; i < texts.length; i++) { + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "fts_" + i + "_" + UUID.randomUUID().toString().substring(0, 8)); + doc.put(PK_FIELD, ftsPk); + doc.put("text", texts[i]); + directFtsContainer.createItem(doc, new PartitionKey(ftsPk), null).block(); + } + + String query = "SELECT TOP 10 * FROM c WHERE FullTextContains(c.text, 'mountain')"; + + QueryResult directResult = drainQuery(directFtsContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(tcFtsContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + assertThat(directResult.results.size()).as("Full-text query should return results").isPositive(); + assertThat(tcResult.results.size()).isEqualTo(directResult.results.size()); + + List directIds = directResult.results.stream().map(d -> d.get("id").asText()).sorted().collect(Collectors.toList()); + List tcIds = tcResult.results.stream().map(d -> d.get("id").asText()).sorted().collect(Collectors.toList()); + assertThat(tcIds).isEqualTo(directIds); + + } finally { + safeDeleteContainer(directFtsContainer); + } + } + + /** + * Creates a container with full-text policy and index, runs ORDER BY RANK FullTextScore query. + * Compares exact ordering between Direct and thin client. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 2) + public void testFullTextScoreRanking() { + String containerId = "ftsRank_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directFtsContainer = db.getContainer(containerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + + CosmosContainerProperties props = new CosmosContainerProperties(containerId, pkDef); + + CosmosFullTextPath ftPath = new CosmosFullTextPath(); + ftPath.setPath("/text"); + ftPath.setLanguage("en-US"); + CosmosFullTextPolicy ftPolicy = new CosmosFullTextPolicy(); + ftPolicy.setDefaultLanguage("en-US"); + ftPolicy.setPaths(Collections.singletonList(ftPath)); + props.setFullTextPolicy(ftPolicy); + + IndexingPolicy idxPolicy = new IndexingPolicy(); + idxPolicy.setIndexingMode(IndexingMode.CONSISTENT); + idxPolicy.setIncludedPaths(Collections.singletonList(new IncludedPath("/*"))); + idxPolicy.setExcludedPaths(Collections.singletonList(new ExcludedPath("/\"_etag\"/?"))); + CosmosFullTextIndex ftIndex = new CosmosFullTextIndex(); + ftIndex.setPath("/text"); + idxPolicy.setCosmosFullTextIndexes(Collections.singletonList(ftIndex)); + props.setIndexingPolicy(idxPolicy); + + db.createContainer(props).block(); + CosmosAsyncContainer tcFtsContainer = thinClient.getDatabase(db.getId()).getContainer(containerId); + + String ftsPk = UUID.randomUUID().toString(); + String[] texts = { + "The quick brown fox jumps over the lazy dog", + "A red bicycle parked near the mountain trail", + "Electronic devices on sale at the downtown store", + "Mountain biking trails with scenic views", + "The lazy cat sleeps on the warm brown couch" + }; + for (int i = 0; i < texts.length; i++) { + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "ftsRank_" + i + "_" + UUID.randomUUID().toString().substring(0, 8)); + doc.put(PK_FIELD, ftsPk); + doc.put("text", texts[i]); + directFtsContainer.createItem(doc, new PartitionKey(ftsPk), null).block(); + } + + String query = "SELECT TOP 5 * FROM c ORDER BY RANK FullTextScore(c.text, 'mountain')"; + + QueryResult directResult = drainQuery(directFtsContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(tcFtsContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + assertThat(directResult.results.size()).as("FullTextScore ranking query should return results").isPositive(); + assertThat(tcResult.results.size()).isEqualTo(directResult.results.size()); + + for (int i = 0; i < directResult.results.size(); i++) { + assertThat(tcResult.results.get(i).get("id").asText()) + .as("FullTextScore ranking result mismatch at position " + i) + .isEqualTo(directResult.results.get(i).get("id").asText()); + } + + } finally { + safeDeleteContainer(directFtsContainer); + } + } + + // ==================== Hybrid Search ==================== + + /** + * Creates a container with vector + full-text policies, runs hybrid RRF query. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT * 2) + public void testHybridSearchGatewayVsThinClient() { + String containerId = "hybridCompare_" + UUID.randomUUID().toString().substring(0, 8); + CosmosAsyncDatabase db = directClient.getDatabase(directContainer.getDatabase().getId()); + CosmosAsyncContainer directHybridContainer = db.getContainer(containerId); + + try { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(Collections.singletonList("/" + PK_FIELD)); + + CosmosContainerProperties props = new CosmosContainerProperties(containerId, pkDef); + + CosmosVectorEmbeddingPolicy vecPolicy = new CosmosVectorEmbeddingPolicy(); + CosmosVectorEmbedding emb = new CosmosVectorEmbedding(); + emb.setPath("/vector"); + emb.setDataType(CosmosVectorDataType.FLOAT32); + emb.setEmbeddingDimensions(3); + emb.setDistanceFunction(CosmosVectorDistanceFunction.COSINE); + vecPolicy.setCosmosVectorEmbeddings(Collections.singletonList(emb)); + props.setVectorEmbeddingPolicy(vecPolicy); + + CosmosFullTextPath ftPath = new CosmosFullTextPath(); + ftPath.setPath("/text"); + ftPath.setLanguage("en-US"); + CosmosFullTextPolicy ftPolicy = new CosmosFullTextPolicy(); + ftPolicy.setDefaultLanguage("en-US"); + ftPolicy.setPaths(Collections.singletonList(ftPath)); + props.setFullTextPolicy(ftPolicy); + + IndexingPolicy idxPolicy = new IndexingPolicy(); + idxPolicy.setIndexingMode(IndexingMode.CONSISTENT); + idxPolicy.setIncludedPaths(Collections.singletonList(new IncludedPath("/*"))); + idxPolicy.setExcludedPaths(Arrays.asList(new ExcludedPath("/vector/*"), new ExcludedPath("/\"_etag\"/?"))); + CosmosVectorIndexSpec vecIdx = new CosmosVectorIndexSpec(); + vecIdx.setPath("/vector"); + vecIdx.setType(CosmosVectorIndexType.FLAT.toString()); + idxPolicy.setVectorIndexes(Collections.singletonList(vecIdx)); + CosmosFullTextIndex ftIndex = new CosmosFullTextIndex(); + ftIndex.setPath("/text"); + idxPolicy.setCosmosFullTextIndexes(Collections.singletonList(ftIndex)); + props.setIndexingPolicy(idxPolicy); + + db.createContainer(props).block(); + CosmosAsyncContainer tcHybridContainer = thinClient.getDatabase(db.getId()).getContainer(containerId); + + String hybridPk = UUID.randomUUID().toString(); + String[] texts = { + "Red bicycle on the mountain trail", + "Blue car parked in the city", + "Green bicycle near the lake" + }; + double[][] vectors = { + {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0} + }; + for (int i = 0; i < texts.length; i++) { + ObjectNode doc = OBJECT_MAPPER.createObjectNode(); + doc.put(ID_FIELD, "hybrid_" + i + "_" + UUID.randomUUID().toString().substring(0, 8)); + doc.put(PK_FIELD, hybridPk); + doc.put("text", texts[i]); + ArrayNode arr = doc.putArray("vector"); + for (double v : vectors[i]) { arr.add(v); } + directHybridContainer.createItem(doc, new PartitionKey(hybridPk), null).block(); + } + + String query = "SELECT TOP 3 * FROM c " + + "ORDER BY RANK RRF(VectorDistance(c.vector, [1.0, 0.0, 0.0]), FullTextScore(c.text, 'bicycle'))"; + + QueryResult directResult = drainQuery(directHybridContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(tcHybridContainer, query, new CosmosQueryRequestOptions(), ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + assertThat(tcResult.results.size()).isEqualTo(directResult.results.size()); + + for (int i = 0; i < directResult.results.size(); i++) { + assertThat(tcResult.results.get(i).get("id").asText()) + .as("Hybrid search result mismatch at position " + i) + .isEqualTo(directResult.results.get(i).get("id").asText()); + } + + } finally { + safeDeleteContainer(directHybridContainer); + } + } + + // ==================== readManyByPartitionKeys Tests ==================== + // These exercise the validation QueryPlan path that bifurcates between Compute Gateway + // and the thin client (Gateway V2) based on the DocumentCollection being passed through. + // For the thin client container we expect the validation QueryPlan call itself to land on + // the :10250 endpoint, just like the per-batch query requests it precedes. + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testReadManyByPartitionKeysNoCustomQuery() { + // No custom query: no QueryPlan is fetched, but the per-batch reads must still + // route through the thin client endpoint. + ReadManyResult gwResult = drainReadMany( + directContainer, Collections.singletonList(new PartitionKey(commonPk)), null, ObjectNode.class); + ReadManyResult tcResult = drainReadMany( + thinClientContainer, Collections.singletonList(new PartitionKey(commonPk)), null, ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()) + .as("readManyByPartitionKeys (no custom query) count mismatch") + .isEqualTo(gwResult.results.size()); + assertThat(idsSorted(tcResult.results)) + .as("readManyByPartitionKeys (no custom query) IDs mismatch") + .isEqualTo(idsSorted(gwResult.results)); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testReadManyByPartitionKeysWithCustomQuery() { + // Custom projection + filter triggers QueryPlan validation. With the bifurcation + // wiring in place the QueryPlan request must travel through Gateway V2 (:10250). + SqlQuerySpec customQuery = new SqlQuerySpec( + "SELECT c.id, c.category, c.status FROM c WHERE c.status = 'active'"); + + ReadManyResult gwResult = drainReadMany( + directContainer, Collections.singletonList(new PartitionKey(commonPk)), customQuery, ObjectNode.class); + ReadManyResult tcResult = drainReadMany( + thinClientContainer, Collections.singletonList(new PartitionKey(commonPk)), customQuery, ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()) + .as("readManyByPartitionKeys (custom query) count mismatch") + .isEqualTo(gwResult.results.size()); + assertThat(tcResult.results.size()) + .as("readManyByPartitionKeys (custom query) expected at least one active doc") + .isGreaterThan(0); + assertThat(idsSorted(tcResult.results)) + .as("readManyByPartitionKeys (custom query) IDs mismatch") + .isEqualTo(idsSorted(gwResult.results)); + + // Projection assertion: every returned doc has the requested fields only. + for (ObjectNode doc : tcResult.results) { + assertThat(doc.has("id")).isTrue(); + assertThat(doc.has("category")).isTrue(); + assertThat(doc.has("status")).isTrue(); + assertThat(doc.get("status").asText()).isEqualTo("active"); + } + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testReadManyByPartitionKeysWithParameterizedCustomQuery() { + // Parameterized custom query — same validation path, exercises SqlParameter binding. + SqlQuerySpec customQuery = new SqlQuerySpec( + "SELECT * FROM c WHERE c.category = @cat", + Collections.singletonList(new SqlParameter("@cat", "electronics"))); + + ReadManyResult gwResult = drainReadMany( + directContainer, Collections.singletonList(new PartitionKey(commonPk)), customQuery, ObjectNode.class); + ReadManyResult tcResult = drainReadMany( + thinClientContainer, Collections.singletonList(new PartitionKey(commonPk)), customQuery, ObjectNode.class); + + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()) + .as("readManyByPartitionKeys (parameterized) count mismatch") + .isEqualTo(gwResult.results.size()); + assertThat(tcResult.results.size()) + .as("readManyByPartitionKeys (parameterized) expected at least one electronics doc") + .isGreaterThan(0); + assertThat(idsSorted(tcResult.results)) + .as("readManyByPartitionKeys (parameterized) IDs mismatch") + .isEqualTo(idsSorted(gwResult.results)); + + for (ObjectNode doc : tcResult.results) { + assertThat(doc.get("category").asText()).isEqualTo("electronics"); + } + } + + // ==================== Assertion & Drain Helpers ==================== + + private static void safeDeleteContainer(CosmosAsyncContainer container) { + if (container != null) { + try { container.delete().block(); } catch (Exception e) { logger.warn("Container cleanup failed: {}", e.getMessage()); } + } + } + + /** Holds query results and per-page diagnostics from a fully drained query. */ + private static class QueryResult { + final List results = new ArrayList<>(); + final List diagnostics = new ArrayList<>(); + } + + /** Holds readManyByPartitionKeys results and per-page diagnostics. */ + private static class ReadManyResult { + final List results = new ArrayList<>(); + final List diagnostics = new ArrayList<>(); + } + + private ReadManyResult drainReadMany( + CosmosAsyncContainer c, + List partitionKeys, + SqlQuerySpec customQuery, + Class type) { + ReadManyResult result = new ReadManyResult<>(); + for (FeedResponse page : c.readManyByPartitionKeys(partitionKeys, customQuery, type) + .byPage() + .toIterable()) { + result.results.addAll(page.getResults()); + result.diagnostics.add(page.getCosmosDiagnostics()); + } + return result; + } + + private static List idsSorted(List docs) { + return docs.stream() + .filter(d -> d.has(ID_FIELD)) + .map(d -> d.get(ID_FIELD).asText()) + .sorted() + .collect(Collectors.toList()); + } + + private CosmosQueryRequestOptions partitionedOptions() { + CosmosQueryRequestOptions opts = new CosmosQueryRequestOptions(); + opts.setPartitionKey(new PartitionKey(commonPk)); + return opts; + } + + private QueryResult drainQuery(CosmosAsyncContainer c, String query, CosmosQueryRequestOptions opts, Class type) { + QueryResult result = new QueryResult<>(); + for (FeedResponse page : c.queryItems(query, opts, type).byPage().toIterable()) { + result.results.addAll(page.getResults()); + result.diagnostics.add(page.getCosmosDiagnostics()); + } + return result; + } + + private QueryResult drainQuery(CosmosAsyncContainer c, SqlQuerySpec qs, CosmosQueryRequestOptions opts, Class type) { + QueryResult result = new QueryResult<>(); + for (FeedResponse page : c.queryItems(qs, opts, type).byPage().toIterable()) { + result.results.addAll(page.getResults()); + result.diagnostics.add(page.getCosmosDiagnostics()); + } + return result; + } + + /** + * Direct vs thin client comparison: run query via both Direct TCP and thin client. + * Assert: (1) thin client used :10250, (2) same count, (3) same document IDs — compared in + * strict sequence for ORDER BY queries and for single-partition queries (partition key set on + * the options), and as sorted sets only for cross-partition non-ORDER-BY queries where ordering + * is undefined. See {@link #isStrictlyOrdered}. + */ + private void assertDirectAndThinClientMatch(String query) { + assertDirectAndThinClientMatch(query, partitionedOptions()); + } + + private void assertDirectAndThinClientMatch(String query, CosmosQueryRequestOptions options) { + QueryResult gwResult = drainQuery(directContainer, query, options, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientContainer, query, options, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Count mismatch: " + query).isEqualTo(gwResult.results.size()); + assertSameDocumentIds(gwResult.results, tcResult.results, isStrictlyOrdered(query, options), query); + } + + private void assertDirectAndThinClientMatch(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { + QueryResult gwResult = drainQuery(directContainer, querySpec, options, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientContainer, querySpec, options, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Count mismatch: " + querySpec.getQueryText()).isEqualTo(gwResult.results.size()); + assertSameDocumentIds(gwResult.results, tcResult.results, isStrictlyOrdered(querySpec.getQueryText(), options), querySpec.getQueryText()); + } + + private void assertScalarDirectAndThinClientMatch(String query, Class resultType) { + assertScalarDirectAndThinClientMatch(query, partitionedOptions(), resultType); + } + + private void assertScalarDirectAndThinClientMatch(String query, CosmosQueryRequestOptions options, Class resultType) { + QueryResult gwResult = drainQuery(directContainer, query, options, resultType); + QueryResult tcResult = drainQuery(thinClientContainer, query, options, resultType); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Scalar count mismatch: " + query).isEqualTo(gwResult.results.size()); + for (int i = 0; i < gwResult.results.size(); i++) { + assertScalarValueEquals(tcResult.results.get(i), gwResult.results.get(i), i + ": " + query); + } + } + + /** Direct vs thin client comparison for GROUP BY where result order may vary — compare as sets. */ + private void assertGroupByDirectAndThinClientMatch(String query, String groupField) { + QueryResult gwResult = drainQuery(directContainer, query, partitionedOptions(), ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientContainer, query, partitionedOptions(), ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("GROUP BY count mismatch: " + query).isEqualTo(gwResult.results.size()); + for (ObjectNode gwRow : gwResult.results) { + String key = gwRow.get(groupField).asText(); + boolean found = tcResult.results.stream().anyMatch(tc -> tc.get(groupField).asText().equals(key) + && jsonEqualsWithTolerance(tc, gwRow)); + assertThat(found).as("GROUP BY row not found in thin client results: " + key).isTrue(); + } + } + + private static final double NUMERIC_TOLERANCE = 1e-6; + + private static boolean isOrderedQuery(String queryText) { + return queryText != null && queryText.toUpperCase(Locale.ROOT).contains("ORDER BY"); + } + + /** + * Decides whether Direct and thin client document IDs must match in strict sequence. Per query + * test best practice we validate order for as many queries as possible to flush out hidden + * ordering/paging bugs, and relax to a sorted-set comparison only when order is genuinely + * undefined: a cross-partition query (no partition key on the options) that lacks an ORDER BY. + * Single-partition queries (partitionedOptions sets a partition key) and any ORDER BY query are + * therefore always compared strictly. + */ + private static boolean isStrictlyOrdered(String queryText, CosmosQueryRequestOptions options) { + boolean crossPartition = options == null || options.getPartitionKey() == null; + return isOrderedQuery(queryText) || !crossPartition; + } + + /** + * Compares the Direct (baseline) and thin client result rows. When {@code ordered} is true the + * sequence is significant (compared position-by-position); otherwise the rows are compared + * order-insensitively (a cross-partition non-ORDER-BY query, where cross-partition / cross-page + * ordering is undefined). See {@link #isStrictlyOrdered}. + *

+ * If every row projects an {@code id} it is used as the comparison key (stable and cheap). + * Otherwise (projection queries that do not select {@code id}, e.g. {@code SELECT AS x}) + * we do NOT drop the check and lose coverage - the full projected rows are compared instead, + * with numeric fields matched within {@link #NUMERIC_TOLERANCE} to avoid float-formatting false + * mismatches between the Direct and thin-client serialization paths. + */ + private static void assertSameDocumentIds(List gwDocs, List tcDocs, boolean ordered, String desc) { + assertThat(tcDocs.size()).as("Row count mismatch: " + desc).isEqualTo(gwDocs.size()); + + boolean allHaveId = !gwDocs.isEmpty() + && gwDocs.stream().allMatch(d -> d.has(ID_FIELD)) + && tcDocs.stream().allMatch(d -> d.has(ID_FIELD)); + + if (allHaveId) { + List gwIds = gwDocs.stream().map(d -> d.get(ID_FIELD).asText()).collect(Collectors.toList()); + List tcIds = tcDocs.stream().map(d -> d.get(ID_FIELD).asText()).collect(Collectors.toList()); + if (ordered) { + assertThat(tcIds).as("IDs mismatch (ordered): " + desc).isEqualTo(gwIds); + } else { + assertThat(tcIds.stream().sorted().collect(Collectors.toList())) + .as("IDs mismatch (unordered): " + desc) + .isEqualTo(gwIds.stream().sorted().collect(Collectors.toList())); + } + return; + } + + // No id projected - compare the full rows so ordering/content coverage is not lost. + if (ordered) { + for (int i = 0; i < gwDocs.size(); i++) { + assertThat(jsonEqualsWithTolerance(tcDocs.get(i), gwDocs.get(i))) + .as("Row mismatch (ordered) at index " + i + ": " + desc + + " | expected=" + gwDocs.get(i) + " actual=" + tcDocs.get(i)) + .isTrue(); + } + } else { + // Multiset comparison: each Direct row must have a distinct matching thin-client row. + List remaining = new ArrayList<>(tcDocs); + for (ObjectNode gwRow : gwDocs) { + int matchIdx = -1; + for (int j = 0; j < remaining.size(); j++) { + if (jsonEqualsWithTolerance(remaining.get(j), gwRow)) { + matchIdx = j; + break; + } + } + assertThat(matchIdx) + .as("Row not found in thin-client results (unordered): " + desc + " | row=" + gwRow) + .isGreaterThanOrEqualTo(0); + remaining.remove(matchIdx); + } + } + } + + /** + * Cross-partition parity check against the dedicated multi-physical-partition container. Asserts + * (1) thin client routing, (2) result parity with Direct (strict order for ORDER BY queries), + * and (3) that the query actually fanned out across more than one partition key range - i.e. the + * seeded data really does span multiple partitions. + */ + private void assertCrossPartitionDirectAndThinClientMatch(String query) { + CosmosQueryRequestOptions crossPartitionOptions = new CosmosQueryRequestOptions(); + QueryResult directResult = drainQuery(directCrossPartitionContainer, query, crossPartitionOptions, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientCrossPartitionContainer, query, crossPartitionOptions, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Count mismatch: " + query).isEqualTo(directResult.results.size()); + assertSameDocumentIds(directResult.results, tcResult.results, isStrictlyOrdered(query, crossPartitionOptions), query); + + // Prove the fixture data genuinely spans more than one physical partition (so this really is + // a cross-partition query). Counted from the Direct diagnostics, which reliably report the + // partition key range per request; the parity check above already proves the thin client + // fanned out and merged all rows from those partitions correctly. + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("cross-partition query must contact more than one partition key range: " + query) + .isGreaterThan(1); + } + + /** + * Cross-partition scalar-aggregate parity check. Runs a {@code SELECT VALUE } query with no + * partition key on the options, so the thin client must fan out to every physical partition, + * collect each partition's partial aggregate, and merge them into the single returned value - + * exercising the cross-partition aggregate MERGE path that the single-partition aggregate tests + * never reach. Asserts (1) thin client routing, (2) value parity with the Direct baseline within + * numeric tolerance, and (3) that the Direct baseline genuinely fanned out across more than one + * partition key range, proving the merge actually happened. + */ + private void assertCrossPartitionScalarMatch(String query, Class resultType) { + CosmosQueryRequestOptions crossPartitionOptions = new CosmosQueryRequestOptions(); + QueryResult directResult = drainQuery(directCrossPartitionContainer, query, crossPartitionOptions, resultType); + QueryResult tcResult = drainQuery(thinClientCrossPartitionContainer, query, crossPartitionOptions, resultType); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Scalar count mismatch: " + query).isEqualTo(directResult.results.size()); + for (int i = 0; i < directResult.results.size(); i++) { + assertScalarValueEquals(tcResult.results.get(i), directResult.results.get(i), i + ": " + query); + } + + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("cross-partition aggregate must contact more than one partition key range: " + query) + .isGreaterThan(1); + } + + /** + * Cross-partition GROUP BY parity check. With no partition key on the options the thin client must + * fan out to every physical partition and merge each partition's partial groups into the final + * grouped result - exercising the cross-partition GROUP BY MERGE path. Because group/page ordering + * across partitions is undefined, rows are compared as a set keyed on {@code groupField} (numeric + * fields matched within tolerance). Asserts thin client routing, set parity with Direct, and + * genuine multi-range fan-out on the Direct baseline. + */ + private void assertCrossPartitionGroupByMatch(String query, String groupField) { + CosmosQueryRequestOptions crossPartitionOptions = new CosmosQueryRequestOptions(); + QueryResult directResult = drainQuery(directCrossPartitionContainer, query, crossPartitionOptions, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientCrossPartitionContainer, query, crossPartitionOptions, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("GROUP BY count mismatch: " + query).isEqualTo(directResult.results.size()); + for (ObjectNode directRow : directResult.results) { + String key = directRow.get(groupField).asText(); + boolean found = tcResult.results.stream().anyMatch(tc -> tc.get(groupField).asText().equals(key) + && jsonEqualsWithTolerance(tc, directRow)); + assertThat(found).as("GROUP BY row not found in thin client results: " + key).isTrue(); + } + + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("cross-partition GROUP BY must contact more than one partition key range: " + query) + .isGreaterThan(1); + } + + /** + * Hierarchical (MULTI_HASH) cross-partition parity check. With no partition key on the options the + * thin client fans out across the /tenantId,/userId container's physical partitions, forcing the + * proxy to emit MULTI-COMPONENT PartitionKeyInternal arrays that drive the MULTI_HASH branch of the + * range conversion. Asserts (1) thin client routing, (2) count parity, (3) row parity (ordered when + * the query is strictly ordered), and (4) genuine multi-range fan-out on the Direct baseline. + */ + private void assertHierarchicalCrossPartitionMatch(String query) { + CosmosQueryRequestOptions crossPartitionOptions = new CosmosQueryRequestOptions(); + QueryResult directResult = drainQuery(directHierarchicalContainer, query, crossPartitionOptions, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientHierarchicalContainer, query, crossPartitionOptions, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(tcResult.results.size()).as("Hierarchical count mismatch: " + query).isEqualTo(directResult.results.size()); + assertSameDocumentIds(directResult.results, tcResult.results, isStrictlyOrdered(query, crossPartitionOptions), query); + + assertThat(distinctPartitionKeyRangesContacted(directResult.diagnostics)) + .as("hierarchical cross-partition query must contact more than one partition key range: " + query) + .isGreaterThan(1); + } + + /** + * Hierarchical SCOPED parity check. Runs the query with a caller-supplied partition key on the + * options - either the full two-component key (a single point range) or just the first component + * (a half-open prefix range) - so the proxy emits the corresponding scoped PartitionKeyInternal + * range. Asserts thin client routing, that the Direct baseline returns {@code expectedCount} rows, + * and full row parity. No multi-range assertion: a scoped query may legitimately hit a single range. + */ + private void assertHierarchicalScopedMatch(String query, CosmosQueryRequestOptions options, int expectedCount) { + QueryResult directResult = drainQuery(directHierarchicalContainer, query, options, ObjectNode.class); + QueryResult tcResult = drainQuery(thinClientHierarchicalContainer, query, options, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(directResult.results.size()).as("Direct scoped count: " + query).isEqualTo(expectedCount); + assertThat(tcResult.results.size()).as("Thin-client scoped count: " + query).isEqualTo(directResult.results.size()); + assertSameDocumentIds(directResult.results, tcResult.results, isStrictlyOrdered(query, options), query); + } + + /** Counts the distinct partition key ranges contacted across all pages' diagnostics. */ + private static int distinctPartitionKeyRangesContacted(List diagnosticsList) { + Set ranges = new HashSet<>(); + for (CosmosDiagnostics diagnostics : diagnosticsList) { + CosmosDiagnosticsContext ctx = diagnostics.getDiagnosticsContext(); + if (ctx == null) { + continue; + } + for (CosmosDiagnosticsRequestInfo requestInfo : ctx.getRequestInfo()) { + String pkRangeId = requestInfo.getPartitionKeyRangeId(); + if (pkRangeId != null && !pkRangeId.isEmpty()) { + ranges.add(pkRangeId); + } + } + } + return ranges.size(); + } + + /** + * Compares scalar aggregate values, treating two numeric values as equal when they fall within + * NUMERIC_TOLERANCE. Avoids false mismatches from floating-point formatting differences (e.g. + * SUM/AVG) between the Direct and thin client paths; falls back to string equality otherwise. + */ + private static void assertScalarValueEquals(T tcValue, T gwValue, String desc) { + String tcStr = String.valueOf(tcValue); + String gwStr = String.valueOf(gwValue); + Double tcNum = tryParseDouble(tcStr); + Double gwNum = tryParseDouble(gwStr); + if (tcNum != null && gwNum != null) { + assertThat(Math.abs(tcNum - gwNum)).as("Scalar numeric mismatch at " + desc).isLessThan(NUMERIC_TOLERANCE); + } else { + assertThat(tcStr).as("Scalar value mismatch at " + desc).isEqualTo(gwStr); + } + } + + private static Double tryParseDouble(String value) { + if (value == null) { + return null; + } + try { + return Double.parseDouble(value.trim()); + } catch (NumberFormatException e) { + return null; + } + } + + /** + * Recursively compares two JSON nodes, treating numeric leaves as equal when they fall within + * NUMERIC_TOLERANCE. Prevents false GROUP BY mismatches caused by floating-point formatting + * differences in aggregate values (e.g. SUM/AVG) between the Direct and thin client paths. + */ + private static boolean jsonEqualsWithTolerance(JsonNode a, JsonNode b) { + if (a == null || b == null) { + return a == b; + } + if (a.isNumber() && b.isNumber()) { + return Math.abs(a.asDouble() - b.asDouble()) < NUMERIC_TOLERANCE; + } + if (a.isObject() && b.isObject()) { + if (a.size() != b.size()) { + return false; + } + Iterator fieldNames = a.fieldNames(); + while (fieldNames.hasNext()) { + String field = fieldNames.next(); + if (!b.has(field) || !jsonEqualsWithTolerance(a.get(field), b.get(field))) { + return false; + } + } + return true; + } + if (a.isArray() && b.isArray()) { + if (a.size() != b.size()) { + return false; + } + for (int i = 0; i < a.size(); i++) { + if (!jsonEqualsWithTolerance(a.get(i), b.get(i))) { + return false; + } + } + return true; + } + return a.equals(b); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java new file mode 100644 index 000000000000..379d12bd3720 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java @@ -0,0 +1,111 @@ +// 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) { + assertThat(diagnostics).isNotNull(); + CosmosDiagnosticsContext ctx = diagnostics.getDiagnosticsContext(); + assertThat(ctx).isNotNull(); + Collection requests = ctx.getRequestInfo(); + assertThat(requests).isNotNull(); + assertThat(requests.size()).isPositive(); + int requestCountAgainstThinClientEndpoint = 0; + for (CosmosDiagnosticsRequestInfo requestInfo : requests) { + if (requestInfo.getEndpoint().contains(THIN_CLIENT_ENDPOINT_INDICATOR)) { + requestCountAgainstThinClientEndpoint++; + } + } + assertThat(requestCountAgainstThinClientEndpoint).isEqualTo(requests.size()); + } + + /** + * 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 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); + } + } +} From 5eb41ee85b5975aebb2c94a933c8a002d330ac54 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 14:34:27 -0400 Subject: [PATCH 03/15] Reuse prefix EPK range on thin-client parallel prefix-query path The earlier fix (499f928) only scoped the prefix MULTI_HASH read when the request carried a non-null partition key. The parallel prefix-query path intentionally nulls the partial hierarchical partition key and instead carries the narrow prefix effective-partition-key range on the request's feedRange, so that guard never fired there and the thin-client read still over-spanned to the whole physical partition. Mark such requests (prefixPartitionKeyQuery) in ParallelDocumentQueryExecutionContextBase and, in ThinClientStoreModel, reuse the already-computed FeedRangeEpkImpl range as the doc-level EPK filter (START_EPK/END_EPK + StartEpkHash/EndEpkHash) instead of recomputing getEPKRangeForPrefixPartitionKey from a partition key we no longer have. This mirrors the Direct/GatewayV1 decision points and honors DRY/SRP. Also relax assertThinClientEndpointUsed so an all-QueryPlan diagnostics context passes: in thin-client mode QueryPlan calls resolve via the classic gateway and are the only requests permitted on a non-thin-client endpoint; any non-QueryPlan (data) request on a non-thin-client endpoint still fails the assertion. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/rx/TestSuiteBase.java | 20 ++++++-- .../RxDocumentServiceRequest.java | 14 ++++++ .../implementation/ThinClientStoreModel.java | 49 ++++++++++++------- ...llelDocumentQueryExecutionContextBase.java | 14 +++++- 4 files changed, 74 insertions(+), 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index 2f4a06dd6046..bda1674d6312 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -2176,16 +2176,30 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnosticsContext ctx) assertThat(requests).isNotNull(); assertThat(requests.size()).isPositive(); + boolean hasNonQueryPlanRequest = false; for (CosmosDiagnosticsRequestInfo requestInfo : requests) { if (requestInfo.getEndpoint() != null && requestInfo.getEndpoint().contains(THIN_CLIENT_ENDPOINT_INDICATOR)) { + // A thin-client request exists -> invariant satisfied. return; } + + // requestType has the form ":" (OperationType.QueryPlan + // stringifies to "QueryPlan"). In thin-client mode QueryPlan calls are resolved via the + // classic gateway, so they are the only requests allowed to target a non-thin-client endpoint. + String requestType = requestInfo.getRequestType(); + if (requestType == null || !requestType.endsWith(":QueryPlan")) { + hasNonQueryPlanRequest = true; + } } - assertThat(false) - .as("No request targeting thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR + ")") - .isTrue(); + // Reaching here means no request targeted the thin-client proxy endpoint. That is acceptable + // only when every request was a QueryPlan call; a non-QueryPlan (data) request on a non-thin-client + // endpoint violates the thin-client routing invariant. + assertThat(hasNonQueryPlanRequest) + .as("No request targeting thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR + + ") and at least one non-QueryPlan request was present") + .isFalse(); } protected static void safeClose(AsyncDocumentClient client) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index 03f27156f62f..e4245b2788cb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -78,6 +78,12 @@ private static CosmosItemSerializer internalDefaultSerializer() { private Range 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; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index 6fe5267ab28a..819b0ead62c8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -14,6 +14,7 @@ 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; @@ -257,12 +258,25 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque // hex-string-as-bytes encoding). Without this filter the proxy resolves the request to the // owning physical partition and returns every co-located document (an over-span). Range prefixEpkRange = null; - if (request.getOperationType() != OperationType.QueryPlan && partitionKey != null) { - PartitionKeyDefinition pkDefinition = request.getPartitionKeyDefinition(); - if (pkDefinition != null - && pkDefinition.getKind() == PartitionKind.MULTI_HASH - && partitionKey.getComponents().size() < pkDefinition.getPaths().size()) { - prefixEpkRange = partitionKey.getEPKRangeForPrefixPartitionKey(pkDefinition); + if (request.getOperationType() != OperationType.QueryPlan) { + if (partitionKey != null) { + PartitionKeyDefinition pkDefinition = request.getPartitionKeyDefinition(); + if (pkDefinition != null + && pkDefinition.getKind() == PartitionKind.MULTI_HASH + && partitionKey.getComponents().size() < pkDefinition.getPaths().size()) { + prefixEpkRange = partitionKey.getEPKRangeForPrefixPartitionKey(pkDefinition); + } + } else if (request.isPrefixPartitionKeyQuery() && request.getFeedRange() instanceof FeedRangeEpkImpl) { + // Parallel prefix-query path: 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 it directly instead of recomputing getEPKRangeForPrefixPartitionKey from a + // partition key we no longer have. This is what scopes the read on this path; without it + // the request falls through to the physical-partition steering below and over-spans. + prefixEpkRange = ((FeedRangeEpkImpl) request.getFeedRange()).getRange(); + } + if (prefixEpkRange != null) { request.getHeaders().put( HttpConstants.HttpHeaders.READ_FEED_KEY_TYPE, ReadFeedKeyType.EffectivePartitionKeyRange.name()); @@ -278,18 +292,17 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque RntbdRequest rntbdRequest = RntbdRequest.from(rntbdRequestArgs); - if (partitionKey != null) { - if (prefixEpkRange != null) { - // Partial (prefix) hierarchical partition key. StartEpk/EndEpk + ReadFeedKeyType were - // already set on the request headers above and serialized into the RNTBD request as the - // backend's doc-level EPK filter. StartEpkHash/EndEpkHash additionally steer the proxy to - // the owning physical partition(s), mirroring the Direct/RNTBD path. - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMin())); - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMax())); - } else { - byte[] epk = partitionKey.getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition()); - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EffectivePartitionKey, epk); - } + if (prefixEpkRange != null) { + // Partial (prefix) hierarchical partition key - either the PK-bearing path (point/scalar + // prefix) or the parallel path that reused request.getFeedRange(). StartEpk/EndEpk + + // ReadFeedKeyType were already set on the request headers above and serialized into the RNTBD + // request as the backend's doc-level EPK filter. StartEpkHash/EndEpkHash additionally steer the + // proxy to the owning physical partition(s), mirroring the Direct/RNTBD path. + 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) { throw new IllegalStateException( "Resolved partition key range should not be null at this point. ResourceType: " diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java index 5774a9dde4fe..e34b46c00f4f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java @@ -81,11 +81,19 @@ 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())) { + if (PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey())) { + // Partial (prefix) hierarchical partition key: partitionKeyInternal stays null (no + // PARTITION_KEY header). The narrow prefix effective-partition-key range is already + // carried on the feedRange passed below; mark the request so downstream store models + // (e.g. ThinClientStoreModel) can reuse it as a doc-level EPK filter instead of + // over-spanning to the whole physical partition. + isPrefixPartitionKeyQuery = true; + } else { partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey()); headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson()); } @@ -93,8 +101,10 @@ protected void initialize( 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>> executeFunc = (request) -> this.executeRequestAsync( From 24921af8f42696a62f9482aee681ce82edf2dae5 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 17:22:19 -0400 Subject: [PATCH 04/15] Consolidate thin-client prefix EPK RNTBD headers, propagate prefix flag on clone, add readAllItems/readMany prefix HPK tests ThinClientStoreModel.wrapInHttpRequest now sets all five EPK-specific RNTBD headers (ReadFeedKeyType, StartEpk/EndEpk, StartEpkHash/EndEpkHash) in a single block directly on the RNTBD request, replacing the split request-header-map + post-construction approach. Wire encoding is unchanged (hex-string UTF-8 bytes for StartEpk/EndEpk, decoded hash bytes for the hash headers). RxDocumentServiceRequest.clone() now copies prefixPartitionKeyQuery so hedged/availability-strategy/retry clones keep the prefix signal and do not over-span. ThinClientQueryE2ETest adds prefix-HPK regression coverage for readAllItems (ReadFeed) and readManyByPartitionKeys, asserting Direct-vs-thin-client parity and per-tenant scoping (no over-span). --- .../cosmos/rx/ThinClientQueryE2ETest.java | 100 ++++++++++++++++++ .../RxDocumentServiceRequest.java | 1 + .../implementation/ThinClientStoreModel.java | 51 +++++---- 3 files changed, 126 insertions(+), 26 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java index 4fcdcde14da7..ff312ed85081 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -1626,6 +1626,94 @@ public void testReadManyByPartitionKeysWithParameterizedCustomQuery() { } } + // ==================== Partial (Prefix) HPK read Tests ==================== + // Regression coverage for the thin-client MULTI_HASH prefix EPK over-span fix. A partial + // hierarchical key (only /tenantId, omitting /userId) must scope the read to that tenant's + // effective-partition-key sub-range [hash(prefix), hash(prefix)+"FF") as a doc-level filter, + // NOT resolve to the whole owning physical partition (which would return co-located documents + // from other tenants). Direct TCP is the baseline; the thin client must match it exactly, so a + // count/ID mismatch is exactly the over-span regression these tests guard against. + + /** + * readAllItems (ReadFeed) with a partial (prefix) hierarchical partition key - only the first + * component (/tenantId). Must return exactly that tenant's {@code HIER_USERS_PER_TENANT} + * documents, matching the Direct baseline. If the thin client over-spanned to the physical + * partition it would return co-located docs from other tenants, which the count/ID parity and the + * per-doc tenant assertion catch. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testHierarchicalReadAllItemsPrefixPartitionKey() { + String tenant0 = hierarchicalKeys.get(0)[0]; + PartitionKey prefixKey = new PartitionKeyBuilder().add(tenant0).build(); + + QueryResult directResult = drainReadAllItems(directHierarchicalContainer, prefixKey, ObjectNode.class); + QueryResult tcResult = drainReadAllItems(thinClientHierarchicalContainer, prefixKey, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(directResult.results.size()) + .as("Direct readAllItems prefix count").isEqualTo(HIER_USERS_PER_TENANT); + assertThat(tcResult.results.size()) + .as("Thin-client readAllItems prefix over-span vs Direct").isEqualTo(directResult.results.size()); + assertThat(idsSorted(tcResult.results)) + .as("readAllItems prefix ID mismatch").isEqualTo(idsSorted(directResult.results)); + // Every returned doc must belong to the prefixed tenant (no over-span into co-located tenants). + for (ObjectNode doc : tcResult.results) { + assertThat(doc.get(TENANT_FIELD).asText()) + .as("readAllItems prefix returned a foreign tenant's document").isEqualTo(tenant0); + } + } + + /** + * readAllItems (ReadFeed) with the FULL two-component hierarchical key resolves to a single + * logical partition - exactly one document - matching the Direct baseline. Guards that the prefix + * handling does not regress the full-key point path. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testHierarchicalReadAllItemsFullPartitionKey() { + String[] firstKey = hierarchicalKeys.get(0); + PartitionKey fullKey = new PartitionKeyBuilder().add(firstKey[0]).add(firstKey[1]).build(); + + QueryResult directResult = drainReadAllItems(directHierarchicalContainer, fullKey, ObjectNode.class); + QueryResult tcResult = drainReadAllItems(thinClientHierarchicalContainer, fullKey, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(directResult.results.size()).as("Direct readAllItems full-key count").isEqualTo(1); + assertThat(tcResult.results.size()) + .as("Thin-client readAllItems full-key vs Direct").isEqualTo(directResult.results.size()); + assertThat(idsSorted(tcResult.results)) + .as("readAllItems full-key ID mismatch").isEqualTo(idsSorted(directResult.results)); + } + + /** + * readManyByPartitionKeys with a partial (prefix) hierarchical key. groupPartitionKeysByPhysicalPartition + * expands the prefix to its effective-partition-key range; the thin client must apply that as a + * doc-level filter and return exactly the tenant's {@code HIER_USERS_PER_TENANT} documents, + * matching Direct rather than over-spanning the physical partition. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testHierarchicalReadManyByPartitionKeysPrefixPartitionKey() { + String tenant0 = hierarchicalKeys.get(0)[0]; + PartitionKey prefixKey = new PartitionKeyBuilder().add(tenant0).build(); + SqlQuerySpec selectAll = new SqlQuerySpec("SELECT * FROM c"); + + ReadManyResult directResult = drainReadMany( + directHierarchicalContainer, Collections.singletonList(prefixKey), selectAll, ObjectNode.class); + ReadManyResult tcResult = drainReadMany( + thinClientHierarchicalContainer, Collections.singletonList(prefixKey), selectAll, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + + assertThat(directResult.results.size()) + .as("Direct readMany prefix count").isEqualTo(HIER_USERS_PER_TENANT); + assertThat(tcResult.results.size()) + .as("Thin-client readMany prefix over-span vs Direct").isEqualTo(directResult.results.size()); + assertThat(idsSorted(tcResult.results)) + .as("readMany prefix ID mismatch").isEqualTo(idsSorted(directResult.results)); + for (ObjectNode doc : tcResult.results) { + assertThat(doc.get(TENANT_FIELD).asText()) + .as("readMany prefix returned a foreign tenant's document").isEqualTo(tenant0); + } + } + // ==================== Assertion & Drain Helpers ==================== private static void safeDeleteContainer(CosmosAsyncContainer container) { @@ -1661,6 +1749,18 @@ private ReadManyResult drainReadMany( return result; } + private QueryResult drainReadAllItems( + CosmosAsyncContainer c, + PartitionKey partitionKey, + Class type) { + QueryResult result = new QueryResult<>(); + for (FeedResponse page : c.readAllItems(partitionKey, type).byPage().toIterable()) { + result.results.addAll(page.getResults()); + result.diagnostics.add(page.getCosmosDiagnostics()); + } + return result; + } + private static List idsSorted(List docs) { return docs.stream() .filter(d -> d.has(ID_FIELD)) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index e4245b2788cb..b70f34cc4d6e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -1095,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; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index 819b0ead62c8..ff999b944b71 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -31,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; @@ -246,17 +247,8 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque PartitionKeyInternal partitionKey = request.getPartitionKeyInternal(); // Detect a partial (prefix) hierarchical partition key BEFORE the RNTBD request is serialized. - // For such a query the prefix's effective-partition-key sub-range - // [hash(prefix), hash(prefix) + "FF") must reach the thin-client proxy as a doc-level EPK - // filter. We convey it by setting StartEpk/EndEpk + ReadFeedKeyType=EffectivePartitionKeyRange - // on the request headers here, so that RntbdRequest.from() -> - // RntbdRequestHeaders.addStartAndEndKeys serializes them into the RNTBD request with the - // correct encoding (the hex string taken as UTF-8 bytes, not the decoded hash) - exactly as - // the Direct/RNTBD path does for the same query (see FeedRangeEpkImpl). The proxy forwards - // these to the backend as the doc-level EPK filter (see TransportSerialization - // AddStartAndEndKeysFromHeaders, where ReadFeedKeyType=EffectivePartitionKeyRange selects the - // hex-string-as-bytes encoding). Without this filter the proxy resolves the request to the - // owning physical partition and returns every co-located document (an over-span). + // 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 prefixEpkRange = null; if (request.getOperationType() != OperationType.QueryPlan) { if (partitionKey != null) { @@ -272,17 +264,9 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque // prefix effective-partition-key range was already computed and carried on the request's // feedRange (DocumentQueryExecutionContextFactory#resolveFeedRangeBasedOnPrefixContainer). // Reuse it directly instead of recomputing getEPKRangeForPrefixPartitionKey from a - // partition key we no longer have. This is what scopes the read on this path; without it - // the request falls through to the physical-partition steering below and over-spans. + // partition key we no longer have. prefixEpkRange = ((FeedRangeEpkImpl) request.getFeedRange()).getRange(); } - if (prefixEpkRange != null) { - request.getHeaders().put( - HttpConstants.HttpHeaders.READ_FEED_KEY_TYPE, - ReadFeedKeyType.EffectivePartitionKeyRange.name()); - request.getHeaders().put(HttpConstants.HttpHeaders.START_EPK, prefixEpkRange.getMin()); - request.getHeaders().put(HttpConstants.HttpHeaders.END_EPK, prefixEpkRange.getMax()); - } } RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs(request); @@ -294,12 +278,27 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque if (prefixEpkRange != null) { // Partial (prefix) hierarchical partition key - either the PK-bearing path (point/scalar - // prefix) or the parallel path that reused request.getFeedRange(). StartEpk/EndEpk + - // ReadFeedKeyType were already set on the request headers above and serialized into the RNTBD - // request as the backend's doc-level EPK filter. StartEpkHash/EndEpkHash additionally steer the - // proxy to the owning physical partition(s), mirroring the Direct/RNTBD path. - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMin())); - rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpkHash, HexConvert.hexToBytes(prefixEpkRange.getMax())); + // 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); From 63641ae03dd9bfc76344758199461b9c567e192c Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 17:32:01 -0400 Subject: [PATCH 05/15] Simplify prefix-query flag assignment in ParallelDocumentQueryExecutionContextBase Assign isPrefixPartitionKeyQuery directly from isPartialPartitionKeyQuery(...) and guard only the PARTITION_KEY-setting branch with !isPrefixPartitionKeyQuery, removing the if/else. Behavior is unchanged: a partial (prefix) key still leaves partitionKeyInternal null so the prefix FeedRangeEpkImpl survives createDocumentServiceRequestWithFeedRange. --- ...llelDocumentQueryExecutionContextBase.java | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java index e34b46c00f4f..5201dc24c796 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java @@ -83,17 +83,15 @@ protected void initialize( 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())) { - // Partial (prefix) hierarchical partition key: partitionKeyInternal stays null (no - // PARTITION_KEY header). The narrow prefix effective-partition-key range is already - // carried on the feedRange passed below; mark the request so downstream store models - // (e.g. ThinClientStoreModel) can reuse it as a doc-level EPK filter instead of - // over-spanning to the whole physical partition. - isPrefixPartitionKeyQuery = true; - } else { + // A partial (prefix) hierarchical partition key: partitionKeyInternal stays null (no + // PARTITION_KEY header, which avoids the feedRange normalization in + // createDocumentServiceRequestWithFeedRange below). The narrow prefix + // effective-partition-key range is already carried on the feedRange passed below; mark + // the request so downstream store models (e.g. ThinClientStoreModel) can reuse it as a + // doc-level EPK filter instead of over-spanning to the whole physical partition. + isPrefixPartitionKeyQuery = + PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey()); + if (!isPrefixPartitionKeyQuery) { partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey()); headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson()); } From 0b746782235bea662a5f8d349425969ac7926926 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 17:35:03 -0400 Subject: [PATCH 06/15] Add CHANGELOG entry for thin-client prefix hierarchical partition key over-span fix --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 25b8545acfff..bc21c6e19a63 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/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). From 817d948c3e3a98382d0906e5e23e4637904f4664 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 17:48:39 -0400 Subject: [PATCH 07/15] Minimize prefix-query flag change vs main in ParallelDocumentQueryExecutionContextBase Keep main's single 'if (!isPartialPartitionKeyQuery)' structure; only capture the predicate in isPrefixPartitionKeyQuery and pass it to request.setPrefixPartitionKeyQuery. Removes the if/else. --- .../ParallelDocumentQueryExecutionContextBase.java | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java index 5201dc24c796..2d4bff199f9a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java @@ -83,14 +83,10 @@ protected void initialize( PartitionKeyInternal partitionKeyInternal = null; boolean isPrefixPartitionKeyQuery = false; if (cosmosQueryRequestOptions.getPartitionKey() != null && cosmosQueryRequestOptions.getPartitionKey() != PartitionKey.NONE) { - // A partial (prefix) hierarchical partition key: partitionKeyInternal stays null (no - // PARTITION_KEY header, which avoids the feedRange normalization in - // createDocumentServiceRequestWithFeedRange below). The narrow prefix - // effective-partition-key range is already carried on the feedRange passed below; mark - // the request so downstream store models (e.g. ThinClientStoreModel) can reuse it as a - // doc-level EPK filter instead of over-spanning to the whole physical partition. - isPrefixPartitionKeyQuery = - PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey()); + // 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. + isPrefixPartitionKeyQuery = PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey()); if (!isPrefixPartitionKeyQuery) { partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey()); headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson()); From 8b5b1a58edd5a9b7863150cec12d22444dd6bbf2 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 18:07:12 -0400 Subject: [PATCH 08/15] Order thin-client prefix detection to prefer the cached feed-range path In ThinClientStoreModel.wrapInHttpRequest, check the cached prefix feed-range path (isPrefixPartitionKeyQuery + FeedRangeEpkImpl) first and fall back to recomputing from the partition key only for the PK-bearing prefix case. Behavior is unchanged (the two cases are mutually exclusive). --- .../implementation/ThinClientStoreModel.java | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index ff999b944b71..83bcd38d9d37 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -251,21 +251,23 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque // as the RNTBD EPK headers in the consolidated block after RntbdRequest.from() below. Range prefixEpkRange = null; if (request.getOperationType() != OperationType.QueryPlan) { - if (partitionKey != null) { + 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); } - } else if (request.isPrefixPartitionKeyQuery() && request.getFeedRange() instanceof FeedRangeEpkImpl) { - // Parallel prefix-query path: 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 it directly instead of recomputing getEPKRangeForPrefixPartitionKey from a - // partition key we no longer have. - prefixEpkRange = ((FeedRangeEpkImpl) request.getFeedRange()).getRange(); } } From c2fa5f140d99c3898976e61b7048699c3c828255 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 1 Jul 2026 18:12:43 -0400 Subject: [PATCH 09/15] Address Copilot review: validate all requests in thin-client endpoint assertion TestSuiteBase.assertThinClientEndpointUsed now validates every request instead of early-returning on the first thin-client match, so a mixed scenario (some data requests via the classic gateway) fails; QueryPlan requests remain exempt and null endpoints are handled. ThinClientTestBase.assertThinClientEndpointUsed delegates to the shared TestSuiteBase implementation, removing duplicated logic that could NPE on a null endpoint and reject valid QueryPlan-via-gateway scenarios. --- .../com/azure/cosmos/rx/TestSuiteBase.java | 34 ++++++++----------- .../azure/cosmos/rx/ThinClientTestBase.java | 17 +++------- 2 files changed, 19 insertions(+), 32 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index bda1674d6312..5f8b17df3c33 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -2176,30 +2176,26 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnosticsContext ctx) assertThat(requests).isNotNull(); assertThat(requests.size()).isPositive(); - boolean hasNonQueryPlanRequest = false; + // 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)) { - // A thin-client request exists -> invariant satisfied. - return; - } - // requestType has the form ":" (OperationType.QueryPlan - // stringifies to "QueryPlan"). In thin-client mode QueryPlan calls are resolved via the - // classic gateway, so they are the only requests allowed to target a non-thin-client endpoint. + // stringifies to "QueryPlan"). String requestType = requestInfo.getRequestType(); - if (requestType == null || !requestType.endsWith(":QueryPlan")) { - hasNonQueryPlanRequest = true; + if (requestType != null && requestType.endsWith(":QueryPlan")) { + continue; } - } - // Reaching here means no request targeted the thin-client proxy endpoint. That is acceptable - // only when every request was a QueryPlan call; a non-QueryPlan (data) request on a non-thin-client - // endpoint violates the thin-client routing invariant. - assertThat(hasNonQueryPlanRequest) - .as("No request targeting thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR - + ") and at least one non-QueryPlan request was present") - .isFalse(); + 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) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java index 379d12bd3720..c680f88c855c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java @@ -76,19 +76,10 @@ protected ObjectNode createTestDocument(String id, String mypk) { * Asserts that all requests in the diagnostics were routed through the thin client endpoint. */ protected static void assertThinClientEndpointUsed(CosmosDiagnostics diagnostics) { - assertThat(diagnostics).isNotNull(); - CosmosDiagnosticsContext ctx = diagnostics.getDiagnosticsContext(); - assertThat(ctx).isNotNull(); - Collection requests = ctx.getRequestInfo(); - assertThat(requests).isNotNull(); - assertThat(requests.size()).isPositive(); - int requestCountAgainstThinClientEndpoint = 0; - for (CosmosDiagnosticsRequestInfo requestInfo : requests) { - if (requestInfo.getEndpoint().contains(THIN_CLIENT_ENDPOINT_INDICATOR)) { - requestCountAgainstThinClientEndpoint++; - } - } - assertThat(requestCountAgainstThinClientEndpoint).isEqualTo(requests.size()); + // 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); } /** From a56db9ad44c4f5a238bb869edb1d585a8ccf9b18 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 2 Jul 2026 17:21:16 -0400 Subject: [PATCH 10/15] Simplify thin-client prefix EPK handling to reuse HTTP EPK headers Replace the explicit MULTI_HASH prefix-detection and pre-serialization range computation with a header-presence branch: when the request already carries START_EPK/END_EPK HTTP headers (populated upstream by FeedRangeEpkImpl for prefix hierarchical partition keys and explicit EPK-range reads), set only the additive StartEpkHash/EndEpkHash RNTBD tokens. The ReadFeedKeyType/StartEpk/ EndEpk string tokens are copied verbatim from those HTTP headers by RntbdRequestHeaders#addStartAndEndKeys during RntbdRequest.from(), so they no longer need to be set here. Removes now-unused imports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../implementation/ThinClientStoreModel.java | 64 ++++--------------- 1 file changed, 13 insertions(+), 51 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index 83bcd38d9d37..a4642c406091 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -14,12 +14,8 @@ 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; @@ -31,7 +27,6 @@ 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; @@ -246,31 +241,6 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque 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 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(); @@ -278,29 +248,21 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque RntbdRequest rntbdRequest = RntbdRequest.from(rntbdRequestArgs); - 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)); + String startEpk = request.getHeaders().get(HttpConstants.HttpHeaders.START_EPK); + String endEpk = request.getHeaders().get(HttpConstants.HttpHeaders.END_EPK); + if (startEpk != null && endEpk != null) { + // The request already carries the effective-partition-key range as HTTP headers (set by + // FeedRangeEpkImpl#populateFeedRangeFilteringHeaders). This covers a partial (prefix) + // hierarchical partition key query as well as any explicit EPK-range read. The + // ReadFeedKeyType/StartEpk/EndEpk RNTBD string tokens are copied verbatim from those HTTP + // headers by RntbdRequestHeaders#addStartAndEndKeys during RntbdRequest.from() above, so only + // the StartEpkHash/EndEpkHash tokens (the decoded hash bytes) are additive here. They steer the + // thin-client proxy to the owning physical partition(s); without them the proxy resolves the + // request to the whole physical partition and returns every co-located document (an over-span). rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.StartEpkHash, - HexConvert.hexToBytes(prefixEpkRange.getMin())); + HexConvert.hexToBytes(startEpk)); rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EndEpkHash, - HexConvert.hexToBytes(prefixEpkRange.getMax())); + HexConvert.hexToBytes(endEpk)); } else if (partitionKey != null) { byte[] epk = partitionKey.getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition()); rntbdRequest.setHeaderValue(RntbdConstants.RntbdRequestHeader.EffectivePartitionKey, epk); From 03a7c672b95180a8fcd9e42516dec7013c913a37 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 2 Jul 2026 19:50:51 -0400 Subject: [PATCH 11/15] Remove dead prefix-query flag and expand HPK thin-client test coverage Following the simplification of the prefix EPK fix to reuse the existing StartEpk/EndEpk HTTP headers (a56db9ad44c), the RxDocumentServiceRequest prefixPartitionKeyQuery flag is no longer read anywhere. Remove the field, its accessors, the clone copy, and the caller assignment in ParallelDocumentQueryExecutionContextBase. Test changes: - Retrofit CosmosMultiHashTest into the thinclient TestNG group so the MULTI_HASH prefix over-span coverage also runs against GatewayV2 (proxy :10250) over HTTP/2, not just the emulator gateway. - ThinClientQueryE2ETest: correct the prefix-scope Javadoc to describe the actual header-copy behavior and add a deterministic prefix-partition-key readAllItems regression. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/CosmosMultiHashTest.java | 40 +++++++- .../cosmos/rx/ThinClientQueryE2ETest.java | 97 +++++++++++++------ .../RxDocumentServiceRequest.java | 15 --- ...llelDocumentQueryExecutionContextBase.java | 1 - 4 files changed, 103 insertions(+), 50 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java index 5ece83be3156..053bd248f57a 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java +++ b/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 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 docs = createItems(); @@ -645,7 +670,12 @@ private void deleteAllItems() { private void testPartialPKContinuationToken() { String requestContinuation = null; List 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"; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java index ff312ed85081..d23dc1b77568 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -1011,9 +1011,12 @@ public void testHierarchicalPartitionKeyCrossPartition() { * Hierarchical prefix (half-open) range. Setting only the FIRST hierarchical component * (tenantId) on the options makes the proxy emit a half-open prefix range whose max is the * "Infinity"-suffixed sibling of the min - the prefix-range boundary case of the conversion. The - * full key (both components) instead resolves to a single point range. Asserts the prefix query - * returns exactly the tenant's users and the full key returns exactly one document, with parity - * to the Direct baseline. + * full key (both components) instead resolves to a single point range. Iterates the prefix + * assertion over ALL seeded tenants so that, with fewer physical partitions than + * {@code HIER_TENANTS}, at least two co-located tenants deterministically exercise the over-span + * case, and adds a nonexistent-tenant prefix that must return zero. Asserts each prefix query + * returns exactly the tenant's users and the full key returns exactly one document, all with + * parity to the Direct baseline. */ @Test(groups = {"thinclient"}, timeOut = TIMEOUT) public void testHierarchicalPrefixHalfOpenRange() { @@ -1025,9 +1028,20 @@ public void testHierarchicalPrefixHalfOpenRange() { fullKeyOptions.setPartitionKey(new PartitionKeyBuilder().add(tenant0).add(user0).build()); assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", fullKeyOptions, 1); - CosmosQueryRequestOptions prefixOptions = new CosmosQueryRequestOptions(); - prefixOptions.setPartitionKey(new PartitionKeyBuilder().add(tenant0).build()); - assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", prefixOptions, HIER_USERS_PER_TENANT); + // Iterate the prefix (half-open range) assertion over every tenant. With fewer physical + // partitions than HIER_TENANTS, at least two tenants are co-located, so the over-span + // regression reproduces deterministically instead of depending on tenant-0's placement. + for (int t = 0; t < HIER_TENANTS; t++) { + String tenant = hierarchicalKeys.get(t * HIER_USERS_PER_TENANT)[0]; + CosmosQueryRequestOptions prefixOptions = new CosmosQueryRequestOptions(); + prefixOptions.setPartitionKey(new PartitionKeyBuilder().add(tenant).build()); + assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", prefixOptions, HIER_USERS_PER_TENANT); + } + + // Negative guard: a nonexistent tenant prefix must scope to zero docs on both paths. + CosmosQueryRequestOptions absentOptions = new CosmosQueryRequestOptions(); + absentOptions.setPartitionKey(new PartitionKeyBuilder().add("tenant-absent-" + UUID.randomUUID()).build()); + assertHierarchicalScopedMatch("SELECT * FROM c ORDER BY c.idx", absentOptions, 0); } /** @@ -1636,31 +1650,49 @@ public void testReadManyByPartitionKeysWithParameterizedCustomQuery() { /** * readAllItems (ReadFeed) with a partial (prefix) hierarchical partition key - only the first - * component (/tenantId). Must return exactly that tenant's {@code HIER_USERS_PER_TENANT} - * documents, matching the Direct baseline. If the thin client over-spanned to the physical - * partition it would return co-located docs from other tenants, which the count/ID parity and the - * per-doc tenant assertion catch. + * component (/tenantId). Iterates over ALL seeded tenants (not just one): at 12000 RU/s the + * container has fewer physical partitions than {@code HIER_TENANTS}, so by the pigeonhole + * principle at least two tenants are co-located on the same physical partition, making the + * over-span regression deterministically reproducible instead of depending on tenant-0's + * placement. Each prefix read must return exactly that tenant's {@code HIER_USERS_PER_TENANT} + * documents, matching the Direct baseline; an over-span to the physical partition would surface + * co-located docs from other tenants, which the count/ID parity and per-doc tenant assertion + * catch. A final nonexistent-tenant prefix must return zero documents - a co-location-independent + * negative guard. */ @Test(groups = {"thinclient"}, timeOut = TIMEOUT) public void testHierarchicalReadAllItemsPrefixPartitionKey() { - String tenant0 = hierarchicalKeys.get(0)[0]; - PartitionKey prefixKey = new PartitionKeyBuilder().add(tenant0).build(); + for (int t = 0; t < HIER_TENANTS; t++) { + String tenant = hierarchicalKeys.get(t * HIER_USERS_PER_TENANT)[0]; + PartitionKey prefixKey = new PartitionKeyBuilder().add(tenant).build(); - QueryResult directResult = drainReadAllItems(directHierarchicalContainer, prefixKey, ObjectNode.class); - QueryResult tcResult = drainReadAllItems(thinClientHierarchicalContainer, prefixKey, ObjectNode.class); - for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } + QueryResult directResult = drainReadAllItems(directHierarchicalContainer, prefixKey, ObjectNode.class); + QueryResult tcResult = drainReadAllItems(thinClientHierarchicalContainer, prefixKey, ObjectNode.class); + for (CosmosDiagnostics d : tcResult.diagnostics) { assertThinClientEndpointUsed(d); } - assertThat(directResult.results.size()) - .as("Direct readAllItems prefix count").isEqualTo(HIER_USERS_PER_TENANT); - assertThat(tcResult.results.size()) - .as("Thin-client readAllItems prefix over-span vs Direct").isEqualTo(directResult.results.size()); - assertThat(idsSorted(tcResult.results)) - .as("readAllItems prefix ID mismatch").isEqualTo(idsSorted(directResult.results)); - // Every returned doc must belong to the prefixed tenant (no over-span into co-located tenants). - for (ObjectNode doc : tcResult.results) { - assertThat(doc.get(TENANT_FIELD).asText()) - .as("readAllItems prefix returned a foreign tenant's document").isEqualTo(tenant0); + assertThat(directResult.results.size()) + .as("Direct readAllItems prefix count for " + tenant).isEqualTo(HIER_USERS_PER_TENANT); + assertThat(tcResult.results.size()) + .as("Thin-client readAllItems prefix over-span vs Direct for " + tenant).isEqualTo(directResult.results.size()); + assertThat(idsSorted(tcResult.results)) + .as("readAllItems prefix ID mismatch for " + tenant).isEqualTo(idsSorted(directResult.results)); + // Every returned doc must belong to the prefixed tenant (no over-span into co-located tenants). + for (ObjectNode doc : tcResult.results) { + assertThat(doc.get(TENANT_FIELD).asText()) + .as("readAllItems prefix returned a foreign tenant's document").isEqualTo(tenant); + } } + + // Negative guard: a prefix matching no seeded tenant must return zero docs on both paths. + // Independent of physical co-location - a broken over-span would surface co-located docs even + // though this tenant was never seeded. + PartitionKey absentPrefix = new PartitionKeyBuilder().add("tenant-absent-" + UUID.randomUUID()).build(); + QueryResult directAbsent = drainReadAllItems(directHierarchicalContainer, absentPrefix, ObjectNode.class); + QueryResult tcAbsent = drainReadAllItems(thinClientHierarchicalContainer, absentPrefix, ObjectNode.class); + for (CosmosDiagnostics d : tcAbsent.diagnostics) { assertThinClientEndpointUsed(d); } + assertThat(directAbsent.results.size()).as("Direct readAllItems nonexistent-tenant prefix count").isEqualTo(0); + assertThat(tcAbsent.results.size()) + .as("Thin-client readAllItems nonexistent-tenant prefix over-span").isEqualTo(0); } /** @@ -1685,10 +1717,17 @@ public void testHierarchicalReadAllItemsFullPartitionKey() { } /** - * readManyByPartitionKeys with a partial (prefix) hierarchical key. groupPartitionKeysByPhysicalPartition - * expands the prefix to its effective-partition-key range; the thin client must apply that as a - * doc-level filter and return exactly the tenant's {@code HIER_USERS_PER_TENANT} documents, - * matching Direct rather than over-spanning the physical partition. + * readManyByPartitionKeys with a partial (prefix) hierarchical key: validates that a prefix + * readMany request routes over the thin client and returns exactly the tenant's documents, + * matching Direct. + * + *

Note: readManyByPartitionKeys groups the supplied keys by physical partition and emits a + * generated {@code WHERE} predicate on the prefix key component(s), so tenant isolation here is + * enforced by that SQL predicate rather than by the RNTBD prefix-EPK range header. This is + * therefore a routing + SQL-predicate parity guard against Direct, not a direct guard of the + * over-span fix. The prefix-EPK header path itself (query/readAll with only a prefix key and no + * generated predicate) is exercised by {@link #testHierarchicalReadAllItemsPrefixPartitionKey()} + * and {@link #testHierarchicalPrefixHalfOpenRange()}. */ @Test(groups = {"thinclient"}, timeOut = TIMEOUT) public void testHierarchicalReadManyByPartitionKeysPrefixPartitionKey() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index b70f34cc4d6e..03f27156f62f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -78,12 +78,6 @@ private static CosmosItemSerializer internalDefaultSerializer() { private Range 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 @@ -912,14 +906,6 @@ 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; } @@ -1095,7 +1081,6 @@ 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; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java index 2d4bff199f9a..18b5f91e56e4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java @@ -97,7 +97,6 @@ protected void initialize( queryOptionsAccessor().setCollectionRid(cosmosQueryRequestOptions, collection.getResourceId()); RxDocumentServiceRequest request = this.createDocumentServiceRequestWithFeedRange(headers, querySpecForInit, partitionKeyInternal, feedRange, collection.getResourceId(), cosmosQueryRequestOptions.getThroughputControlGroupName()); - request.setPrefixPartitionKeyQuery(isPrefixPartitionKeyQuery); return request; }; From 848c57ef863d39c2b3c8a707c01d2773acbeccc3 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 2 Jul 2026 20:25:29 -0400 Subject: [PATCH 12/15] Revert ParallelDocumentQueryExecutionContextBase to main The prefix-query flag on RxDocumentServiceRequest was removed in 03a7c672b95 after the fix was simplified to reuse the StartEpk/EndEpk HTTP headers. The only functional edit in this file (the setter call) went away with it, leaving behind unnecessary refactor noise. Restore the file to its upstream/main state so the hotfix touches no query execution code. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../query/ParallelDocumentQueryExecutionContextBase.java | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java index 18b5f91e56e4..5774a9dde4fe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java @@ -81,13 +81,11 @@ 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. - isPrefixPartitionKeyQuery = PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey()); - if (!isPrefixPartitionKeyQuery) { + if (!PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey())) { partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey()); headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson()); } @@ -95,9 +93,8 @@ protected void initialize( queryOptionsAccessor().setPartitionKeyDefinition(cosmosQueryRequestOptions, collection.getPartitionKey()); queryOptionsAccessor().setCollectionRid(cosmosQueryRequestOptions, collection.getResourceId()); - RxDocumentServiceRequest request = this.createDocumentServiceRequestWithFeedRange(headers, querySpecForInit, partitionKeyInternal, feedRange, + return this.createDocumentServiceRequestWithFeedRange(headers, querySpecForInit, partitionKeyInternal, feedRange, collection.getResourceId(), cosmosQueryRequestOptions.getThroughputControlGroupName()); - return request; }; Function>> executeFunc = (request) -> this.executeRequestAsync( From 3481541eebf0a66f360ecb65394ef0d8e13adfb0 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 5 Jul 2026 15:06:03 -0400 Subject: [PATCH 13/15] Address final review comments on thin-client EPK over-span fix Comment 1: Tighten the StartEpkHash/EndEpkHash guard in ThinClientStoreModel#wrapInHttpRequest to require READ_FEED_KEY_TYPE == EffectivePartitionKeyRange in addition to the StartEpk/EndEpk headers. RntbdRequestHeaders#addStartAndEndKeys only emits the ReadFeedKeyType RNTBD token when that header is present, so the consumer guard now matches the producer contract and keeps the string and binary EPK tokens coherent. Comment 2: Add testExplicitFeedRangeShardedReadParity, which exercises the explicit user-supplied FeedRange (EPK-range) read path directly by sharding the multi-physical-partition container into its physical feed ranges and asserting per-shard parity with the Direct baseline plus a clean, non-overlapping union over the full fan-out. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../cosmos/rx/ThinClientQueryE2ETest.java | 59 +++++++++++++++++++ .../implementation/ThinClientStoreModel.java | 16 +++-- 2 files changed, 70 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java index d23dc1b77568..39f63510ff80 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -27,6 +27,7 @@ import com.azure.cosmos.models.CosmosVectorIndexSpec; import com.azure.cosmos.models.CosmosVectorIndexType; import com.azure.cosmos.models.ExcludedPath; +import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.IncludedPath; import com.azure.cosmos.models.IndexingMode; @@ -994,6 +995,64 @@ public void testFullRangeInfinityBoundary() { .isGreaterThan(1); } + /** + * Explicit user-supplied {@link FeedRange} (EPK-range) scoped read parity. + *

+ * The over-span fix in {@code ThinClientStoreModel#wrapInHttpRequest} keys off the presence of the + * StartEpk/EndEpk headers (paired with {@code READ_FEED_KEY_TYPE = EffectivePartitionKeyRange}), + * which {@code FeedRangeEpkImpl} emits whenever a caller supplies an explicit {@link FeedRange} on + * the query options. The prefix-HPK tests exercise that header path indirectly through the query + * plan; this test exercises it DIRECTLY by sharding the multi-physical-partition container into its + * physical feed ranges and querying each one, so a broken EPK-range guard would over-span a shard + * and surface documents belonging to a sibling partition. + *

+ * For every physical feed range asserts: (1) the thin client routed through the :10250 endpoint, + * and (2) the thin-client rows for the shard exactly equal the Direct baseline for the SAME shard - + * an over-span would make the thin-client set a strict superset of the Direct set. Across all shards + * it further asserts the union exactly reconstructs the full unsharded fan-out with no duplicates, + * proving the shards are a clean, non-overlapping partition of the key space. + */ + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testExplicitFeedRangeShardedReadParity() { + List feedRanges = directCrossPartitionContainer.getFeedRanges().block(); + assertThat(feedRanges.size()) + .as("multi-physical-partition container must expose more than one feed range") + .isGreaterThan(1); + + List unionThinClientIds = new ArrayList<>(); + for (FeedRange feedRange : feedRanges) { + CosmosQueryRequestOptions directOptions = new CosmosQueryRequestOptions().setFeedRange(feedRange); + CosmosQueryRequestOptions tcOptions = new CosmosQueryRequestOptions().setFeedRange(feedRange); + + QueryResult directShard = + drainQuery(directCrossPartitionContainer, "SELECT * FROM c", directOptions, ObjectNode.class); + QueryResult tcShard = + drainQuery(thinClientCrossPartitionContainer, "SELECT * FROM c", tcOptions, ObjectNode.class); + for (CosmosDiagnostics d : tcShard.diagnostics) { assertThinClientEndpointUsed(d); } + + // Per-shard parity: the thin client must return EXACTLY the Direct baseline's documents for + // this feed range. A broken EPK-range guard would over-span and pull in sibling-shard docs, + // making the thin-client set a strict superset of the Direct set. + assertThat(tcShard.results.size()) + .as("Thin-client feed-range shard count vs Direct for " + feedRange) + .isEqualTo(directShard.results.size()); + assertThat(idsSorted(tcShard.results)) + .as("Thin-client feed-range shard over-span vs Direct for " + feedRange) + .isEqualTo(idsSorted(directShard.results)); + + unionThinClientIds.addAll(idsSorted(tcShard.results)); + } + + // The shards must be a clean partition of the whole container: their union reconstructs the full + // unsharded fan-out exactly, with no duplicates (which would indicate two shards overlap). + QueryResult fullFanOut = drainQuery( + directCrossPartitionContainer, "SELECT * FROM c", new CosmosQueryRequestOptions(), ObjectNode.class); + List unionSorted = unionThinClientIds.stream().sorted().collect(Collectors.toList()); + assertThat(unionSorted) + .as("union of per-feed-range thin-client shards must equal the full fan-out with no duplicates") + .isEqualTo(idsSorted(fullFanOut.results)); + } + /** * Hierarchical (MULTI_HASH) partition key cross-partition parity. Cross-partition queries * over the /tenantId,/userId container force the proxy to emit MULTI-COMPONENT PartitionKeyInternal diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index a4642c406091..815d88995afe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -250,12 +250,18 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque String startEpk = request.getHeaders().get(HttpConstants.HttpHeaders.START_EPK); String endEpk = request.getHeaders().get(HttpConstants.HttpHeaders.END_EPK); - if (startEpk != null && endEpk != null) { - // The request already carries the effective-partition-key range as HTTP headers (set by + String readFeedKeyType = request.getHeaders().get(HttpConstants.HttpHeaders.READ_FEED_KEY_TYPE); + if (startEpk != null && endEpk != null + && ReadFeedKeyType.EffectivePartitionKeyRange.name().equalsIgnoreCase(readFeedKeyType)) { + // The request already carries the effective-partition-key range as HTTP headers (set together by // FeedRangeEpkImpl#populateFeedRangeFilteringHeaders). This covers a partial (prefix) - // hierarchical partition key query as well as any explicit EPK-range read. The - // ReadFeedKeyType/StartEpk/EndEpk RNTBD string tokens are copied verbatim from those HTTP - // headers by RntbdRequestHeaders#addStartAndEndKeys during RntbdRequest.from() above, so only + // hierarchical partition key query as well as any explicit EPK-range read. All three headers + // (ReadFeedKeyType=EffectivePartitionKeyRange, StartEpk, EndEpk) are required for the backend to + // interpret the range: RntbdRequestHeaders#addStartAndEndKeys reads each independently and only + // emits the ReadFeedKeyType RNTBD token when READ_FEED_KEY_TYPE is present, so we guard on it here + // as well to keep the StartEpk/EndEpk string tokens and the StartEpkHash/EndEpkHash binary tokens + // coherent. The ReadFeedKeyType/StartEpk/EndEpk RNTBD string tokens are copied verbatim from those + // HTTP headers by RntbdRequestHeaders#addStartAndEndKeys during RntbdRequest.from() above, so only // the StartEpkHash/EndEpkHash tokens (the decoded hash bytes) are additive here. They steer the // thin-client proxy to the owning physical partition(s); without them the proxy resolves the // request to the whole physical partition and returns every co-located document (an over-span). From 274dc69f7f2f957797356d1e5db697231728be7f Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 5 Jul 2026 16:46:58 -0400 Subject: [PATCH 14/15] Add mode-aware thin-client endpoint routing assertions to CosmosMultiHashTest Verify data requests route through the thin-client endpoint (:10250) when the thinclient group is active, and through the gateway otherwise. Also removes redundant in-code thin-client enablement (owned by the TestNG config / Maven profile). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/CosmosMultiHashTest.java | 41 +++++++++++++++---- .../cosmos/rx/ThinClientQueryE2ETest.java | 2 - .../azure/cosmos/rx/ThinClientTestBase.java | 14 +------ 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java index 053bd248f57a..9dfc522bcf02 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java @@ -43,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 { @@ -54,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) { @@ -62,6 +67,7 @@ public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) { @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) public void before_CosmosMultiHashTest() { + thinClientEnabled = false; client = getClientBuilder().buildClient(); initDatabaseAndContainers(); } @@ -69,10 +75,11 @@ public void before_CosmosMultiHashTest() { // 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. + // 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() { - System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); + thinClientEnabled = true; client = createGatewayRxDocumentClient( TestConfigurations.HOST, null, true, null, true, true, true).buildClient(); initDatabaseAndContainers(); @@ -125,7 +132,16 @@ public void afterClass_thinClient() { logger.info("starting cleanup (thin client)...."); safeDeleteSyncDatabase(createdDatabase); safeCloseSyncClient(client); - System.clearProperty("COSMOS.THINCLIENT_ENABLED"); + } + + // 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) @@ -145,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()); } @@ -211,6 +228,7 @@ public void readManySupportsNestedPartitionKeyPaths() { FeedResponse documentFeedResponse = createdNestedPathContainer.readMany(itemList, ObjectNode.class); validateResponse(documentFeedResponse, itemList); + assertEndpointRouting(documentFeedResponse.getCosmosDiagnostics()); } @Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT) @@ -228,7 +246,13 @@ public void readAllItemsSupportsNestedPartitionKeyPaths() { CosmosPagedIterable readAllResults = createdNestedPathContainer.readAllItems(new PartitionKey(city), ObjectNode.class); - assertThat(readAllResults.stream().map(item -> item.get("id").asText()).collect(Collectors.toList())) + List readAllIds = new ArrayList<>(); + for (FeedResponse 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()); } @@ -555,6 +579,7 @@ private void multiHashQueryTests() { FeedResponse 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(); @@ -670,10 +695,11 @@ private void deleteAllItems() { private void testPartialPKContinuationToken() { String requestContinuation = null; List receivedDocuments = new ArrayList<>(); - // 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. + // 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 = - Boolean.parseBoolean(System.getProperty("COSMOS.THINCLIENT_ENABLED")) + thinClientEnabled ? createGatewayRxDocumentClient(TestConfigurations.HOST, null, true, null, true, true, true).buildAsyncClient() : getClientBuilder().buildAsyncClient(); CosmosAsyncDatabase cosmosAsyncDatabase = new CosmosAsyncDatabase(createdDatabase.getId(), asyncClient); @@ -701,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(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java index 39f63510ff80..f99477d4aecb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -128,7 +128,6 @@ public void before_ThinClientQueryE2ETest() { this.directContainer = getSharedMultiPartitionCosmosContainer(this.directClient); // 2. Gateway V2 thin client (system under test) - ThinClientTestBase.enableThinClientForTest(); CosmosClientBuilder thinClientBuilder = createGatewayRxDocumentClient( TestConfigurations.HOST, null, true, null, true, true, true); this.thinClient = thinClientBuilder.buildAsyncClient(); @@ -301,7 +300,6 @@ public void afterClass() { logger.warn("Bulk delete of seeded docs failed: {}", e.getMessage()); } } - ThinClientTestBase.clearThinClientForTest(); if (directCrossPartitionContainer != null) { safeDeleteContainer(directCrossPartitionContainer); } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java index c680f88c855c..6f4a3b85e419 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java @@ -38,7 +38,6 @@ protected ThinClientTestBase(CosmosClientBuilder 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); @@ -48,20 +47,11 @@ public void before_ThinClientTest() { @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). */ @@ -75,7 +65,7 @@ protected ObjectNode createTestDocument(String id, String mypk) { /** * Asserts that all requests in the diagnostics were routed through the thin client endpoint. */ - protected static void assertThinClientEndpointUsed(CosmosDiagnostics diagnostics) { + public 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. @@ -86,7 +76,7 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnostics 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) { + public static void assertGatewayEndpointUsed(CosmosDiagnostics diagnostics) { assertThat(diagnostics).isNotNull(); CosmosDiagnosticsContext ctx = diagnostics.getDiagnosticsContext(); assertThat(ctx).isNotNull(); From 0c607e615a51271fc53c55e73c38a69fd77cc004 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 6 Jul 2026 13:04:31 -0400 Subject: [PATCH 15/15] Address xinlian12 review comments: docstring accuracy, dead scaffolding cleanup, container-leak fix - ThinClientQueryE2ETest: clarify prefix-HPK test docstring as routing+SQL-predicate parity guard; cross-reference genuine EPK-header guards - ThinClientTestBase: convert to stateless final helper holder (drop unused lifecycle scaffolding) - ThinClientQueryE2ETest: best-effort container cleanup in setup catch block to avoid leaking provisioned containers Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../cosmos/rx/ThinClientQueryE2ETest.java | 50 +++++++++++------ .../azure/cosmos/rx/ThinClientTestBase.java | 55 ++++--------------- 2 files changed, 43 insertions(+), 62 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java index f99477d4aecb..ac9e6083665c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientQueryE2ETest.java @@ -147,7 +147,15 @@ public void before_ThinClientQueryE2ETest() { // multi-component QueryPlan range-conversion coverage. seedHierarchicalData(); } catch (Exception e) { - // Clean up any clients that were successfully created before the failure + // Best-effort cleanup of any resources created before the failure. Dedicated container + // deletes must run before the owning clients are closed, and any cleanup failure must + // not mask the original setup exception. + try { + if (this.directCrossPartitionContainer != null) { safeDeleteContainer(this.directCrossPartitionContainer); } + if (this.directHierarchicalContainer != null) { safeDeleteContainer(this.directHierarchicalContainer); } + } catch (Exception cleanupError) { + logger.warn("Cleanup after setup failure did not fully succeed: {}", cleanupError.getMessage()); + } if (this.thinClient != null) { this.thinClient.close(); this.thinClient = null; } if (this.directClient != null) { this.directClient.close(); this.directClient = null; } throw e; @@ -1698,24 +1706,32 @@ public void testReadManyByPartitionKeysWithParameterizedCustomQuery() { } // ==================== Partial (Prefix) HPK read Tests ==================== - // Regression coverage for the thin-client MULTI_HASH prefix EPK over-span fix. A partial - // hierarchical key (only /tenantId, omitting /userId) must scope the read to that tenant's - // effective-partition-key sub-range [hash(prefix), hash(prefix)+"FF") as a doc-level filter, - // NOT resolve to the whole owning physical partition (which would return co-located documents - // from other tenants). Direct TCP is the baseline; the thin client must match it exactly, so a - // count/ID mismatch is exactly the over-span regression these tests guard against. + // Direct-vs-thin-client parity coverage for partial hierarchical keys (only /tenantId, omitting + // /userId). Direct TCP is the baseline; the thin client must return exactly the same documents. + // + // The isolation mechanism differs by API and is called out per-test below: + // - readAllItems / readManyByPartitionKeys with a prefix generate a SQL WHERE predicate on the + // present component(s) and route by the resolved physical PK-range id, so foreign co-located + // tenants are filtered by that predicate (a routing + SQL-predicate parity guard). + // - The RNTBD prefix-EPK range path (a raw query with only a prefix key and no generated + // predicate, where isolation is purely by the [hash(prefix), hash(prefix)+"FF") sub-range) is + // the direct guard of the thin-client MULTI_HASH prefix EPK over-span fix. It is exercised by + // testHierarchicalPrefixHalfOpenRange() and testExplicitFeedRangeShardedReadParity(). /** * readAllItems (ReadFeed) with a partial (prefix) hierarchical partition key - only the first * component (/tenantId). Iterates over ALL seeded tenants (not just one): at 12000 RU/s the * container has fewer physical partitions than {@code HIER_TENANTS}, so by the pigeonhole - * principle at least two tenants are co-located on the same physical partition, making the - * over-span regression deterministically reproducible instead of depending on tenant-0's - * placement. Each prefix read must return exactly that tenant's {@code HIER_USERS_PER_TENANT} - * documents, matching the Direct baseline; an over-span to the physical partition would surface - * co-located docs from other tenants, which the count/ID parity and per-doc tenant assertion - * catch. A final nonexistent-tenant prefix must return zero documents - a co-location-independent - * negative guard. + * principle at least two tenants are co-located on the same physical partition. Each prefix read + * must return exactly that tenant's {@code HIER_USERS_PER_TENANT} documents, matching the Direct + * baseline; a final nonexistent-tenant prefix must return zero documents. + * + *

Note: readAllItems with a prefix key generates a SQL {@code WHERE} predicate on the present + * component(s) ({@code createLogicalPartitionScanQuerySpec}) and routes by the resolved physical + * PK-range id, so tenant isolation here is enforced by that predicate rather than by an RNTBD + * prefix-EPK range header. This is therefore a routing + SQL-predicate parity guard against + * Direct, not a direct guard of the prefix-EPK over-span fix (see + * {@link #testHierarchicalPrefixHalfOpenRange()} for that path). */ @Test(groups = {"thinclient"}, timeOut = TIMEOUT) public void testHierarchicalReadAllItemsPrefixPartitionKey() { @@ -1782,9 +1798,9 @@ public void testHierarchicalReadAllItemsFullPartitionKey() { * generated {@code WHERE} predicate on the prefix key component(s), so tenant isolation here is * enforced by that SQL predicate rather than by the RNTBD prefix-EPK range header. This is * therefore a routing + SQL-predicate parity guard against Direct, not a direct guard of the - * over-span fix. The prefix-EPK header path itself (query/readAll with only a prefix key and no - * generated predicate) is exercised by {@link #testHierarchicalReadAllItemsPrefixPartitionKey()} - * and {@link #testHierarchicalPrefixHalfOpenRange()}. + * over-span fix. The prefix-EPK header path itself (a raw query with only a prefix key and no + * generated predicate) is exercised by {@link #testHierarchicalPrefixHalfOpenRange()} and + * {@link #testExplicitFeedRangeShardedReadParity()}. */ @Test(groups = {"thinclient"}, timeOut = TIMEOUT) public void testHierarchicalReadManyByPartitionKeysPrefixPartitionKey() { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java index 6f4a3b85e419..ceb703cc5154 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientTestBase.java @@ -2,64 +2,29 @@ // 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. + * Stateless helper holder for thin client E2E tests: shared constants and static endpoint-routing + * assertions used across thin-client test classes. This is intentionally NOT a TestNG base class - + * no test class extends it and it owns no client/container lifecycle; consumers only static-import + * the assertion helpers and reference the shared constants. */ -public abstract class ThinClientTestBase extends TestSuiteBase { +public final class ThinClientTestBase { - 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(); + static final String THIN_CLIENT_ENDPOINT_INDICATOR = ":10250/"; + static final String ID_FIELD = "id"; + static final String PARTITION_KEY_FIELD = "mypk"; + 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(); - 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() { - if (this.client != null) { - this.client.close(); - } - } - - /** - * 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; + private ThinClientTestBase() { } /**