diff --git a/sdk/cosmos/azure-cosmos-tests/docs/THINCLIENT_TEST_MATRIX.md b/sdk/cosmos/azure-cosmos-tests/docs/THINCLIENT_TEST_MATRIX.md new file mode 100644 index 000000000000..3f7b79c2360d --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/docs/THINCLIENT_TEST_MATRIX.md @@ -0,0 +1,203 @@ +# Thin Client (Gateway V2) QueryPlan — E2E Test Specification + +**Primary suite:** `sdk/cosmos/azure-cosmos-tests/.../rx/ThinClientQueryE2ETest.java` (97 `@Test` methods) +**Sibling suites:** `ThinClientPointOperationE2ETest` (3), `ThinClientChangeFeedE2ETest` (3), `ThinClientStoredProcedureE2ETest` (3) +**TestNG group:** `thinclient` · **Status:** reverse-engineered from the committed test code + +> This document is the human-readable test-design spec reconstructed from the test code so it can be reviewed independently of the implementation. It states **what is being proven, how, and what is intentionally not covered** — i.e. exactly what a reviewer needs to sign off on the test strategy. + +--- + +## 1. Objective + +Prove that **routing query execution (and its prerequisite QueryPlan retrieval) through the Gateway V2 "thin client" proxy (`:10250`) produces results that are functionally identical to the established Direct (TCP, `:443`) path**, across the full SQL surface area that the QueryPlan/feature-flag contract advertises. + +The thin client changes two things the tests must defend: +1. **Request routing** — a new RNTBD `QueryPlan` operation (`0x0042`) plus `SupportedQueryFeatures` / `QueryVersion` headers are emitted to the proxy. +2. **Response deserialization** — the proxy returns `queryRanges` as **PartitionKeyInternal JSON arrays** (not EPK hex), which the client must convert to sorted EPK ranges before building the pipeline. + +If either is wrong, results diverge from Direct, or the wrong endpoint is hit, or an error is mis-mapped (e.g. `statusCode 0`). The suite is designed to catch all three classes. + +--- + +## 2. System Under Test vs. Baseline (the oracle) + +| Role | Client | Transport | Endpoint | +|------|--------|-----------|----------| +| **Baseline / oracle** | `directClient` | Direct TCP | `:443` | +| **SUT** | `thinClient` | Gateway V2 (HTTP/2 → proxy) | `:10250` | + +**Oracle model — differential testing.** Each functional test runs the *same* query/options against **both** clients over the **same** physical container and asserts the SUT output matches the baseline. The Direct path is the trusted reference implementation; the thin client must match it exactly. This avoids hand-maintained expected-result fixtures and means every test is automatically a parity test. + +Both clients address the **same** container (seeded once via Direct), so any divergence is attributable to the thin-client path, not to data differences. + +--- + +## 3. Data Model (fixture) + +Seeded once per class (`@BeforeClass`), 10 documents, single logical partition `commonPk` (`/mypk`): + +| Field | Type | Values / shape | +|-------|------|----------------| +| `id` | string | `tcdoc--` | +| `mypk` | string | `commonPk` (shared) | +| `category` | string | electronics / books / clothing / toys (skewed distribution) | +| `status` | string | active / inactive | +| `age` | int | 8–61 (incl. boundary values 17, 18-adjacent, 42, 61) | +| `price` | double | 7.50–549.99 (drives SUM/AVG/MIN/MAX tolerance tests) | +| `idx` | int | 0–9 (stable ORDER BY key) | +| `isActive` | bool | derived from `status` | +| `address` | object | `{ city: Seattle|Portland, zip }` (nested-property tests) | +| `scores` | int[] | `[i*10, i*10+5]` (JOIN / ARRAY_CONTAINS / EXISTS) | +| `tags` | string[] | category + conditional `on-sale` / `featured` (JOIN / EXISTS string match) | + +Design intent: the fixture is deliberately heterogeneous (nested objects, arrays of two types, booleans, skewed categorical distribution, numeric range with decimals) so a **single shared container** exercises filtering, projection, aggregation, JOIN, EXISTS, and type-checking without per-test seeding. Tests needing a different topology (multi-partition, vector, full-text, hybrid) create and tear down their **own** container. + +Cleanup: `@AfterClass` bulk-deletes seeded docs and closes both clients; per-test containers use `safeDeleteContainer` in a `finally`. + +--- + +## 4. Assertion Contracts (the heart of the spec) + +Every comparison helper enforces **endpoint provenance first**, then result equivalence. A reviewer should focus here — these helpers define what "identical" means. + +### 4.1 Endpoint provenance — applied in *every* test +``` +for (CosmosDiagnostics d : tcResult.diagnostics) assertThinClientEndpointUsed(d); +``` +Asserts the thin-client pages actually traversed `:10250`. This guards against a silent fallback to Direct/Gateway V1 that would make a parity test pass for the wrong reason. The invalid-query test asserts the same on the **exception** diagnostics. + +### 4.2 Document-set equality — `assertSameDocumentIds(gw, tc, ordered, desc)` +- **Count** must match first. +- Whether the comparison is order-sensitive is decided by **`isStrictlyOrdered(queryText, options)`**: + - **strict sequence** when the query has an `ORDER BY` **or** runs against a single logical partition (a partition key is set on the options, i.e. `partitionedOptions()`). The large majority of §5 tests fall here. + - **sorted set** only when order is genuinely undefined — a **cross-partition** query (no partition key on the options) **without** an `ORDER BY`. +- **Comparison key:** when every row projects an `id` it is used as the key; otherwise (projections that don't select `id`, e.g. `SELECT AS x`) the **full rows** are compared instead — numeric leaves within `NUMERIC_TOLERANCE` — so coverage is not silently dropped. The unordered full-row path is a multiset match (each baseline row consumes one distinct thin-client row). + +This split is the key correctness nuance: it asserts ordering exactly where it is defined (ORDER BY, or single-partition) and relaxes to set comparison only where cross-partition / cross-page ordering is genuinely undefined — catching real ordering bugs without false failures. + +### 4.3 Scalar aggregates — `assertScalarValueEquals` (COUNT/SUM/AVG/MIN/MAX) +Numeric values compared within `NUMERIC_TOLERANCE = 1e-6`; falls back to string equality for non-numerics. Prevents false mismatches from float formatting differences between the two serialization paths while still catching real aggregate errors. + +### 4.4 GROUP BY — `assertGroupByDirectAndThinClientMatch(query, groupField)` +Rows compared as a **set keyed by `groupField`**, with full-row deep-equality via `jsonEqualsWithTolerance` (recursive, numeric-tolerant). GROUP BY output order is not guaranteed, so set semantics are correct; deep tolerant compare still validates every aggregate in the row. + +### 4.5 Drain model — `drainQuery` / `byPage().toIterable()` +All helpers **fully drain** the feed (iterate every page, accumulate results + per-page diagnostics). This means continuation-token handling is exercised implicitly on *every* test, not just the dedicated draining test. + +--- + +## 5. Test Matrix — `ThinClientQueryE2ETest` (97) + +Unless noted, each test calls `assertDirectAndThinClientMatch(...)` (§4.1–4.2) with `partitionedOptions()` (single-PK), so per `isStrictlyOrdered` (§4.2) results are compared in **strict sequence** unless the row says otherwise. + +| # | Category | Tests | Count | Notable assertion | +|---|----------|-------|-------|-------------------| +| 1 | **Filtering (WHERE)** | SelectAll, WhereEquality, WhereEqualityParameterized, WhereRangeGreaterThan, WhereRangeLessThanOrEqual, WhereRangeBetween, WhereIn, WhereCompoundAndOr, WhereNotEqual, WhereBooleanField, WhereIsDefined, WhereStartsWith, WhereContains, WhereArrayContains, WhereNestedProperty, Between | 16 | strict-seq (single-PK) | +| 2 | **Projection** | SelectSpecificFields, SelectComputedAlias, SelectValueObject, SelectValueScalar | 4 | field/scalar parity; strict-seq | +| 3 | **ORDER BY** | OrderByAsc, OrderByDesc, MultipleOrderBy | 3 | **sequence**; MultipleOrderBy uses a (category ASC, age DESC) composite-indexed container | +| 4 | **DISTINCT** | DistinctValue, DistinctValueBoolean | 2 | scalar set | +| 5 | **TOP** | Top, TopWithOrderBy | 2 | strict-seq | +| 6 | **Aggregates** | Count, Sum, Avg, Min, Max, DCount | 6 | scalar ±1e-6 | +| 7 | **GROUP BY** | GroupByCount, GroupBySumAvg | 2 | keyed-set + tolerant deep-equal | +| 8 | **OFFSET/LIMIT** | OffsetLimit | 1 | **sequence** (has ORDER BY) | +| 9 | **JOIN (self, arrays)** | JoinScoresArray, JoinWithFilter, JoinTagsArray | 3 | strict-seq | +| 10 | **EXISTS subquery** | ExistsSubquery, ExistsSubqueryWithStringMatch, ExistsAliasInProjection | 3 | strict-seq | +| 11 | **LIKE (basic)** | LikePrefix, LikeSuffix, LikeContains | 3 | strict-seq | +| 12 | **String functions** | Concat, EndsWith, Lower, Upper, Length, Substring, Replace, IndexOf, Left, Reverse, Trim, RegexMatch | 12 | strict-seq | +| 13 | **Type-check functions** | IsArray, IsBool, IsNull, IsNumber, IsString, IsObject | 6 | strict-seq | +| 14 | **Math functions** | MathAbs, MathCeilingFloor, MathRound, MathPower, MathSqrt | 5 | strict-seq | +| 15 | **Array functions** | ArrayLength, ArraySlice | 2 | strict-seq | +| 16 | **Conditional** | Iif | 1 | strict-seq | +| 17 | **Date/Time** | GetCurrentDateTime | 1 | ISO-8601 shape only (values differ by design) | +| 18 | **QueryOracle — LIKE patterns** | LikeSingleCharWildcard, LikeCharacterClassRange, LikeNegatedCharacterClass, NotLike | 4 | strict-seq | +| 19 | **QueryOracle — scalar expressions** | CoalesceOperator, ComputedMemberIndexer, ArrayLiteralProjection, UnaryNegation, ModuloOperator, TernaryConditional | 6 | scalar / strict-seq | +| 20 | **Query-plan caching** | CachedQueryPlanFromProxyExecutesCorrectly | 1 | see §6.5 | +| 21 | **Cross-partition** | CrossPartitionSelectAll, CrossPartitionWhereFilter | 2 | no PK filter; ORDER BY idx → sequence; asserts fan-out >1 PKRange | +| 22 | **Multi-range (own container, 24k RU)** | MultiRangeIN(3/5), MultiRangeOR(2/3), MultiRangeMany(10/10) | 3 | exact count + sorted-ID equality; Many asserts >1 PKRange | +| 23 | **Continuation draining** | ContinuationTokenDraining | 1 | see §6.1 | +| 24 | **Error handling** | InvalidQueryReturnsBadRequest | 1 | see §6.2 | +| 25 | **Vector / Full-Text / Hybrid** | VectorSearch, FullTextSearch, FullTextScoreRanking, HybridSearch | 4 | see §6.3 | +| 26 | **readManyByPartitionKeys (validation QueryPlan path)** | NoCustomQuery, WithCustomQuery, WithParameterizedCustomQuery | 3 | see §6.4 | + +**Total: 97.** + +--- + +## 6. Special-case tests (hardened — F1–F6) + +These were strengthened beyond simple parity because they guard the specific failure modes the thin-client change introduces. A reviewer should weigh these most heavily. + +### 6.1 `testContinuationTokenDraining` (F-drain) +Drains the thin client with **page size 3** to force multiple continuations, against a fully-drained Direct baseline. +Asserts: (a) `pageCount > 1` (continuations genuinely occurred — not a single page that vacuously passes), (b) drained count == baseline, (c) sorted-ID set == baseline, (d) **no duplicate IDs** across pages (`HashSet.size() == list.size()`). Directly covers the "lower page size and drain" scenario and proves continuation tokens round-trip correctly through the proxy without dropping or duplicating rows. + +### 6.2 `testInvalidQueryReturnsBadRequest` (F-error) +Sends `SELEC * FORM c`. Asserts **`statusCode == 400`** (explicitly, with a message naming the `statusCode 0` thin-client decode regression as the thing being guarded), and asserts the failing request used the thin-client endpoint. This is the regression test for the production fix in `RxGatewayStoreModel.validateOrThrow` (non-JSON / NUL-padded proxy error frame must surface as a real 400, never 0). + +### 6.3 Vector / Full-Text / Hybrid (F-search) +- **Vector**: TOP-5 `VectorDistance` ORDER BY, **cosine** and **euclidean** variants; asserts size==5, **position-by-position ID equality** with Direct, top result is the planted nearest neighbor, and score > 0.99 for the cosine match. +- **FullTextContains**: asserts non-empty baseline and sorted-ID parity. +- **FullTextScore ranking** (`ORDER BY RANK FullTextScore`): asserts **exact ranked order** parity. +- **Hybrid** (`ORDER BY RANK RRF(VectorDistance, FullTextScore)`): asserts **exact ranked order** parity for TOP-3. + +These prove the streaming/non-streaming ORDER BY and HybridSearch features advertised in `SupportedQueryFeatures` actually produce correct *ordering* through the proxy — the highest-risk area, since hybrid/non-streaming plans are where the advertised feature set and the deserialized `queryRanges` interact. + +### 6.4 `readManyByPartitionKeys` (F-validation) +Exercises the **validation QueryPlan path** that bifurcates between Compute Gateway (V1) and thin client based on whether `DocumentCollection` metadata is available: +- **NoCustomQuery**: no QueryPlan fetched, but per-batch reads must still hit `:10250`. +- **WithCustomQuery**: custom projection+filter triggers QueryPlan validation; asserts parity **and** projection shape (`id`,`category`,`status` present, `status=='active'`). +- **WithParameterizedCustomQuery**: `@cat` binding; asserts parity and every row `category=='electronics'`. +This is the test that proves the routing rule `useGatewayMode = (partitionKeyDefinition == null)` behaves correctly for both branches. + +### 6.5 `testCachedQueryPlanFromProxyExecutesCorrectly` (F-cache) +Runs a single-partition GROUP BY (which requires a query plan) **twice**. Because a partition key is set on the options, the proxy-generated query plan is cached on the client after the first execution and reused on the second. Both executions assert GROUP BY parity with Direct — proving a cached, proxy-generated `QueryPlan` continues to execute correctly (does not go stale or get mis-applied) on reuse. + +--- + +## 7. Query-feature coverage vs. the advertised contract + +`SupportedQueryFeatures` advertised by the client (per `QueryPlanRetriever`): Aggregate, CompositeAggregate, MultipleOrderBy, MultipleAggregates, OrderBy, OffsetAndLimit, Distinct, GroupBy, Top, DCount, NonValueAggregate, NonStreamingOrderBy, HybridSearch, WeightedRankFusion. + +| Advertised feature | Covered by | Gap? | +|--------------------|-----------|------| +| Aggregate / NonValueAggregate | §6 Aggregates (Count/Sum/Avg/Min/Max) | — | +| GroupBy | GroupByCount, GroupBySumAvg | — | +| OrderBy / MultipleOrderBy / NonStreamingOrderBy | OrderBy*, MultipleOrderBy (composite index), vector/FTS ranking | — | +| OffsetAndLimit | OffsetLimit | — | +| Distinct | DistinctValue* | — | +| Top | Top, TopWithOrderBy | — | +| HybridSearch / WeightedRankFusion | HybridSearch (RRF) | — | +| DCount | DCount | — | +| CompositeAggregate / MultipleAggregates | GroupBySumAvg (two aggs) | partial | + +**Intentionally NOT advertised** (documented in code, should be noted to the reviewer): +- `CountIf` — Java has no CountIf aggregator in `SingleGroupAggregator` yet. +- `ListAndSetAggregate` — Java has no MAKELIST/MAKESET support. +- `HybridSearchSkipOrderByRewrite` — Java hybrid pipeline cannot yet consume the optimized plan (would 400 / SC1001). + +--- + +## 8. Known gaps / limitations (for reviewer sign-off) + +1. **Cross-partition aggregate/GROUP BY** — aggregates run within a single logical partition (`partitionedOptions`); cross-partition aggregate merge through the proxy is not directly asserted. +2. **Large result drain** — continuation draining uses 10 docs / page size 3 (multiple pages, but small). No high-cardinality (hundreds/thousands of docs) drain to stress merge/continuation under realistic volume. +3. **Geospatial (`ST_DISTANCE`/`ST_WITHIN`) and UDF query categories** are intentionally not covered — they require a geospatial-indexed container and registered user-defined functions respectively, which the shared seeded fixture does not provide (noted in the QueryOracle-derived section of the test). +4. **Tolerance-based scalar/JSON compare** could in principle mask a genuine sub-1e-6 aggregate discrepancy; acceptable but worth noting. +5. **Mixed/empty-result edge cases** (query matching zero docs, null-projection rows) are implicit at best. +6. Sibling suites (point-op, change feed, sproc) are summarized in the header but specced separately. + +--- + +## 9. What a reviewer (e.g. Aditya Sarpotdar) should verify + +- [ ] The **oracle is sound**: Direct path is an acceptable reference for every advertised feature (esp. NonStreamingOrderBy / HybridSearch). +- [ ] The **ordered-vs-unordered** split (§4.2) matches the service ordering contract for each query class. +- [ ] **Endpoint provenance** (§4.1) is genuinely sufficient to prove no Direct fallback (i.e. `assertThinClientEndpointUsed` inspects the actual transport, not a cached/first-hop signal). +- [ ] The **error contract** test (§6.2) asserts the right status/substatus for the proxy's actual error frame. +- [ ] Agreement on the **gap list** (§8) — especially whether cross-partition aggregate merge and a high-volume drain must be added before merge. +- [ ] Multi-range test (row 22) actually spans **multiple physical partitions** (24k RU container) so EPK-range sort correctness through `convertToSortedEpkRanges` is exercised. + +--- + +*Reverse-engineered from the committed test code. No behavioral claims here are asserted beyond what the test assertions enforce.* diff --git a/sdk/cosmos/azure-cosmos-tests/pom.xml b/sdk/cosmos/azure-cosmos-tests/pom.xml index 4ee203075178..ca7c5145062b 100644 --- a/sdk/cosmos/azure-cosmos-tests/pom.xml +++ b/sdk/cosmos/azure-cosmos-tests/pom.xml @@ -930,6 +930,8 @@ Licensed under the MIT License. true + true + true 1 256 paranoid diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanRoutingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanRoutingTest.java new file mode 100644 index 000000000000..4bee2b714176 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanRoutingTest.java @@ -0,0 +1,162 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.query.DocumentQueryExecutionContextFactory; +import com.azure.cosmos.implementation.query.IDocumentQueryClient; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.SqlQuerySpec; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; + +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests that pin the bifurcation of the validation-only QueryPlan request issued by + * {@code readManyByPartitionKeys} when a custom query is supplied. + * + *

The fork lives in {@code QueryPlanRetriever.getQueryPlanThroughGatewayAsync}: + * 

queryPlanRequest.useGatewayMode = !(queryClient.useThinClient(queryPlanRequest)
+ *     && Configs.isThinClientQueryPlanEnabled());
+ * + *

Routing is driven by the thin-client opt-in ({@link IDocumentQueryClient#useThinClient}), + * not by collection metadata. These tests assert both halves of that contract: + *

    + *
  • thin-client not eligible ⇒ request pinned to Gateway V1 (Compute Gateway);
    • + *
  • thin-client eligible ⇒ request is thin-client eligible (Gateway V2 / proxy) — the + * {@code useGatewayMode=false} flag is the prerequisite the store-model layer reads + * to bifurcate.
    • + *
+ */ +public class ReadManyByPartitionKeyQueryPlanRoutingTest { + + @Test(groups = { "unit" }) + public void validationQueryPlanRoutesToGatewayV1WhenThinClientNotEligible() { + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(RxDocumentServiceRequest.class); + IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ false); + + DocumentQueryExecutionContextFactory + .fetchQueryPlanForValidation( + Mockito.mock(DiagnosticsClientContext.class), + queryClient, + new SqlQuerySpec("SELECT * FROM c"), + "dbs/db/colls/col", + new CosmosQueryRequestOptions(), + collectionWithPartitionKey(), + /* queryPlanCachingEnabled */ false, + Collections.emptyMap()) + .block(); + + assertThat(requestCaptor.getAllValues()) + .as("a single validation query-plan request must be issued") + .hasSize(1); + assertThat(requestCaptor.getValue().useGatewayMode) + .as("validation query-plan must pin to Gateway V1 when thin client is not eligible") + .isTrue(); + } + + @Test(groups = { "unit" }) + public void validationQueryPlanIsThinClientEligibleWhenOptedIn() { + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(RxDocumentServiceRequest.class); + IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ true); + + DocumentQueryExecutionContextFactory + .fetchQueryPlanForValidation( + Mockito.mock(DiagnosticsClientContext.class), + queryClient, + new SqlQuerySpec("SELECT * FROM c"), + "dbs/db/colls/col", + new CosmosQueryRequestOptions(), + collectionWithPartitionKey(), + /* queryPlanCachingEnabled */ false, + Collections.emptyMap()) + .block(); + + assertThat(requestCaptor.getAllValues()) + .as("a single validation query-plan request must be issued") + .hasSize(1); + assertThat(requestCaptor.getValue().useGatewayMode) + .as("validation query-plan must be thin-client eligible (useGatewayMode=false) " + + "when the thin-client opt-in is satisfied") + .isFalse(); + } + + @Test(groups = { "unit" }) + public void validationQueryPlanForcedToGatewayV1WhenKillSwitchDisabled() { + // The kill-switch (Configs.isThinClientQueryPlanEnabled(), backed by the + // COSMOS.THINCLIENT_QUERY_PLAN_ENABLED system property / COSMOS_THINCLIENT_QUERY_PLAN_ENABLED + // env var) must dominate the thin-client opt-in: even when the request is thin-client eligible + // (useThinClient=true), disabling the switch pins the validation query-plan back onto Gateway V1 + // (useGatewayMode=true). This is the inverse of validationQueryPlanIsThinClientEligibleWhenOptedIn. + String previous = System.getProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED"); + System.setProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED", "false"); + try { + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(RxDocumentServiceRequest.class); + IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ true); + + DocumentQueryExecutionContextFactory + .fetchQueryPlanForValidation( + Mockito.mock(DiagnosticsClientContext.class), + queryClient, + new SqlQuerySpec("SELECT * FROM c"), + "dbs/db/colls/col", + new CosmosQueryRequestOptions(), + collectionWithPartitionKey(), + /* queryPlanCachingEnabled */ false, + Collections.emptyMap()) + .block(); + + assertThat(requestCaptor.getAllValues()) + .as("a single validation query-plan request must be issued") + .hasSize(1); + assertThat(requestCaptor.getValue().useGatewayMode) + .as("validation query-plan must be forced onto Gateway V1 (useGatewayMode=true) when " + + "the thin-client query-plan kill-switch is disabled, even though the request is " + + "thin-client eligible") + .isTrue(); + } finally { + if (previous == null) { + System.clearProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED"); + } else { + System.setProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED", previous); + } + } + } + + private static DocumentCollection collectionWithPartitionKey() { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition(); + partitionKeyDefinition.setPaths(Collections.singletonList("/pk")); + DocumentCollection collection = new DocumentCollection(); + collection.setPartitionKey(partitionKeyDefinition); + return collection; + } + + private static IDocumentQueryClient mockQueryClient( + ArgumentCaptor requestCaptor, + boolean useThinClient) { + IDocumentQueryClient queryClient = Mockito.mock(IDocumentQueryClient.class); + Mockito.when(queryClient.useThinClient(Mockito.any())).thenReturn(useThinClient); + // executeFeedOperationWithAvailabilityStrategy is a generic method; use doAnswer so the + // returned Mono.empty() is supplied at invocation time without forcing a generic cast on the + // stubbing call site. We only care about the captured request, not the downstream payload. + Mockito + .doAnswer(invocation -> Mono.empty()) + .when(queryClient) + .executeFeedOperationWithAvailabilityStrategy( + Mockito.any(), + Mockito.any(), + Mockito.any(), + requestCaptor.capture(), + Mockito.any(), + Mockito.anyString()); + return queryClient; + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java index e72643f54d07..9f75b4d3cfbc 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java @@ -15,7 +15,9 @@ import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.buffer.ByteBuf; import io.netty.buffer.PooledByteBufAllocator; +import io.netty.buffer.Unpooled; import io.netty.channel.ConnectTimeoutException; +import io.netty.handler.codec.http.HttpResponseStatus; import io.netty.handler.timeout.ReadTimeoutException; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; @@ -28,6 +30,7 @@ import java.lang.reflect.Method; import java.net.SocketException; import java.net.URI; +import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.HashMap; import java.util.Map; @@ -1041,6 +1044,94 @@ private static int getGatewayRetryWithTimeoutInSeconds(RxGatewayStoreModel store return (int) getGatewayRetryWithTimeoutInSeconds.invoke(storeModel); } + /** + * Error bodies that {@code validateOrThrow} must turn into a CosmosException carrying the + * real HTTP status (400 here) rather than leaking the parse failure as statusCode 0. + * + * Columns: description, rawBody, expectMessageContains (null = don't assert message content). + * The Gateway V2 / thin-client proxy can return a non-JSON body (optionally NUL-padded) or a + * valid-but-non-object JSON value (scalar / array) for query-plan failures; both must fall back + * cleanly. A valid CosmosError JSON body must still parse and preserve its message. + */ + @DataProvider(name = "validateOrThrowErrorBodyProvider") + public Object[][] validateOrThrowErrorBodyProvider() { + return new Object[][]{ + {"NUL-padded plain text", "Query plan generation failed\u0000\u0000\u0000\u0000", "Query plan generation failed"}, + {"JSON scalar string", "\"BadRequest\"", null}, + {"JSON array", "[1,2,3]", null}, + {"JSON number", "42", null}, + {"valid CosmosError JSON", "{\"code\":\"BadRequest\",\"message\":\"bad query syntax\"}", "bad query syntax"}, + {"empty body", "", null}, + }; + } + + /** + * Locks the contract of the {@code validateOrThrow} fallback: regardless of the shape of the + * error body the proxy returns, the thrown {@link CosmosException} must carry the real HTTP + * status code (400) and never the leaked statusCode 0. The non-object-JSON rows (scalar/array/ + * number) specifically guard the ClassCastException path that previously escaped the catch. + */ + @Test(groups = "unit", dataProvider = "validateOrThrowErrorBodyProvider") + public void validateOrThrowPreservesStatusCodeForMalformedBodies( + String description, + String rawBody, + String expectMessageContains) throws Exception { + + RxGatewayStoreModel storeModel = new RxGatewayStoreModel( + mockDiagnosticsClientContext(), + Mockito.mock(ISessionContainer.class), + ConsistencyLevel.SESSION, + QueryCompatibilityMode.Default, + new UserAgentContainer(), + Mockito.mock(GlobalEndpointManager.class), + Mockito.mock(HttpClient.class), + null, + null); + + RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName( + mockDiagnosticsClientContext(), + OperationType.QueryPlan, + "/dbs/db/colls/col", + ResourceType.Document); + dsr.requestContext = new DocumentServiceRequestContext(); + + ByteBuf bodyBuf = Unpooled.wrappedBuffer(rawBody.getBytes(StandardCharsets.UTF_8)); + + try { + invokeValidateOrThrow(storeModel, dsr, HttpResponseStatus.BAD_REQUEST, new HttpHeaders(), bodyBuf); + fail("validateOrThrow should have thrown a CosmosException for: " + description); + } catch (java.lang.reflect.InvocationTargetException invocationTargetException) { + Throwable cause = invocationTargetException.getCause(); + assertThat(cause) + .as("validateOrThrow should throw CosmosException for: " + description) + .isInstanceOf(CosmosException.class); + CosmosException cosmosException = (CosmosException) cause; + assertThat(cosmosException.getStatusCode()) + .as("status code must be preserved (never statusCode 0) for: " + description) + .isEqualTo(HttpConstants.StatusCodes.BADREQUEST); + if (expectMessageContains != null) { + assertThat(cosmosException.getMessage()) + .as("error message should be preserved for: " + description) + .contains(expectMessageContains); + } + } + } + + private static void invokeValidateOrThrow(RxGatewayStoreModel storeModel, + RxDocumentServiceRequest request, + HttpResponseStatus status, + HttpHeaders headers, + ByteBuf retainedBodyAsByteBuf) throws Exception { + Method validateOrThrow = RxGatewayStoreModel.class.getDeclaredMethod( + "validateOrThrow", + RxDocumentServiceRequest.class, + HttpResponseStatus.class, + HttpHeaders.class, + ByteBuf.class); + validateOrThrow.setAccessible(true); + validateOrThrow.invoke(storeModel, request, status, headers, retainedBodyAsByteBuf); + } + enum SessionTokenType { NONE, // no session token applied USER, // userControlled session token diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java deleted file mode 100644 index 3965c3ea8e07..000000000000 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java +++ /dev/null @@ -1,379 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. -package com.azure.cosmos.implementation; - -import com.azure.cosmos.ConsistencyLevel; -import com.azure.cosmos.CosmosAsyncClient; -import com.azure.cosmos.CosmosAsyncContainer; -import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.CosmosDiagnostics; -import com.azure.cosmos.CosmosDiagnosticsContext; -import com.azure.cosmos.CosmosDiagnosticsRequestInfo; -import com.azure.cosmos.FlakyTestRetryAnalyzer; -import com.azure.cosmos.models.CosmosBatch; -import com.azure.cosmos.models.CosmosBatchResponse; -import com.azure.cosmos.models.CosmosBulkItemResponse; -import com.azure.cosmos.models.CosmosBulkOperationResponse; -import com.azure.cosmos.models.CosmosBulkOperations; -import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; -import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.FeedRange; -import com.azure.cosmos.models.PartitionKey; -import com.azure.cosmos.models.SqlQuerySpec; -import com.azure.cosmos.models.SqlParameter; -import com.azure.cosmos.models.FeedResponse; -import com.azure.cosmos.models.CosmosContainerProperties; -import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.CosmosItemResponse; -import com.azure.cosmos.models.CosmosItemRequestOptions; -import com.azure.cosmos.models.CosmosPatchOperations; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ObjectNode; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.testng.annotations.Test; -import reactor.core.publisher.Flux; - -import java.util.Arrays; -import java.util.Collection; -import java.util.List; -import java.util.UUID; - -import static org.assertj.core.api.AssertionsForClassTypes.assertThat; -import static org.assertj.core.api.Fail.fail; -import static com.azure.cosmos.rx.TestSuiteBase.createCollection; - -// End to end sanity tests for basic thin client functionality. -public class ThinClientE2ETest { - private static final Logger logger = LoggerFactory.getLogger(ThinClientE2ETest.class); - private static final String thinClientEndpointIndicator = ":10250/"; - - @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class) - public void testThinClientQuery() { - CosmosAsyncClient client = null; - try { - // If running locally, uncomment these lines - // System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); - // System.setProperty("COSMOS.HTTP2_ENABLED", "true"); - - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .gatewayMode() - .consistencyLevel(ConsistencyLevel.SESSION) - .buildAsyncClient(); - - CosmosAsyncContainer container = client.getDatabase("db1").getContainer("c2"); - String idName = "id"; - String partitionKeyName = "partitionKey"; - ObjectMapper mapper = new ObjectMapper(); - ObjectNode doc = mapper.createObjectNode(); - String idValue = UUID.randomUUID().toString(); - doc.put(idName, idValue); - doc.put(partitionKeyName, idValue); - - container.createItem(doc, new PartitionKey(idValue), null).block(); - - String query = "select * from c WHERE c." + partitionKeyName + "=@id"; - SqlQuerySpec querySpec = new SqlQuerySpec(query); - querySpec.setParameters(Arrays.asList(new SqlParameter("@id", idValue))); - CosmosQueryRequestOptions requestOptions = - new CosmosQueryRequestOptions().setPartitionKey(new PartitionKey(idValue)); - FeedResponse response = container - .queryItems(querySpec, requestOptions, ObjectNode.class) - .byPage() - .blockFirst(); - - ObjectNode docFromResponse = response.getResults().get(0); - assertThat(docFromResponse.get(partitionKeyName).textValue()).isEqualTo(idValue); - assertThat(docFromResponse.get(idName).textValue()).isEqualTo(idValue); - assertThinClientEndpointUsed(response.getCosmosDiagnostics()); - - } finally { - if (client != null) { - client.close(); - } - } - } - - @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class) - public void testThinClientBulk() { - CosmosAsyncClient client = null; - try { - // If running locally, uncomment these lines - // System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); - // System.setProperty("COSMOS.HTTP2_ENABLED", "true"); - - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .gatewayMode() - .consistencyLevel(ConsistencyLevel.EVENTUAL) - .buildAsyncClient(); - - CosmosAsyncContainer container = client.getDatabase("db1").getContainer("c2"); - String idName = "id"; - String partitionKeyName = "partitionKey"; - ObjectMapper mapper = new ObjectMapper(); - ObjectNode doc = mapper.createObjectNode(); - String idValue = UUID.randomUUID().toString(); - doc.put(idName, idValue); - doc.put(partitionKeyName, idValue); - - Flux> responsesFlux = container.executeBulkOperations(Flux.just( - CosmosBulkOperations.getCreateItemOperation(doc, new PartitionKey(idValue)) - )); - - List> responses = responsesFlux.collectList().block(); - - assertThat(responses.size()).isEqualTo(1); - assertThat(responses.get(0).getException()).isNull(); - CosmosBulkItemResponse bulkResponse = responses.get(0).getResponse(); - assertThat(bulkResponse.isSuccessStatusCode()).isEqualTo(true); - assertThinClientEndpointUsed(bulkResponse.getCosmosDiagnostics()); - } finally { - if (client != null) { - client.close(); - } - } - } - - @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class) - public void testThinClientBatch() { - CosmosAsyncClient client = null; - try { - // If running locally, uncomment these lines - // System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); - // System.setProperty("COSMOS.HTTP2_ENABLED", "true"); - - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .gatewayMode() - .consistencyLevel(ConsistencyLevel.SESSION) - .buildAsyncClient(); - - CosmosAsyncContainer container = client.getDatabase("db1").getContainer("c2"); - String idName = "id"; - String partitionKeyName = "partitionKey"; - ObjectMapper mapper = new ObjectMapper(); - String pkValue = UUID.randomUUID().toString(); - ObjectNode doc1 = mapper.createObjectNode(); - String idValue1 = UUID.randomUUID().toString(); - doc1.put(idName, idValue1); - doc1.put(partitionKeyName, pkValue); - - ObjectNode doc2 = mapper.createObjectNode(); - String idValue2 = UUID.randomUUID().toString(); - doc2.put(idName, idValue2); - doc2.put(partitionKeyName, pkValue); - - CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(pkValue)); - batch.createItemOperation(doc1); - batch.createItemOperation(doc2); - - CosmosBatchResponse response = container - .executeCosmosBatch(batch) - .block(); - - assertThat(response.getStatusCode()).isEqualTo(200); - assertThinClientEndpointUsed(response.getDiagnostics()); - } finally { - if (client != null) { - client.close(); - } - } - } - - @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class) - public void testThinClientIncrementalChangeFeed() { - CosmosAsyncClient client = null; - try { - // If running locally, uncomment these lines -// System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); -// System.setProperty("COSMOS.HTTP2_ENABLED", "true"); - - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .gatewayMode() - .consistencyLevel(ConsistencyLevel.SESSION) - .buildAsyncClient(); - - CosmosAsyncContainer container = client.getDatabase("db1").getContainer("c2"); - String idName = "id"; - String partitionKeyName = "partitionKey"; - ObjectMapper mapper = new ObjectMapper(); - String pkValue = UUID.randomUUID().toString(); - ObjectNode doc1 = mapper.createObjectNode(); - String idValue1 = UUID.randomUUID().toString(); - doc1.put(idName, idValue1); - doc1.put(partitionKeyName, pkValue); - - ObjectNode doc2 = mapper.createObjectNode(); - String idValue2 = UUID.randomUUID().toString(); - doc2.put(idName, idValue2); - doc2.put(partitionKeyName, pkValue); - - CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(pkValue)); - batch.createItemOperation(doc1); - batch.createItemOperation(doc2); - - CosmosBatchResponse response = container - .executeCosmosBatch(batch) - .block(); - - FeedResponse changeFeedResponse = container - .queryChangeFeed(CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(FeedRange.forFullRange()), ObjectNode.class) - .byPage() - .blockFirst(); - - assertThat(changeFeedResponse).isNotNull(); - assertThat(changeFeedResponse.getResults()).isNotNull(); - assertThat(changeFeedResponse.getResults().size()).isGreaterThanOrEqualTo(1); - assertThinClientEndpointUsed(changeFeedResponse.getCosmosDiagnostics()); - } finally { - if (client != null) { - client.close(); - } - } - } - - private 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(); - - for (CosmosDiagnosticsRequestInfo requestInfo : requests) { - logger.info( - "Endpoint: {}, RequestType: {}, Partition: {}/{}, ActivityId: {}", - requestInfo.getEndpoint(), - requestInfo.getRequestType(), - requestInfo.getPartitionId(), - requestInfo.getPartitionKeyRangeId(), - requestInfo.getActivityId()); - if (requestInfo.getEndpoint().contains(thinClientEndpointIndicator)) { - return; - } - } - - fail("No request targeting thin client proxy endpoint."); - } - - - @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class) - public void testThinClientDocumentPointOperations() { - CosmosAsyncClient client = null; - try { - // if running locally, uncomment these lines - // System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); - // System.setProperty("COSMOS.HTTP2_ENABLED", "true"); - - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .gatewayMode() - .consistencyLevel(ConsistencyLevel.SESSION) - .buildAsyncClient(); - - String idName = "id"; - String partitionKeyName = "partitionKey"; - - client.createDatabaseIfNotExists("db1").block(); - - CosmosContainerProperties containerDef = - new CosmosContainerProperties("c2", "/" + partitionKeyName); - CosmosAsyncContainer container = createCollection( - client.getDatabase("db1"), - containerDef, - new CosmosContainerRequestOptions(), - 35_000); - - ObjectMapper mapper = new ObjectMapper(); - ObjectNode doc = mapper.createObjectNode(); - String idValue = UUID.randomUUID().toString(); - doc.put(idName, idValue); - doc.put(partitionKeyName, idValue); - - // create - CosmosItemResponse createResponse = container.createItem(doc).block(); - assertThat(createResponse.getStatusCode()).isEqualTo(201); - assertThat(createResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(createResponse.getDiagnostics()); - - // read - CosmosItemResponse readResponse = container.readItem(idValue, new PartitionKey(idValue), ObjectNode.class).block(); - assertThat(readResponse.getStatusCode()).isEqualTo(200); - assertThat(readResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(readResponse.getDiagnostics()); - - ObjectNode doc2 = mapper.createObjectNode(); - String idValue2 = UUID.randomUUID().toString(); - doc2.put(idName, idValue2); - doc2.put(partitionKeyName, idValue); - - // replace - CosmosItemResponse replaceResponse = container.replaceItem(doc2, idValue, new PartitionKey(idValue)).block(); - assertThat(replaceResponse.getStatusCode()).isEqualTo(200); - assertThat(replaceResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(replaceResponse.getDiagnostics()); - - CosmosItemResponse readAfterReplaceResponse = container.readItem(idValue2, new PartitionKey(idValue), ObjectNode.class).block(); - assertThat(readAfterReplaceResponse.getStatusCode()).isEqualTo(200); - ObjectNode replacedItemFromRead = readAfterReplaceResponse.getItem(); - assertThat(replacedItemFromRead.get(idName).asText()).isEqualTo(idValue2); - assertThat(replacedItemFromRead.get(partitionKeyName).asText()).isEqualTo(idValue); - assertThinClientEndpointUsed(readAfterReplaceResponse.getDiagnostics()); - - ObjectNode doc3 = mapper.createObjectNode(); - doc3.put(idName, idValue2); - doc3.put(partitionKeyName, idValue); - doc3.put("newField", "newValue"); - - // upsert - CosmosItemResponse upsertResponse = container.upsertItem(doc3, new PartitionKey(idValue), new CosmosItemRequestOptions()).block(); - assertThat(upsertResponse.getStatusCode()).isEqualTo(200); - assertThat(upsertResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(upsertResponse.getDiagnostics()); - - CosmosItemResponse readAfterUpsertResponse = container.readItem(idValue2, new PartitionKey(idValue), ObjectNode.class).block(); - ObjectNode upsertedItemFromRead = readAfterUpsertResponse.getItem(); - assertThat(upsertedItemFromRead.get(idName).asText()).isEqualTo(idValue2); - assertThat(upsertedItemFromRead.get(partitionKeyName).asText()).isEqualTo(idValue); - assertThat(upsertedItemFromRead.get("newField").asText()).isEqualTo("newValue"); - assertThinClientEndpointUsed(readAfterUpsertResponse.getDiagnostics()); - - // patch - CosmosPatchOperations patchOperations = CosmosPatchOperations.create(); - patchOperations.add("/anotherNewField", "anotherNewValue"); - patchOperations.replace("/newField", "patchedNewField"); - CosmosItemResponse patchResponse = container.patchItem(idValue2, new PartitionKey(idValue), patchOperations, ObjectNode.class).block(); - assertThat(patchResponse.getStatusCode()).isEqualTo(200); - assertThat(patchResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(patchResponse.getDiagnostics()); - - CosmosItemResponse readAfterPatchResponse = container.readItem(idValue2, new PartitionKey(idValue), ObjectNode.class).block(); - ObjectNode patchedItemFromRead = readAfterPatchResponse.getItem(); - assertThat(patchedItemFromRead.get(idName).asText()).isEqualTo(idValue2); - assertThat(patchedItemFromRead.get(partitionKeyName).asText()).isEqualTo(idValue); - assertThat(patchedItemFromRead.get("newField").asText()).isEqualTo("patchedNewField"); - assertThat(patchedItemFromRead.get("anotherNewField").asText()).isEqualTo("anotherNewValue"); - assertThinClientEndpointUsed(readAfterPatchResponse.getDiagnostics()); - - // delete - CosmosItemResponse deleteResponse = container.deleteItem(idValue2, new PartitionKey(idValue)).block(); - assertThat(deleteResponse.getStatusCode()).isEqualTo(204); - assertThat(deleteResponse.getRequestCharge()).isGreaterThan(0.0); - assertThinClientEndpointUsed(deleteResponse.getDiagnostics()); - } finally { - if (client != null) { - client.close(); - } - } - } -} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/PartitionKeyInternalTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/PartitionKeyInternalTest.java index 427d8c763a7d..9710003552c1 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/PartitionKeyInternalTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/PartitionKeyInternalTest.java @@ -14,12 +14,18 @@ import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyInternalUtils; +import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.guava25.collect.Lists; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.fasterxml.jackson.databind.node.ArrayNode; + import java.util.ArrayList; +import java.util.List; import java.util.function.BiFunction; import static org.assertj.core.api.AssertionsForClassTypes.assertThat; @@ -27,6 +33,8 @@ public class PartitionKeyInternalTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + /** * Tests serialization of empty partition key. */ @@ -473,4 +481,126 @@ private static void verifyEffectivePartitionKeyEncoding(String buffer, int lengt PartitionKeyInternal pk = PartitionKeyInternalUtils.createPartitionKeyInternal(buffer.substring(0, length)); assertThat(PartitionKeyInternalHelper.getEffectivePartitionKeyString(pk, pkDefinition)).isEqualTo(expectedValue); } + + // ==================== convertToSortedEpkRanges Unit Tests ==================== + + private static PartitionKeyDefinition singleHashPkDef() { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setPaths(ImmutableList.of("/pk")); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setKind(PartitionKind.HASH); + return pkDef; + } + + @Test(groups = "unit") + public void convertToSortedEpkRangesSingleValueRange() { + // Single PK value range: min=["testValue"], max=["testValue"] + // This is what ServiceInterop returns for WHERE pk = 'testValue' + ObjectNode json = MAPPER.createObjectNode(); + ArrayNode ranges = json.putArray("queryRanges"); + ObjectNode range = ranges.addObject(); + range.putArray("min").add("testValue"); + range.putArray("max").add("testValue"); + range.put("isMinInclusive", true); + range.put("isMaxInclusive", true); + + List> result = PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + + assertThat(result.size()).isEqualTo(1); + // EPK for a string value is a non-empty hex hash + assertThat(result.get(0).getMin()).isNotNull(); + assertThat(result.get(0).getMin().length()).isGreaterThan(0); + assertThat(result.get(0).getMin()).isEqualTo(result.get(0).getMax()); + assertThat(result.get(0).isMinInclusive()).isTrue(); + assertThat(result.get(0).isMaxInclusive()).isTrue(); + } + + @Test(groups = "unit") + public void convertToSortedEpkRangesMultipleRangesSorted() { + // Multiple ranges that need sorting after EPK conversion + ObjectNode json = MAPPER.createObjectNode(); + ArrayNode ranges = json.putArray("queryRanges"); + + // Add two ranges with different PK values — EPK hash order may differ from insertion order + ObjectNode range1 = ranges.addObject(); + range1.putArray("min").add("zzzValue"); + range1.putArray("max").add("zzzValue"); + range1.put("isMinInclusive", true); + range1.put("isMaxInclusive", true); + + ObjectNode range2 = ranges.addObject(); + range2.putArray("min").add("aaaValue"); + range2.putArray("max").add("aaaValue"); + range2.put("isMinInclusive", true); + range2.put("isMaxInclusive", true); + + List> result = PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + + assertThat(result.size()).isEqualTo(2); + // Verify sorted by min EPK (ascending) + assertThat(result.get(0).getMin().compareTo(result.get(1).getMin())).isLessThanOrEqualTo(0); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesMissingQueryRangesThrows() { + // Missing queryRanges property entirely + ObjectNode json = MAPPER.createObjectNode(); + json.put("queryInfo", "someValue"); + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesNonArrayQueryRangesThrows() { + // queryRanges is a string instead of array + ObjectNode json = MAPPER.createObjectNode(); + json.put("queryRanges", "notAnArray"); + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesEmptyArrayThrows() { + // queryRanges is an empty array — would otherwise silently yield zero results + ObjectNode json = MAPPER.createObjectNode(); + json.putArray("queryRanges"); + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesNonObjectElementThrows() { + // queryRanges array contains a string instead of object + ObjectNode json = MAPPER.createObjectNode(); + json.putArray("queryRanges").add("notAnObject"); + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesNullMinBoundaryThrows() { + // Range with null min boundary + ObjectNode json = MAPPER.createObjectNode(); + ArrayNode ranges = json.putArray("queryRanges"); + ObjectNode range = ranges.addObject(); + range.putNull("min"); + range.putArray("max").add("value"); + range.put("isMinInclusive", true); + range.put("isMaxInclusive", false); + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } + + @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + public void convertToSortedEpkRangesMissingInclusiveFieldsThrows() { + // Range missing isMinInclusive and isMaxInclusive + ObjectNode json = MAPPER.createObjectNode(); + ArrayNode ranges = json.putArray("queryRanges"); + ObjectNode range = ranges.addObject(); + range.putArray("min").add("value"); + range.putArray("max").add("value"); + // intentionally no isMinInclusive or isMaxInclusive + + PartitionKeyInternalHelper.convertToSortedEpkRanges("queryRanges", json, singleHashPkDef()); + } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/QueryPlanRetrieverSupportedFeaturesTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/QueryPlanRetrieverSupportedFeaturesTest.java new file mode 100644 index 000000000000..1cf0b3174dfd --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/QueryPlanRetrieverSupportedFeaturesTest.java @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation.query; + +import org.testng.annotations.Test; + +import java.lang.reflect.Field; + +import static org.assertj.core.api.Assertions.assertThat; + +public class QueryPlanRetrieverSupportedFeaturesTest { + + @Test(groups = {"unit"}) + public void supportedFeaturesExcludeUnsupportedAggregates() throws Exception { + String supportedFeatures = getQueryPlanRetrieverString("SUPPORTED_QUERY_FEATURES"); + + assertThat(supportedFeatures) + .doesNotContain(QueryFeature.CountIf.name()) + .doesNotContain("HybridSearchSkipOrderByRewrite") + .doesNotContain("ListAndSetAggregate"); + } + + private static String getQueryPlanRetrieverString(String fieldName) throws Exception { + Class queryPlanRetrieverClass = + Class.forName("com.azure.cosmos.implementation.query.QueryPlanRetriever"); + Field field = queryPlanRetrieverClass.getDeclaredField(fieldName); + field.setAccessible(true); + return (String) field.get(null); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/GatewayReadConsistencyStrategySpyWireTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/GatewayReadConsistencyStrategySpyWireTest.java index 99fc00b69f6f..11d9a4f962c6 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/GatewayReadConsistencyStrategySpyWireTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/GatewayReadConsistencyStrategySpyWireTest.java @@ -728,6 +728,13 @@ private HttpRequest findFeedRequest(String mode, List requests) { String uri = request.uri().toString(); if (isGatewayV2(mode)) { if ("POST".equalsIgnoreCase(request.httpMethod().toString()) && uri.contains(":10250")) { + // Skip query-plan precursor requests: thin-client routes QueryPlan via V2 too, and those + // RNTBD frames intentionally carry no ReadConsistencyStrategy / ConsistencyLevel tokens + // because QueryPlan generation doesn't depend on read-consistency. Mirror the V1 filter + // below by inspecting the RNTBD operation type encoded in the framed body. + if (isRntbdQueryPlanFrame(collectHttpBody(request))) { + continue; + } return request; } } else { @@ -797,5 +804,23 @@ private static RntbdRequest decodeRntbdFrame(byte[] rntbdFrame) { return RntbdRequest.decode(sliced); } + /** + * Peeks at an RNTBD frame's operation-type token (offset 6 LE after the 4-byte length prefix + * and 2-byte resource type) to detect QueryPlan precursor frames. QueryPlan == 0x0042 in + * {@link RntbdConstants.RntbdOperationType}. + */ + private static boolean isRntbdQueryPlanFrame(byte[] rntbdFrame) { + if (rntbdFrame == null || rntbdFrame.length < 8) { + return false; + } + ByteBuf buffer = Unpooled.wrappedBuffer(rntbdFrame); + try { + short operationType = buffer.getShortLE(buffer.readerIndex() + Integer.BYTES + Short.BYTES); + return operationType == RntbdConstants.RntbdOperationType.QueryPlan.id(); + } catch (RuntimeException ignored) { + return false; + } + } + // endregion } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/SinglePartitionDocumentQueryTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/SinglePartitionDocumentQueryTest.java index 0aebb62bfcee..30fc83874361 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/SinglePartitionDocumentQueryTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/SinglePartitionDocumentQueryTest.java @@ -123,9 +123,16 @@ public void querySinglePartitionDocuments() throws Exception { // In gateway mode, serverstoremodel is GatewayStoreModel/ThinClientStoreModel so below passes // In direct mode, serverStoreModel is ServerStoreModel. So queryPlan goes through gatewayProxy and the query // goes through the serverStoreModel - Mockito.verify(spyProxy, Mockito.times(1)).processMessage(Mockito.any()); - if (asyncDocumentClient.getConnectionPolicy().getConnectionMode() == ConnectionMode.DIRECT) { - Mockito.verify(spyServerStoreModel, Mockito.times(1)).processMessage(Mockito.any()); + if (asyncDocumentClient.useThinClient()) { + // In thin client mode both the QueryPlan request and the data query are routed through the thin proxy + // (RxDocumentClientImpl.useThinClientStoreModel now includes OperationType.QueryPlan), so the proxy is + // invoked twice: once for the QueryPlan and once for the query. + Mockito.verify(spyProxy, Mockito.times(2)).processMessage(Mockito.any()); + } else { + Mockito.verify(spyProxy, Mockito.times(1)).processMessage(Mockito.any()); + if (asyncDocumentClient.getConnectionPolicy().getConnectionMode() == ConnectionMode.DIRECT) { + Mockito.verify(spyServerStoreModel, Mockito.times(1)).processMessage(Mockito.any()); + } } } 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 0e41cb2c491e..4e903e502f62 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 @@ -1184,6 +1184,7 @@ private static boolean isRetryableCollectionReadinessFailure(Throwable error) { || (statusCode == HttpConstants.StatusCodes.NOTFOUND && (cosmosException.getSubStatusCode() == HttpConstants.SubStatusCodes.UNKNOWN || cosmosException.getSubStatusCode() == 1013 + || cosmosException.getSubStatusCode() == HttpConstants.SubStatusCodes.OWNER_RESOURCE_NOT_EXISTS || cosmosException.getSubStatusCode() == HttpConstants.SubStatusCodes.INCORRECT_CONTAINER_RID_SUB_STATUS)); } @@ -2710,22 +2711,36 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnosticsContext ctx) assertThat(requests).isNotNull(); assertThat(requests.size()).isPositive(); - // Validate every request rather than early-returning on the first thin-client match: a mixed - // scenario (some data requests via the thin-client endpoint, some via the classic gateway) must - // fail. Every non-QueryPlan (data) request must route through the thin-client proxy endpoint; - // QueryPlan calls are resolved via the classic gateway in thin-client mode, so they are the only - // requests allowed to target a non-thin-client endpoint. + // When thin client is opted in with HTTP/2, QueryPlan calls are routed to Gateway V2 as well + // (Configs.isThinClientQueryPlanEnabled() defaults to true), so every request -- including + // QueryPlan -- must target the thin-client proxy endpoint. + // + // The kill-switch (COSMOS.THINCLIENT_QUERY_PLAN_ENABLED / COSMOS_THINCLIENT_QUERY_PLAN_ENABLED) + // flips only the QueryPlan routing: when disabled, QueryPlan requests are forced back onto + // Gateway V1 while all data-plane requests continue through the thin-client proxy. In that mixed + // state QueryPlan requests are exempt from the :10250 requirement (and are asserted to NOT use it), + // whereas every non-QueryPlan request must still target the thin-client endpoint. + // + // Validate every request rather than early-returning on the first thin-client match: an + // unexpected mixed scenario (some data requests via the classic gateway) must still fail. + boolean queryPlanRoutedToThinClient = Configs.isThinClientQueryPlanEnabled(); for (CosmosDiagnosticsRequestInfo requestInfo : requests) { - // requestType has the form ":" (OperationType.QueryPlan - // stringifies to "QueryPlan"). String requestType = requestInfo.getRequestType(); - if (requestType != null && requestType.endsWith(":QueryPlan")) { + String endpoint = requestInfo.getEndpoint(); + boolean isQueryPlan = requestType != null && requestType.endsWith(":QueryPlan"); + boolean usesThinClient = endpoint != null && endpoint.contains(THIN_CLIENT_ENDPOINT_INDICATOR); + + if (isQueryPlan && !queryPlanRoutedToThinClient) { + assertThat(usesThinClient) + .as("QueryPlan request must be forced onto Gateway V1 (NOT the thin-client proxy " + + THIN_CLIENT_ENDPOINT_INDICATOR + ") when the thin-client query-plan kill-switch " + + "is disabled, but endpoint was: " + endpoint) + .isFalse(); continue; } - String endpoint = requestInfo.getEndpoint(); - assertThat(endpoint != null && endpoint.contains(THIN_CLIENT_ENDPOINT_INDICATOR)) - .as("Non-QueryPlan request must target the thin client proxy endpoint (" + assertThat(usesThinClient) + .as("Request must target the thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR + "), but was: " + endpoint + " (requestType: " + requestType + ")") .isTrue(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientChangeFeedE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientChangeFeedE2ETest.java new file mode 100644 index 000000000000..8a24adb46121 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientChangeFeedE2ETest.java @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.rx; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnostics; +import com.azure.cosmos.models.CosmosBatch; +import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; +import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; + +/** + * Thin client E2E tests for change feed operations. + * Container is truncated in {@code @BeforeClass} — no per-test cleanup needed. + */ +public class ThinClientChangeFeedE2ETest extends ThinClientTestBase { + + @Factory(dataProvider = "clientBuildersWithGatewayAndHttp2") + public ThinClientChangeFeedE2ETest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientIncrementalChangeFeed() { + String pkValue = UUID.randomUUID().toString(); + ObjectNode doc1 = createTestDocument(UUID.randomUUID().toString(), pkValue); + ObjectNode doc2 = createTestDocument(UUID.randomUUID().toString(), pkValue); + + CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(pkValue)); + batch.createItemOperation(doc1); + batch.createItemOperation(doc2); + container.executeCosmosBatch(batch).block(); + + // Scope change feed to the specific logical partition to avoid + // consuming changes from other tests or partitions. + CosmosChangeFeedRequestOptions options = CosmosChangeFeedRequestOptions + .createForProcessingFromBeginning(FeedRange.forLogicalPartition(new PartitionKey(pkValue))); + + List changeFeedResults = new ArrayList<>(); + List allDiag = new ArrayList<>(); + for (FeedResponse page : container.queryChangeFeed(options, ObjectNode.class).byPage().toIterable()) { + changeFeedResults.addAll(page.getResults()); + allDiag.add(page.getCosmosDiagnostics()); + if (page.getResults().isEmpty()) { + break; + } + } + + assertThat(changeFeedResults.size()).isGreaterThanOrEqualTo(2); + for (CosmosDiagnostics d : allDiag) { + assertThinClientEndpointUsed(d); + } + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientChangeFeedFullRange() { + // Insert docs across two different partition keys so the full-range feed spans multiple partitions. + String pk1 = "cfFullRange1_" + UUID.randomUUID().toString().substring(0, 8); + String pk2 = "cfFullRange2_" + UUID.randomUUID().toString().substring(0, 8); + container.createItem(createTestDocument(UUID.randomUUID().toString(), pk1)).block(); + container.createItem(createTestDocument(UUID.randomUUID().toString(), pk2)).block(); + + CosmosChangeFeedRequestOptions options = CosmosChangeFeedRequestOptions + .createForProcessingFromBeginning(FeedRange.forFullRange()); + + List changeFeedResults = new ArrayList<>(); + List allDiag = new ArrayList<>(); + for (FeedResponse page : container.queryChangeFeed(options, ObjectNode.class).byPage().toIterable()) { + changeFeedResults.addAll(page.getResults()); + allDiag.add(page.getCosmosDiagnostics()); + if (page.getResults().isEmpty()) { + break; + } + } + + assertThat(changeFeedResults.size()).isGreaterThanOrEqualTo(2); + for (CosmosDiagnostics d : allDiag) { + assertThinClientEndpointUsed(d); + } + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientChangeFeedPartitionKey() { + String pkValue = "cfPk_" + UUID.randomUUID().toString().substring(0, 8); + container.createItem(createTestDocument(UUID.randomUUID().toString(), pkValue)).block(); + container.createItem(createTestDocument(UUID.randomUUID().toString(), pkValue)).block(); + + CosmosChangeFeedRequestOptions options = CosmosChangeFeedRequestOptions + .createForProcessingFromBeginning(FeedRange.forLogicalPartition(new PartitionKey(pkValue))); + + List changeFeedResults = new ArrayList<>(); + List allDiag = new ArrayList<>(); + for (FeedResponse page : container.queryChangeFeed(options, ObjectNode.class).byPage().toIterable()) { + changeFeedResults.addAll(page.getResults()); + allDiag.add(page.getCosmosDiagnostics()); + if (page.getResults().isEmpty()) { + break; + } + } + + // Should only see the 2 docs from this partition key + assertThat(changeFeedResults.size()).isEqualTo(2); + for (ObjectNode result : changeFeedResults) { + assertThat(result.get(PARTITION_KEY_FIELD).asText()).isEqualTo(pkValue); + } + for (CosmosDiagnostics d : allDiag) { + assertThinClientEndpointUsed(d); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientPointOperationE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientPointOperationE2ETest.java new file mode 100644 index 000000000000..1233f72dacfb --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientPointOperationE2ETest.java @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.rx; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.models.CosmosBatch; +import com.azure.cosmos.models.CosmosBatchResponse; +import com.azure.cosmos.models.CosmosBulkItemResponse; +import com.azure.cosmos.models.CosmosBulkOperationResponse; +import com.azure.cosmos.models.CosmosBulkOperations; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; + +/** + * Thin client E2E tests for point operations: Create, Read, Replace, Upsert, Patch, Delete, Bulk, Batch. + * Container is truncated in {@code @BeforeClass} — no per-test cleanup needed. + */ +public class ThinClientPointOperationE2ETest extends ThinClientTestBase { + + @Factory(dataProvider = "clientBuildersWithGatewayAndHttp2") + public ThinClientPointOperationE2ETest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientDocumentPointOperations() { + String idValue = UUID.randomUUID().toString(); + ObjectNode doc = createTestDocument(idValue, idValue); + + // create + CosmosItemResponse createResponse = container.createItem(doc).block(); + assertThat(createResponse.getStatusCode()).isEqualTo(201); + assertThat(createResponse.getRequestCharge()).isGreaterThan(0.0); + assertThinClientEndpointUsed(createResponse.getDiagnostics()); + + // read + CosmosItemResponse readResponse = container.readItem(idValue, new PartitionKey(idValue), ObjectNode.class).block(); + assertThat(readResponse.getStatusCode()).isEqualTo(200); + assertThinClientEndpointUsed(readResponse.getDiagnostics()); + + String idValue2 = UUID.randomUUID().toString(); + ObjectNode doc2 = createTestDocument(idValue2, idValue); + + // replace + CosmosItemResponse replaceResponse = container.replaceItem(doc2, idValue, new PartitionKey(idValue)).block(); + assertThat(replaceResponse.getStatusCode()).isEqualTo(200); + assertThinClientEndpointUsed(replaceResponse.getDiagnostics()); + + // upsert + ObjectNode doc3 = createTestDocument(idValue2, idValue); + doc3.put("newField", "newValue"); + CosmosItemResponse upsertResponse = container.upsertItem(doc3, new PartitionKey(idValue), new CosmosItemRequestOptions()).block(); + assertThat(upsertResponse.getStatusCode()).isEqualTo(200); + assertThinClientEndpointUsed(upsertResponse.getDiagnostics()); + + CosmosItemResponse readAfterUpsertResponse = container.readItem(idValue2, new PartitionKey(idValue), ObjectNode.class).block(); + assertThat(readAfterUpsertResponse.getItem().get("newField").asText()).isEqualTo("newValue"); + + // patch + CosmosPatchOperations patchOperations = CosmosPatchOperations.create(); + patchOperations.add("/anotherNewField", "anotherNewValue"); + patchOperations.replace("/newField", "patchedNewField"); + CosmosItemResponse patchResponse = container.patchItem(idValue2, new PartitionKey(idValue), patchOperations, ObjectNode.class).block(); + assertThat(patchResponse.getStatusCode()).isEqualTo(200); + assertThinClientEndpointUsed(patchResponse.getDiagnostics()); + + CosmosItemResponse readAfterPatchResponse = container.readItem(idValue2, new PartitionKey(idValue), ObjectNode.class).block(); + assertThat(readAfterPatchResponse.getItem().get("newField").asText()).isEqualTo("patchedNewField"); + assertThat(readAfterPatchResponse.getItem().get("anotherNewField").asText()).isEqualTo("anotherNewValue"); + + // delete + CosmosItemResponse deleteResponse = container.deleteItem(idValue2, new PartitionKey(idValue)).block(); + assertThat(deleteResponse.getStatusCode()).isEqualTo(204); + assertThinClientEndpointUsed(deleteResponse.getDiagnostics()); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientBulk() { + String idValue = UUID.randomUUID().toString(); + ObjectNode doc = createTestDocument(idValue, idValue); + + Flux> responsesFlux = container.executeBulkOperations(Flux.just( + CosmosBulkOperations.getCreateItemOperation(doc, new PartitionKey(idValue)) + )); + + List> responses = responsesFlux.collectList().block(); + assertThat(responses.size()).isEqualTo(1); + CosmosBulkItemResponse bulkResponse = responses.get(0).getResponse(); + assertThat(bulkResponse.isSuccessStatusCode()).isEqualTo(true); + assertThinClientEndpointUsed(bulkResponse.getCosmosDiagnostics()); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientBatch() { + String pkValue = UUID.randomUUID().toString(); + String idValue1 = UUID.randomUUID().toString(); + String idValue2 = UUID.randomUUID().toString(); + ObjectNode doc1 = createTestDocument(idValue1, pkValue); + ObjectNode doc2 = createTestDocument(idValue2, pkValue); + + CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(pkValue)); + batch.createItemOperation(doc1); + batch.createItemOperation(doc2); + + CosmosBatchResponse response = container.executeCosmosBatch(batch).block(); + assertThat(response.getStatusCode()).isEqualTo(200); + assertThinClientEndpointUsed(response.getDiagnostics()); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientStoredProcedureE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientStoredProcedureE2ETest.java new file mode 100644 index 000000000000..6af552954082 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ThinClientStoredProcedureE2ETest.java @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.rx; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.models.CosmosStoredProcedureProperties; +import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions; +import com.azure.cosmos.models.CosmosStoredProcedureResponse; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.Arrays; +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.Fail.fail; + +/** + * Thin client E2E tests for stored procedure execution. + * Container is truncated in {@code @BeforeClass} — no per-test cleanup needed. + */ +public class ThinClientStoredProcedureE2ETest extends ThinClientTestBase { + + @Factory(dataProvider = "clientBuildersWithGatewayAndHttp2") + public ThinClientStoredProcedureE2ETest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientStoredProcedure() { + String sprocId = "createDocSproc_" + UUID.randomUUID(); + String pkValue = UUID.randomUUID().toString(); + String docId = UUID.randomUUID().toString(); + + CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties( + sprocId, + "function createDocument(docToCreate) {" + + "var context = getContext();" + + "var container = context.getCollection();" + + "var response = context.getResponse();" + + "var accepted = container.createDocument(" + + " container.getSelfLink()," + + " docToCreate," + + " function(err, docCreated) {" + + " if (err) throw new Error('Error creating document: ' + err.message);" + + " response.setBody(docCreated);" + + " });" + + "if (!accepted) throw new Error('Document creation was not accepted');" + + "}" + ); + + CosmosStoredProcedureResponse createResponse = container.getScripts() + .createStoredProcedure(storedProcedureDef).block(); + assertThat(createResponse).isNotNull(); + assertThat(createResponse.getStatusCode()).isEqualTo(201); + + CosmosStoredProcedureRequestOptions options = new CosmosStoredProcedureRequestOptions(); + options.setPartitionKey(new PartitionKey(pkValue)); + + ObjectNode docToCreate = createTestDocument(docId, pkValue); + + CosmosStoredProcedureResponse executeResponse = container.getScripts() + .getStoredProcedure(sprocId) + .execute(Arrays.asList(docToCreate), options).block(); + + assertThat(executeResponse).isNotNull(); + assertThat(executeResponse.getStatusCode()).isEqualTo(200); + assertThat(executeResponse.getRequestCharge()).isGreaterThan(0.0); + assertThinClientEndpointUsed(executeResponse.getDiagnostics()); + + CosmosItemResponse readResponse = container.readItem(docId, new PartitionKey(pkValue), ObjectNode.class).block(); + assertThat(readResponse).isNotNull(); + assertThat(readResponse.getItem().get(ID_FIELD).asText()).isEqualTo(docId); + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testStoredProcedureExecutionWithoutPartitionKeyThrows() { + String sprocId = "noPartitionKeySproc_" + UUID.randomUUID(); + + CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties( + sprocId, "function() { getContext().getResponse().setBody('Hello'); }"); + + container.getScripts().createStoredProcedure(storedProcedureDef).block(); + + CosmosStoredProcedureRequestOptions options = new CosmosStoredProcedureRequestOptions(); + + try { + container.getScripts().getStoredProcedure(sprocId).execute(null, options).block(); + fail("Expected UnsupportedOperationException for sproc execution without partition key"); + } catch (UnsupportedOperationException e) { + assertThat(e.getMessage()).contains("PartitionKey value must be supplied"); + } + } + + @Test(groups = {"thinclient"}, timeOut = TIMEOUT) + public void testThinClientStoredProcedureWithPartitionKeyNone() { + String sprocId = "pkNoneSproc_" + UUID.randomUUID(); + + CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties( + sprocId, "function() { getContext().getResponse().setBody('Hello from PK.NONE'); }"); + + container.getScripts().createStoredProcedure(storedProcedureDef).block(); + + CosmosStoredProcedureRequestOptions options = new CosmosStoredProcedureRequestOptions(); + options.setPartitionKey(PartitionKey.NONE); + + CosmosStoredProcedureResponse executeResponse = container.getScripts() + .getStoredProcedure(sprocId).execute(null, options).block(); + + assertThat(executeResponse).isNotNull(); + assertThat(executeResponse.getStatusCode()).isEqualTo(200); + assertThat(executeResponse.getRequestCharge()).isGreaterThan(0.0); + assertThinClientEndpointUsed(executeResponse.getDiagnostics()); + } +} 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 ceb703cc5154..e4d5faea76e1 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,38 +2,79 @@ // 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; /** - * 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. + * Base class for thin client E2E tests. Provides shared setup/teardown, + * constants, and helper methods common to all thin client test classes. */ -public final class ThinClientTestBase { +public abstract class ThinClientTestBase extends TestSuiteBase { - 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 static final String ID_FIELD = "id"; + protected static final String PARTITION_KEY_FIELD = "mypk"; + protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private ThinClientTestBase() { + protected CosmosAsyncClient client; + protected CosmosAsyncContainer container; + + protected ThinClientTestBase(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT) + public void before_ThinClientTest() { + // Thin client is enabled JVM-wide by the "thinclient" CI lane via + // -DCOSMOS.THINCLIENT_ENABLED=true (+ -DCOSMOS.HTTP2_ENABLED=true). Every test/lifecycle + // method in these classes is scoped to groups={"thinclient"}, so they only ever run in that + // lane. We therefore rely on the ambient flag rather than mutating the JVM-global property + // per class (which is read lazily per client build and would otherwise leak across the + // classes that share this JVM). + 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; } /** - * Asserts that all requests in the diagnostics were routed through the thin client endpoint. + * Asserts that all data requests in the diagnostics were routed through the thin client endpoint. */ 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. + // (every request via the thin-client endpoint -- including QueryPlan, which is routed to + // Gateway V2 when thin client + HTTP/2 are opted in) and null-endpoint handling are applied + // consistently across all thin-client tests. TestSuiteBase.assertThinClientEndpointUsed(diagnostics); } diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index e20aee6e0b70..29e3537406c9 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.82.0-beta.1 (Unreleased) #### Features Added +* Added support for QueryPlan and Execute Stored Procedure requests to be routed to Gateway V2. - See [PR 47759](https://github.com/Azure/azure-sdk-for-java/pull/47759) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index cf1a812e10f2..b7d1097fe443 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -55,6 +55,12 @@ public class Configs { private static final String THINCLIENT_ENABLED = "COSMOS.THINCLIENT_ENABLED"; private static final String THINCLIENT_ENABLED_VARIABLE = "COSMOS_THINCLIENT_ENABLED"; + // Kill-switch to opt out of routing QueryPlan requests through the thin client (Gateway V2). + // Defaults to enabled; set to false to force QueryPlan requests back onto Gateway V1. + private static final boolean DEFAULT_THINCLIENT_QUERY_PLAN_ENABLED = true; + private static final String THINCLIENT_QUERY_PLAN_ENABLED = "COSMOS.THINCLIENT_QUERY_PLAN_ENABLED"; + private static final String THINCLIENT_QUERY_PLAN_ENABLED_VARIABLE = "COSMOS_THINCLIENT_QUERY_PLAN_ENABLED"; + private static final boolean DEFAULT_NETTY_HTTP_CLIENT_METRICS_ENABLED = false; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED = "COSMOS.NETTY_HTTP_CLIENT_METRICS_ENABLED"; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED_VARIABLE = "COSMOS_NETTY_HTTP_CLIENT_METRICS_ENABLED"; @@ -585,6 +591,20 @@ public static boolean isThinClientEnabled() { return DEFAULT_THINCLIENT_ENABLED; } + public static boolean isThinClientQueryPlanEnabled() { + String valueFromSystemProperty = System.getProperty(THINCLIENT_QUERY_PLAN_ENABLED); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + return Boolean.parseBoolean(valueFromSystemProperty); + } + + String valueFromEnvVariable = System.getenv(THINCLIENT_QUERY_PLAN_ENABLED_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + return Boolean.parseBoolean(valueFromEnvVariable); + } + + return DEFAULT_THINCLIENT_QUERY_PLAN_ENABLED; + } + public static boolean isNettyHttpClientMetricsEnabled() { return Boolean.parseBoolean( System.getProperty(NETTY_HTTP_CLIENT_METRICS_ENABLED, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 049d18a3dd1e..e5fcbfeebe74 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4691,7 +4691,7 @@ private Flux> readManyByPartitionKeys( Mono queryValidationMono; if (customQuery != null) { queryValidationMono = validateCustomQueryForReadManyByPartitionKeys( - customQuery, resourceLink, state.getQueryOptions()); + customQuery, resourceLink, state.getQueryOptions(), collection); } else { queryValidationMono = Mono.empty(); } @@ -5135,7 +5135,8 @@ private static final class BatchDescriptor { private Mono validateCustomQueryForReadManyByPartitionKeys( SqlQuerySpec customQuery, String resourceLink, - CosmosQueryRequestOptions queryRequestOptions) { + CosmosQueryRequestOptions queryRequestOptions, + DocumentCollection collection) { IDocumentQueryClient queryClient = documentQueryClientImpl( RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); @@ -5147,6 +5148,7 @@ private Mono validateCustomQueryForReadManyByPartitionKeys( customQuery, resourceLink, queryRequestOptions, + collection, Configs.isQueryPlanCachingEnabled(), this.getQueryPlanCache()) .doOnNext(RxDocumentClientImpl::validateQueryPlanForReadManyByPartitionKeys) @@ -5663,6 +5665,11 @@ public GlobalEndpointManager getGlobalEndpointManager() { public GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker getGlobalPartitionEndpointManagerForCircuitBreaker() { return RxDocumentClientImpl.this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker; } + + @Override + public boolean useThinClient(RxDocumentServiceRequest request) { + return RxDocumentClientImpl.this.useThinClientStoreModel(request); + } }; } @@ -7444,7 +7451,8 @@ private RxStoreModel getStoreProxy(RxDocumentServiceRequest request) { resourceType == ResourceType.ClientEncryptionKey || resourceType.isScript() && operationType != OperationType.ExecuteJavaScript || resourceType == ResourceType.PartitionKeyRange || - resourceType == ResourceType.PartitionKey && operationType == OperationType.Delete) { + resourceType == ResourceType.PartitionKey && operationType == OperationType.Delete || + operationType == OperationType.QueryPlan) { return this.gatewayProxy; } @@ -7482,7 +7490,7 @@ private RxStoreModel getStoreProxy(RxDocumentServiceRequest request) { if ((operationType == OperationType.Query || operationType == OperationType.SqlQuery || operationType == OperationType.ReadFeed) && - Utils.isCollectionChild(request.getResourceType())) { + Utils.isCollectionChild(request.getResourceType())) { // Go to gateway only when partition key range and partition key are not set. This should be very rare if (request.getPartitionKeyRangeIdentity() == null && request.getHeaders().get(HttpConstants.HttpHeaders.PARTITION_KEY) == null) { @@ -9044,7 +9052,8 @@ public boolean useThinClient() { private boolean useThinClientStoreModel(RxDocumentServiceRequest request) { if (!useThinClient || !this.globalEndpointManager.hasThinClientReadLocations() - || request.getResourceType() != ResourceType.Document) { + || (request.getResourceType() != ResourceType.Document + && !request.isExecuteStoredProcedureBasedRequest())) { return false; } @@ -9054,7 +9063,10 @@ private boolean useThinClientStoreModel(RxDocumentServiceRequest request) { return operationType.isPointOperation() || operationType == OperationType.Query || operationType == OperationType.Batch - || request.isChangeFeedRequest() && !request.isAllVersionsAndDeletesChangeFeedMode(); + || (request.isChangeFeedRequest() + && !request.isAllVersionsAndDeletesChangeFeedMode()) + || request.isExecuteStoredProcedureBasedRequest() + || operationType == OperationType.QueryPlan; } private DocumentClientRetryPolicy getRetryPolicyForPointOperation( 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..1e2958040978 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 @@ -968,6 +968,10 @@ public boolean isChangeFeedRequest() { return this.headers.containsKey(HttpConstants.HttpHeaders.A_IM); } + public boolean isExecuteStoredProcedureBasedRequest() { + return this.resourceType == ResourceType.StoredProcedure && this.operationType == OperationType.ExecuteJavaScript; + } + public boolean isAllVersionsAndDeletesChangeFeedMode() { String aImHeader = this.headers.get(HttpConstants.HttpHeaders.A_IM); return this.headers.containsKey(HttpConstants.HttpHeaders.A_IM) && HttpConstants.A_IMHeaderValues.FULL_FIDELITY_FEED.equals(aImHeader); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index 8d9b324d92fb..06e6353fd27e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -861,7 +861,24 @@ private void validateOrThrow(RxDocumentServiceRequest request, safeSilentRelease(retainedBodyAsByteBuf); CosmosError cosmosError; - cosmosError = (StringUtils.isNotEmpty(body)) ? new CosmosError(body) : new CosmosError(); + if (StringUtils.isNotEmpty(body)) { + try { + cosmosError = new CosmosError(body); + } catch (IllegalArgumentException | ClassCastException jsonParseError) { + // Gateway V2 / thin-client error responses (notably query-plan generation + // failures) can carry a raw, non-JSON body (optionally NUL-padded) or a non-object + // JSON value (scalar, array) instead of a serialized CosmosError. Parsing it throws + // (IllegalArgumentException for non-JSON, ClassCastException for valid non-object + // JSON) and, upstream, the real HTTP status code is lost and the failure surfaces + // as statusCode 0. Fall back to the raw (sanitized) body as the message so the + // actual status code is preserved on the thrown CosmosException. + logger.debug("Failed to parse gateway error body as CosmosError; " + + "falling back to raw body. statusCode: {}", statusCodeString, jsonParseError); + cosmosError = new CosmosError(statusCodeString, sanitizeErrorBody(body)); + } + } else { + cosmosError = new CosmosError(); + } cosmosError = new CosmosError(statusCodeString, String.format("%s, StatusCode: %s", cosmosError.getMessage(), statusCodeString), cosmosError.getPartitionedQueryExecutionInfo()); @@ -872,6 +889,13 @@ private void validateOrThrow(RxDocumentServiceRequest request, } } + private static String sanitizeErrorBody(String body) { + // Proxies may NUL-pad fixed-size buffers; strip padding and surrounding whitespace so the + // preserved error message is readable. If stripping leaves nothing, keep the original body. + String sanitized = body.replace("\u0000", "").trim(); + return sanitized.isEmpty() ? body : sanitized; + } + private static HttpMethod getHttpMethod(RxDocumentServiceRequest request) { switch (request.getOperationType()) { case Create: 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 815d88995afe..95973424f24b 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,7 +251,13 @@ 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); String readFeedKeyType = request.getHeaders().get(HttpConstants.HttpHeaders.READ_FEED_KEY_TYPE); - if (startEpk != null && endEpk != null + if (request.getOperationType() == OperationType.QueryPlan) { + // QueryPlan is collection-scoped on the thin-client proxy: it carries no + // EffectivePartitionKey and no StartEpkHash/EndEpkHash headers - the proxy fans out + // across partitions itself. Keep this explicit so the contract is self-documenting and + // cannot be violated if a resolved partition key range is ever present on the request. + //noinspection StatementWithEmptyBody + } else 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) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index 15b4b50df559..f9eb191fc20c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -283,7 +283,8 @@ public enum RntbdOperationType { PreReplaceValidation((short) 0x0020, OperationType.PreReplaceValidation), AddComputeGatewayRequestCharges((short) 0x0021, OperationType.AddComputeGatewayRequestCharges), MigratePartition((short) 0x0022, OperationType.MigratePartition), - Batch((short) 0x0025, OperationType.Batch); + Batch((short) 0x0025, OperationType.Batch), + QueryPlan((short) 0x0042, OperationType.QueryPlan); private final short id; private final OperationType type; @@ -367,6 +368,8 @@ public static RntbdOperationType fromId(final short id) { return RntbdOperationType.MigratePartition; case 0x0025: return RntbdOperationType.Batch; + case 0x0042: + return RntbdOperationType.QueryPlan; default: throw new DecoderException(String.format("expected byte value matching %s value, not %s", RntbdOperationType.class.getSimpleName(), @@ -444,6 +447,8 @@ public static RntbdOperationType fromType(OperationType type) { return RntbdOperationType.AddComputeGatewayRequestCharges; case Batch: return RntbdOperationType.Batch; + case QueryPlan: + return RntbdOperationType.QueryPlan; default: throw new IllegalArgumentException(String.format("unrecognized operation type: %s", type)); } @@ -619,7 +624,10 @@ public enum RntbdRequestHeader implements RntbdHeader { ThroughputBucket((short)0x00DB, RntbdTokenType.Byte, false), WorkloadId((short)0x00DC, RntbdTokenType.Byte, false), HubRegionProcessingOnly((short)0x00EF, RntbdTokenType.Byte , false), - ReadConsistencyStrategy((short)0x00FE, RntbdTokenType.Byte, false); + ReadConsistencyStrategy((short)0x00FE, RntbdTokenType.Byte, false), + // QueryPlan headers for proxy — IDs match server-side RntbdConstants.cs + SupportedQueryFeatures((short) 0x00FF, RntbdTokenType.String, false), + QueryVersion((short) 0x0100, RntbdTokenType.SmallString, false); public static final List thinClientHeadersInOrderList = Arrays.asList( EffectivePartitionKey, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java index 8eb632c48c8c..294b7399d6f5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java @@ -204,6 +204,8 @@ private static RntbdOperationType map(final OperationType operationType) { return RntbdOperationType.AddComputeGatewayRequestCharges; case Batch: return RntbdOperationType.Batch; + case QueryPlan: + return RntbdOperationType.QueryPlan; default: final String reason = String.format("Unrecognized operation type: %s", operationType); throw new UnsupportedOperationException(reason); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index 38de8021c44f..d00f397790c9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -195,6 +195,11 @@ private static ImplementationBridgeHelpers.PriorityLevelHelper.PriorityLevelAcce // and BE will respect the per-request value. this.fillTokenFromHeader(headers, this::getClientVersion, HttpHeaders.VERSION); + + // QueryPlan headers — needed for proxy to extract supported features and query version + // from the RNTBD body (IDs match server-side proxy) + this.fillTokenFromHeader(headers, this::getSupportedQueryFeatures, HttpHeaders.SUPPORTED_QUERY_FEATURES); + this.fillTokenFromHeader(headers, this::getQueryVersion, HttpHeaders.QUERY_VERSION); } private RntbdRequestHeaders(ByteBuf in) { @@ -653,6 +658,14 @@ private RntbdToken getChangeFeedWireFormatVersion() { return this.get(RntbdRequestHeader.ChangeFeedWireFormatVersion); } + private RntbdToken getSupportedQueryFeatures() { + return this.get(RntbdRequestHeader.SupportedQueryFeatures); + } + + private RntbdToken getQueryVersion() { + return this.get(RntbdRequestHeader.QueryVersion); + } + private void addAimHeader(final Map headers) { final String value = headers.get(HttpHeaders.A_IM); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index 11528ed4f7ea..c71e704d8c0e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -127,6 +127,7 @@ private static Mono getPartitionKeyRangesAn query, resourceLink, cosmosQueryRequestOptions, + collection, queryPlanCachingEnabled, queryPlanCache) .flatMap( @@ -145,6 +146,7 @@ private static Mono fetchQueryPlan( SqlQuerySpec query, String resourceLink, CosmosQueryRequestOptions cosmosQueryRequestOptions, + DocumentCollection collection, boolean queryPlanCachingEnabled, Map queryPlanCache) { @@ -163,7 +165,8 @@ private static Mono fetchQueryPlan( client, query, resourceLink, - cosmosQueryRequestOptions) + cosmosQueryRequestOptions, + collection) .doOnNext(partitionedQueryExecutionInfo -> { if (queryPlanCachingEnabled && isScopedToSinglePartition(cosmosQueryRequestOptions)) { tryCacheQueryPlan(query, partitionedQueryExecutionInfo, queryPlanCache); @@ -346,6 +349,7 @@ public static Mono fetchQueryPlanForValidation( SqlQuerySpec sqlQuerySpec, String resourceLink, CosmosQueryRequestOptions queryRequestOptions, + DocumentCollection collection, boolean queryPlanCachingEnabled, Map queryPlanCache) { @@ -355,6 +359,7 @@ public static Mono fetchQueryPlanForValidation( sqlQuerySpec, resourceLink, queryRequestOptions, + collection, queryPlanCachingEnabled, queryPlanCache); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java index 8dd42f5aa32c..830a37a2dda4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java @@ -4,6 +4,7 @@ package com.azure.cosmos.implementation.query; import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.ClientSideRequestStatistics; import com.azure.cosmos.implementation.DiagnosticsClientContext; @@ -173,6 +174,7 @@ private void initialize( aggregatedGlobalStatistics = Flux.fromIterable(globalStatsProducers) .flatMap(producer -> producer.produceAsync() .map(documentProducerFeedResponse -> { + this.captureClientSideRequestStatistics(documentProducerFeedResponse.pageResult); List results = documentProducerFeedResponse.pageResult.getResults(); return new GlobalFullTextSearchQueryStatistics(results.get(0)); })) @@ -443,10 +445,27 @@ private Flux getComponentQueryResults(List targetFee return Flux.fromIterable(componentProducers) .flatMap(DocumentProducer::produceAsync) + .doOnNext(response -> this.captureClientSideRequestStatistics(response.pageResult)) .flatMap(response -> Flux.fromIterable(response.pageResult.getResults())); }); } + private void captureClientSideRequestStatistics(FeedResponse response) { + if (response == null || response.getCosmosDiagnostics() == null) { + return; + } + + CosmosDiagnostics diagnostics = response.getCosmosDiagnostics(); + Collection requestStatistics = + diagAccessor().getFeedResponseDiagnostics(diagnostics) != null + ? diagAccessor().getClientSideRequestStatisticsForQueryPipelineAggregations(diagnostics) + : diagAccessor().getClientSideRequestStatistics(diagnostics); + + if (requestStatistics != null && !requestStatistics.isEmpty()) { + this.clientSideRequestStatistics.addAll(requestStatistics); + } + } + private Flux retrieveRewrittenQueryInfos(List componentQueryInfos) { return aggregatedGlobalStatistics.hasElement().flatMapMany(globalStatistics -> { if (globalStatistics != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IDocumentQueryClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IDocumentQueryClient.java index 91679d40490e..bfc2293f6aeb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IDocumentQueryClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IDocumentQueryClient.java @@ -110,4 +110,6 @@ Mono addPartitionLevelUnavailableRegionsOnRequest( GlobalEndpointManager getGlobalEndpointManager(); GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker getGlobalPartitionEndpointManagerForCircuitBreaker(); + + boolean useThinClient(RxDocumentServiceRequest request); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PartitionedQueryExecutionInfo.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PartitionedQueryExecutionInfo.java index 4967d2b1e6be..49599e0cbbd1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PartitionedQueryExecutionInfo.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PartitionedQueryExecutionInfo.java @@ -5,10 +5,11 @@ import com.azure.cosmos.implementation.RequestTimeline; import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo; +import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.implementation.JsonSerializable; -import com.azure.cosmos.implementation.Constants; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.azure.cosmos.models.PartitionKeyDefinition; import java.util.List; @@ -25,23 +26,31 @@ public final class PartitionedQueryExecutionInfo extends JsonSerializable { private RequestTimeline queryPlanRequestTimeline; private HybridSearchQueryInfo hybridSearchQueryInfo; - PartitionedQueryExecutionInfo(QueryInfo queryInfo, List> queryRanges) { - this.queryInfo = queryInfo; - this.queryRanges = queryRanges; - - this.set( - PartitionedQueryExecutionInfoInternal.PARTITIONED_QUERY_EXECUTION_INFO_VERSION_PROPERTY, - Constants.PartitionedQueryExecutionInfo.VERSION_1 - ); - } - + /** + * Constructs with EPK hex string format expected for queryRanges. + */ public PartitionedQueryExecutionInfo(ObjectNode content, RequestTimeline queryPlanRequestTimeline) { super(content); this.queryPlanRequestTimeline = queryPlanRequestTimeline; } - public PartitionedQueryExecutionInfo(String jsonString) { - super(jsonString); + /** + * Constructs with PartitionKeyInternal array format expected for queryRanges. + * The {@code partitionKeyDefinition} is required for converting PartitionKeyInternal + * values to EPK hex strings. + * + * @param partitionKeyDefinition the container's partition key definition, must not be null. + */ + PartitionedQueryExecutionInfo(ObjectNode content, RequestTimeline queryPlanRequestTimeline, PartitionKeyDefinition partitionKeyDefinition) { + super(content); + if (partitionKeyDefinition == null) { + throw new IllegalArgumentException("partitionKeyDefinition must not be null"); + } + this.queryPlanRequestTimeline = queryPlanRequestTimeline; + this.queryRanges = PartitionKeyInternalHelper.convertToSortedEpkRanges( + PartitionedQueryExecutionInfoInternal.QUERY_RANGES_PROPERTY, + content, + partitionKeyDefinition); } public int getVersion() { @@ -54,10 +63,26 @@ public QueryInfo getQueryInfo() { PartitionedQueryExecutionInfoInternal.QUERY_INFO_PROPERTY, QueryInfo.class)); } + /** + * Returns the query ranges as sorted EPK hex string ranges. + *

+ * Two formats exist: + *

    + *
  • Gateway V1: {@code min}/{@code max} are EPK hex strings and are + * deserialized directly via {@code getList()}.
    • + *
  • Thin client proxy: {@code min}/{@code max} are PartitionKeyInternal + * JSON arrays and are converted to EPK hex by the thin-client constructor.
    • + *
+ */ public List> getQueryRanges() { - return this.queryRanges != null ? this.queryRanges - : (this.queryRanges = super.getList( - PartitionedQueryExecutionInfoInternal.QUERY_RANGES_PROPERTY, QUERY_RANGES_CLASS)); + if (this.queryRanges != null) { + return this.queryRanges; + } + + // EPK hex string format — direct deserialization + this.queryRanges = super.getList( + PartitionedQueryExecutionInfoInternal.QUERY_RANGES_PROPERTY, QUERY_RANGES_CLASS); + return this.queryRanges; } public RequestTimeline getQueryPlanRequestTimeline() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryInfo.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryInfo.java index b8ea415ebf2f..e65101b02d51 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryInfo.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryInfo.java @@ -16,6 +16,7 @@ import java.time.Duration; import java.time.Instant; import java.util.Collection; +import java.util.HashMap; import java.util.List; import java.util.Map; @@ -166,10 +167,23 @@ public boolean hasNonStreamingOrderBy() { return this.nonStreamingOrderBy; } - public Map getGroupByAliasToAggregateType(){ - Map groupByAliasToAggregateMap; - groupByAliasToAggregateMap = super.getMap("groupByAliasToAggregateType"); - return groupByAliasToAggregateMap; + public Map getGroupByAliasToAggregateType() { + Map rawMap = super.getMap("groupByAliasToAggregateType"); + if (rawMap == null) { + return null; + } + + Map groupByAliasToAggregateMap = new HashMap<>(rawMap.size()); + rawMap.forEach((key, value) -> { + if (value == null || (value instanceof String && ((String) value).isEmpty())) { + groupByAliasToAggregateMap.put(key, null); + } else if (value instanceof AggregateOperator) { + groupByAliasToAggregateMap.put(key, (AggregateOperator) value); + } else { + groupByAliasToAggregateMap.put(key, AggregateOperator.valueOf(String.valueOf(value))); + } + }); + return groupByAliasToAggregateMap; } public List getGroupByAliases() { @@ -267,4 +281,3 @@ public int hashCode() { return super.hashCode(); } } - diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java index 591b601eb7c4..d2f53bb47878 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java @@ -8,13 +8,18 @@ import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DiagnosticsClientContext; +import com.azure.cosmos.implementation.DocumentCollection; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.PathsHelper; +import com.azure.cosmos.implementation.RequestTimeline; import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; +import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; +import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyDefinition; import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.implementation.BackoffRetryUtility; @@ -50,6 +55,9 @@ private static ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.Cosmo // new NonStreamingOrderBy query feature the client might run into some issue of not being able to recognize this, // and throw a 400 exception. If the environment variable `AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY` is set to // True to opt out of this new query feature, we will return the OLD query features to operate correctly. + // TODO: Add ListAndSetAggregate after Java supports MAKELIST/MAKESET query aggregation. + // TODO: Add HybridSearchSkipOrderByRewrite after the Java hybrid query pipeline can consume the optimized plan. + // TODO: Add CountIf after Java implements a CountIf aggregator in SingleGroupAggregator. private static final String SUPPORTED_QUERY_FEATURES = QueryFeature.Aggregate.name() + ", " + QueryFeature.CompositeAggregate.name() + ", " + QueryFeature.MultipleOrderBy.name() + ", " + @@ -81,7 +89,8 @@ static Mono getQueryPlanThroughGatewayAsync(Diagn IDocumentQueryClient queryClient, SqlQuerySpec sqlQuerySpec, String resourceLink, - CosmosQueryRequestOptions initialQueryRequestOptions) { + CosmosQueryRequestOptions initialQueryRequestOptions, + DocumentCollection collection) { CosmosQueryRequestOptions nonNullRequestOptions = initialQueryRequestOptions != null ? initialQueryRequestOptions @@ -89,6 +98,8 @@ static Mono getQueryPlanThroughGatewayAsync(Diagn PartitionKey partitionKey = nonNullRequestOptions.getPartitionKey(); + PartitionKeyDefinition partitionKeyDefinition = collection != null ? collection.getPartitionKey() : null; + final Map requestHeaders = new HashMap<>(); requestHeaders.put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON); requestHeaders.put(HttpConstants.HttpHeaders.IS_QUERY_PLAN_REQUEST, TRUE); @@ -106,7 +117,14 @@ static Mono getQueryPlanThroughGatewayAsync(Diagn ResourceType.Document, resourceLink, requestHeaders); - queryPlanRequest.useGatewayMode = true; + // Route the query-plan request based on the thin-client opt-in, not on collection + // metadata. When thin client is enabled and eligible for this request, leave gateway + // mode off so the request is routed through the Gateway V2 thin-client proxy; + // otherwise force Gateway V1 (preserving the legacy query-plan routing). + // Configs.isThinClientQueryPlanEnabled() is a kill-switch (system property / + // environment variable) that forces query-plan requests back onto Gateway V1. + queryPlanRequest.useGatewayMode = + !(queryClient.useThinClient(queryPlanRequest) && Configs.isThinClientQueryPlanEnabled()); // Create a defensive copy to prevent concurrent modification of the shared // SqlQuerySpec's internal ObjectNode when multiple threads retrieve the query @@ -138,10 +156,29 @@ static Mono getQueryPlanThroughGatewayAsync(Diagn return BackoffRetryUtility.executeRetry(() -> { retryPolicyInstance.onBeforeSendRequest(req); return queryClient.executeQueryAsync(req).flatMap(rxDocumentServiceResponse -> { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = - new PartitionedQueryExecutionInfo( - (ObjectNode) rxDocumentServiceResponse.getResponseBody(), - rxDocumentServiceResponse.getGatewayHttpRequestTimeline()); + ObjectNode responseBody = (ObjectNode) rxDocumentServiceResponse.getResponseBody(); + RequestTimeline timeline = rxDocumentServiceResponse.getGatewayHttpRequestTimeline(); + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + + // In thin client mode, the proxy returns queryRanges in PartitionKeyInternal + // format (e.g., {"min": ["value"], "max": ["Infinity"]}). Convert to sorted + // List> with EPK hex strings and pass directly to the DTO — + // avoiding a redundant JSON round-trip. + if (req.useThinClientMode && partitionKeyDefinition == null) { + throw new IllegalStateException( + "PartitionKeyDefinition must not be null in thin client mode. " + + "Ensure DocumentCollection is resolved before calling getQueryPlanThroughGatewayAsync."); + } + + if (req.useThinClientMode) { + partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo( + responseBody, timeline, partitionKeyDefinition); + } else { + partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo( + responseBody, timeline); + } + return Mono.just(partitionedQueryExecutionInfo); }); }, retryPolicyInstance); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternalHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternalHelper.java index 2803536084be..fd606a20aa4e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternalHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternalHelper.java @@ -9,13 +9,24 @@ import com.azure.cosmos.implementation.ByteBufferOutputStream; import com.azure.cosmos.implementation.Bytes; import com.azure.cosmos.implementation.RMResources; +import com.azure.cosmos.implementation.Utils; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; import java.io.IOException; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.List; +import java.util.Spliterators; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; public class PartitionKeyInternalHelper { + private static final Range.MinComparator MIN_COMPARATOR = new Range.MinComparator<>(); + public static final String MinimumInclusiveEffectivePartitionKey = toHexEncodedBinaryString(PartitionKeyInternal.EmptyPartitionKey.components); public static final byte[] MinimumInclusiveEffectivePartitionKeyBytes = toBinary(PartitionKeyInternal.EmptyPartitionKey.components); public static final String MaximumExclusiveEffectivePartitionKey = toHexEncodedBinaryString(PartitionKeyInternal.InfinityPartitionKey.components); @@ -294,4 +305,111 @@ static public Range getEPKRangeForPrefixPartitionKey( String maxEPK = minEPK + MaximumExclusiveEffectivePartitionKey; return new Range<>(minEPK, maxEPK, true, false); } + + /** + * Converts query ranges from PartitionKeyInternal JSON format to sorted EPK hex string ranges. + * + *

The thin client proxy returns queryRanges as PartitionKeyInternal JSON arrays + * (e.g., {@code {"min": ["value"], "max": ["Infinity"]}}). This method parses each range, + * computes the EPK hex string via {@link #getEffectivePartitionKeyString}, and sorts the + * result using {@link Range.MinComparator} to satisfy + * {@link RoutingMapProviderHelper#getOverlappingRanges} which requires sorted, non-overlapping input. + * + * @param queryRangesProperty the name of the JSON property containing the ranges array + * @param queryPlanJson the raw query plan JSON containing PartitionKeyInternal ranges + * @param partitionKeyDefinition the container's partition key definition + * @return sorted list of EPK hex string ranges; empty list if the property is absent + */ + public static List> convertToSortedEpkRanges( + String queryRangesProperty, + ObjectNode queryPlanJson, + PartitionKeyDefinition partitionKeyDefinition) { + + JsonNode queryRangesNode = queryPlanJson.get(queryRangesProperty); + if (queryRangesNode == null || !queryRangesNode.isArray()) { + String actualType = queryRangesNode == null ? "null (property absent)" : queryRangesNode.getNodeType().name(); + throw new IllegalStateException( + "Thin client proxy query plan response has missing or invalid '" + queryRangesProperty + "' property. " + + "Expected: JSON array of {min, max, isMinInclusive, isMaxInclusive} range objects. " + + "Actual node type: " + actualType + ". " + + "Response keys: " + StreamSupport.stream( + Spliterators.spliteratorUnknownSize(queryPlanJson.fieldNames(), 0), false) + .collect(Collectors.joining(", ")) + ". " + + "This indicates a protocol mismatch between the SDK and the thin client proxy."); + } + + ArrayNode rawRanges = (ArrayNode) queryRangesNode; + if (rawRanges.isEmpty()) { + throw new IllegalStateException( + "Thin client proxy query plan response '" + queryRangesProperty + "' array must not be empty. " + + "Expected: at least one range covering the partitions to fan out to. " + + "An empty array would silently produce zero query results."); + } + + List> epkRanges = new ArrayList<>(rawRanges.size()); + + for (JsonNode rangeNode : rawRanges) { + if (!rangeNode.isObject()) { + throw new IllegalStateException( + "Thin client proxy query plan response contains a non-object element in queryRanges array. " + + "Expected: JSON object with {min, max, isMinInclusive, isMaxInclusive}. " + + "Array size: " + rawRanges.size() + ". " + + "Actual node type: " + rangeNode.getNodeType().name() + "."); + } + ObjectNode rangeObj = (ObjectNode) rangeNode; + + String minEpk = partitionKeyInternalToEpkString(rangeObj.get("min"), partitionKeyDefinition); + String maxEpk = partitionKeyInternalToEpkString(rangeObj.get("max"), partitionKeyDefinition); + + JsonNode minInclusiveNode = rangeObj.get("isMinInclusive"); + JsonNode maxInclusiveNode = rangeObj.get("isMaxInclusive"); + if (minInclusiveNode == null || maxInclusiveNode == null) { + throw new IllegalStateException( + "Thin client proxy query plan range missing required fields. " + + "Expected: isMinInclusive and isMaxInclusive. " + + "Range object fields: " + StreamSupport.stream( + Spliterators.spliteratorUnknownSize(rangeObj.fieldNames(), 0), false) + .collect(Collectors.joining(", ")) + "."); + } + boolean isMinInclusive = minInclusiveNode.asBoolean(); + boolean isMaxInclusive = maxInclusiveNode.asBoolean(); + + epkRanges.add(new Range<>(minEpk, maxEpk, isMinInclusive, isMaxInclusive)); + } + + epkRanges.sort(MIN_COMPARATOR); + return epkRanges; + } + + /** + * Converts a single PartitionKeyInternal JSON node to its EPK hex string representation. + * + * @param rangeBoundaryNode the JSON node representing a range boundary (min or max) in PartitionKeyInternal format + * @param partitionKeyDefinition the container's partition key definition + * @return the EPK hex string + */ + private static String partitionKeyInternalToEpkString(JsonNode rangeBoundaryNode, PartitionKeyDefinition partitionKeyDefinition) { + if (rangeBoundaryNode == null || rangeBoundaryNode.isNull()) { + throw new IllegalStateException( + "Thin client proxy query plan range has null boundary value. " + + "Expected: PartitionKeyInternal JSON array (e.g., [\"value\"] or [{\"type\":\"Infinity\"}])."); + } + + PartitionKeyInternal partitionKey; + try { + partitionKey = Utils.getSimpleObjectMapper().treeToValue(rangeBoundaryNode, PartitionKeyInternal.class); + } catch (JsonProcessingException e) { + throw new IllegalStateException( + "Failed to parse PartitionKeyInternal from range boundary. " + + "Boundary node type: " + rangeBoundaryNode.getNodeType().name() + ".", e); + } + + if (partitionKey.getComponents() == null) { + throw new IllegalStateException( + "Thin client proxy query plan range boundary deserialized to NonePartitionKey (null components). " + + "Boundary node type: " + rangeBoundaryNode.getNodeType().name() + "."); + } + + return getEffectivePartitionKeyString(partitionKey, partitionKeyDefinition); + } }