From 6eb72a6aea326e5a780c3dce782b989d01ba88c3 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 11 Aug 2026 15:43:12 -0700 Subject: [PATCH 1/8] Add binary encoding for query pages Negotiate Cosmos binary JSON for SQL query and document read-feed operations while keeping query-plan and change-feed requests text-only. Normalize binary pages at pipeline ingest and restore the planned response format at emit so DISTINCT and streaming ORDER BY retain their existing text-based internals. Preserve the response encoding selected during planning even when execution options differ. Make emulator responses and feed item slices binary-aware, pin numeric DISTINCT invariants, and add text/binary parity coverage to emulator, fuzzer, and live split test surfaces. Update the binary encoding design and changelogs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 1 + .../tests/binary_roundtrip_fuzzer.rs | 109 +++++++++++- .../cosmos_query_distinct_split.rs | 13 +- .../azure_data_cosmos_driver/CHANGELOG.md | 2 + .../docs/BINARY_ENCODING_HLD.md | 11 +- .../src/driver/cosmos_driver.rs | 142 +++++++++------ .../src/driver/dataflow/distinct_hash.rs | 13 ++ .../src/driver/dataflow/pipeline.rs | 15 ++ .../src/driver/dataflow/query_response.rs | 116 +++++++++++- .../src/in_memory_emulator/operations.rs | 23 ++- .../src/models/cosmos_response.rs | 10 ++ .../src/models/mod.rs | 18 +- .../src/models/response_body.rs | 60 +++++-- .../in_memory_emulator_tests/distinct.rs | 168 +++++++++++++++++- .../in_memory_emulator_tests/read_feed.rs | 34 ++++ 15 files changed, 628 insertions(+), 107 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index e8e88ac88ec..78148236d88 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -11,6 +11,7 @@ - Reading and querying items by RID now works end-to-end, including a parent-database cross-check that rejects a container RID belonging to a different database. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687)) - Added cross-partition `DISTINCT` query support. `SELECT DISTINCT` now deduplicates structurally equal values across every physical partition and page, rather than failing as an unsupported query feature. A `DISTINCT` query whose `ORDER BY` prefixes its projection (for example `SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city`) is resumable from a continuation token; one without such an `ORDER BY` is not, and requesting a token for it returns an error explaining how to rewrite the query. ([#5026](https://github.com/Azure/azure-sdk-for-rust/pull/5026)) - Added an SDK-generated `x-ms-client-id` header that remains stable for each `CosmosClient`. ([#4844](https://github.com/Azure/azure-sdk-for-rust/pull/4844)) +- Added opt-in Cosmos binary JSON encoding for SQL query pages (including cross-partition `DISTINCT` and streaming `ORDER BY`) and document read feeds. Query plans and change feeds remain text. - Added opt-in Cosmos binary JSON encoding for item operations (`create`/`read`/`replace`/`upsert`). Enable it via `CosmosClientBuilder::with_binary_encoding_options` (or the `AZURE_COSMOS_BINARY_ENCODING_ENABLED` environment-variable fallback). Off by default; when disabled, requests and responses are byte-for-byte unchanged. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671)) - Added `FeedOptions::max_fan_out` (and `FeedOptions::with_max_fan_out`) to cap how many physical partitions a cross-partition query or change feed may fan out to. Applies to `ContainerClient::query_items` and `ContainerClient::query_change_feed`. The cap is enforced only at initial query setup; a partition that splits mid-execution and pushes the fan-out higher does not abort the operation. ([#4855](https://github.com/Azure/azure-sdk-for-rust/pull/4855)) - Added resumable cross-partition streaming `ORDER BY` query support. ([#4800](https://github.com/Azure/azure-sdk-for-rust/pull/4800)) diff --git a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs index c2703146036..b45695623d1 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs @@ -45,9 +45,11 @@ use azure_data_cosmos::options::{ OperationOptions, Region, ServerCertificateValidation, }; use azure_data_cosmos::{ - AccountEndpoint, AccountReference, CosmosClient, CosmosRuntime, RoutingStrategy, SubStatusCode, + AccountEndpoint, AccountReference, CosmosClient, CosmosRuntime, FeedScope, Query, + RoutingStrategy, SubStatusCode, }; use azure_data_cosmos_driver::models::ConnectionString; +use futures::TryStreamExt; use serde::{Deserialize, Serialize}; use serde_json::{Map, Number, Value}; use sha2::{Digest, Sha256}; @@ -1964,11 +1966,65 @@ where ); tokio::time::sleep(backoff).await; } + Err(e) => return Err(format!("{context}: {op_name} failed: {e}").into()), } } } +async fn query_values( + container: &ContainerClient, + sql: &str, + run_id: &str, + context: &str, +) -> Result, Box> { + let mut attempt = 0; + loop { + attempt += 1; + let query = Query::from(sql).with_parameter("@run", run_id)?; + let result = match container + .query_items(query, FeedScope::full_container(), None) + .await + { + Ok(iterator) => Box::pin(iterator.try_collect()).await, + Err(err) => Err(err), + }; + match result { + Ok(values) => return Ok(values), + Err(err) if is_transient(&err) && attempt < MAX_OP_ATTEMPTS => { + let backoff = + std::time::Duration::from_millis(200u64 * (1u64 << (attempt - 1)).min(16)); + eprintln!( + "{context}: query transient failure (attempt {attempt}/{MAX_OP_ATTEMPTS}), \ + retrying in {backoff:?}: {err}" + ); + tokio::time::sleep(backoff).await; + } + Err(err) => return Err(format!("{context}: query failed: {err}").into()), + } + } +} + +fn canonical_query_results(values: Vec, ordered: bool) -> Vec { + let mut canonical: Vec = values + .into_iter() + .map(|value| { + let value = match value { + Value::Object(mut map) => { + strip_reserved_fields(&mut map); + Value::Object(map) + } + value => value, + }; + canonicalize(&value) + }) + .collect(); + if !ordered { + canonical.sort(); + } + canonical +} + // ───────────────────────────────────────────────────────────────────────────── // The fuzzer test // ───────────────────────────────────────────────────────────────────────────── @@ -2027,7 +2083,8 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { // below — the same document stored under multiple configs would // otherwise collide on the `(pk, id)` key and fail with 409 Conflict. let base_doc = gen_object(&mut rng, &cfg); - let pk = format!("pk-{}", rng.below(16)); + let partition_bucket = rng.below(16); + let run_id = format!("fuzz-{:016x}-{iter}", cfg.seed); // Optionally print the generated document (pretty JSON) so a run can be // eyeballed. Enable with `AZURE_COSMOS_FUZZ_PRINT=true`. @@ -2050,10 +2107,12 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { // touching the document RNG stream so a rerun with the same // AZURE_COSMOS_FUZZ_SEED reproduces the exact document *and* its // canonical form (a random `Uuid` here would defeat that promise). - let id = format!("fuzz-{:016x}-{iter}-{config_idx}", cfg.seed); + let id = format!("{run_id}-{config_idx}"); + let pk = format!("{run_id}-pk-{partition_bucket}-{config_idx}"); let mut doc = base_doc.clone(); doc.insert("id".to_string(), Value::String(id.clone())); doc.insert("pk".to_string(), Value::String(pk.clone())); + doc.insert("fuzzRun".to_string(), Value::String(run_id.clone())); // Never send Cosmos-reserved system properties (`_rid`, `_self`, // `_etag`, `_ts`, `_attachments`): the service **owns** these and // overwrites/assigns them, so a random value we send would come @@ -2188,13 +2247,55 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { } } + let query_cases = [ + ( + "SELECT * FROM c WHERE c.fuzzRun = @run", + false, + "select-all", + ), + ( + "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run", + false, + "distinct", + ), + ( + "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run \ + ORDER BY c._sampler.int", + true, + "distinct-order-by", + ), + ]; + for (sql, ordered, phase) in query_cases { + let mut expected: Option> = None; + for (label, client) in &clients { + let container = client + .database_client(&database_name) + .container_client(&container_name) + .await?; + let context = format!("iter={iter} config={label} query={phase} seed={}", cfg.seed); + let actual = canonical_query_results( + query_values(&container, sql, &run_id, &context).await?, + ordered, + ); + if let Some(expected) = &expected { + assert_eq!( + &actual, expected, + "{context}: query result diverged across encoding configurations" + ); + } else { + expected = Some(actual); + } + checked += 1; + } + } + if (iter + 1) % 100 == 0 { println!("... {} iterations, {checked} round-trips OK", iter + 1); } } println!( - "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops = {checked} round-trips, all canonical-equal (seed={})", + "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops + 3 queries/config = {checked} canonical comparisons, all equal (seed={})", cfg.iterations, configs.len(), cfg.seed diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs index acee3ebb685..a37a4bca540 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -//! Live-only split coverage for cross-partition `DISTINCT`. +//! Live-only binary-encoding split coverage for cross-partition `DISTINCT`. //! //! Neither .NET nor Java tests `DISTINCT` against a real partition split //! (.NET's `FullPipelineTests.TestMerge` covers `ORDER BY` only), so this is @@ -39,7 +39,7 @@ use azure_data_cosmos::{ clients::ContainerClient, feed::FeedScope, models::{ContainerProperties, ThroughputProperties}, - options::{MaxItemCountHint, QueryOptions}, + options::{BinaryEncodingOptions, MaxItemCountHint, QueryOptions}, }; use framework::{TestClient, TestOptions}; use futures::StreamExt; @@ -124,7 +124,8 @@ where not(test_category = "split"), ignore = "requires test_category 'split'" )] -pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), Box> { +pub async fn binary_distinct_query_across_split_returns_each_value_once( +) -> Result<(), Box> { TestClient::run_with_unique_db( async |run_context, db_client| { let properties = @@ -271,7 +272,11 @@ pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), }, // A real split takes minutes; the 80s default would abort mid-poll. // Matches the other split tests in this directory. - Some(TestOptions::new().with_timeout(Duration::from_secs(40 * 60))), + Some( + TestOptions::new() + .with_timeout(Duration::from_secs(40 * 60)) + .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)), + ), ) .await } diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 1347c6adfae..b99942851a6 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Added Cosmos binary JSON negotiation for SQL query pages and document read feeds. Query plans and change feeds remain text. +- Extended the binary round trip fuzzer with plain, `DISTINCT`, and `DISTINCT` + `ORDER BY` query parity across text, binary, and binary-with-text-response modes. - Added binary round trip fuzzer. As a part of the implementation, binary JSON responses now deserialize a service-echoed integral `Double` into a signed or unsigned integer field (previously a type error); this is intentionally lossy for integers the service cannot represent exactly, while a fractional `Double` remains a type error. Does not yet cover integer elements inside a uniform `Float64` array or an enum variant. ([#4976](https://github.com/Azure/azure-sdk-for-rust/pull/4976)) - Added driver-internal resolution of containers by resource id (RID). `CosmosDriver::resolve_container_by_rid` reads a container's metadata addressing it purely by RID (deriving the parent database RID from the container RID, so no `read_database` round-trip is needed) and caches the result in a by-RID index. References are validated for consistent name/RID addressing in `plan_operation` — the single choke point every executable operation passes through, including multi-page queries — returning a deterministic `CLIENT_MIXED_NAME_RID_ADDRESSING` error before signing instead of letting the gateway reject a mixed name/RID request with an opaque `401`. The new `CLIENT_INVALID_RESOURCE_ID` and `CLIENT_MIXED_NAME_RID_ADDRESSING` client statuses carry searchable names for diagnostics. ([#4663](https://github.com/Azure/azure-sdk-for-rust/pull/4663)) - Added `models::is_database_rid`, which reports whether a RID string decodes to a database-level RID (4 bytes). Lets callers that reuse a supplied RID as a database identity reject a wrong-hierarchy RID before it addresses the wrong resource. ([#4640](https://github.com/Azure/azure-sdk-for-rust/pull/4640)) diff --git a/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md b/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md index 678862f0863..9e29b85ebe7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md +++ b/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md @@ -22,7 +22,7 @@ Because the option lives on the driver and is schema-agnostic, the driver perfor A self-contained, in-tree **end-to-end validation loop** is included via the in-memory emulator (no Docker, no live account, no external test vectors). -> **Scope:** item operations (`create` / `replace` / `upsert` / `read`). Query, patch, transactional batch, and bulk are intentionally deferred (see [Deferred work](#deferred-work)). +> **Scope:** item operations (`create` / `replace` / `upsert` / `read`), SQL query pages, and document read feeds. Change feed, patch, transactional batch, and bulk are intentionally deferred (see [Deferred work](#deferred-work)). --- @@ -55,7 +55,7 @@ FFI hosts set the equivalent flat fields on the C ABI `cosmos_operation_options_ ### Two transcodes, both in the driver (schema-agnostic) -When `binary_encoding.enabled` is set, `CosmosDriver::execute_operation` owns the wire format both ways: +When `binary_encoding.enabled` is set, `CosmosDriver::plan_operation` and `execute_plan` own the wire format both ways: * **Request** (`apply_request_binary_encoding`) — transcodes a **text** request body to Cosmos binary JSON via `binary_json::transcode_to_binary` (`serde_json::from_slice` → `encode`) and advertises `JsonText,CosmosBinary`. An **already-binary** or empty body passes through unchanged, so a caller that pre-encoded pays nothing. * **Response** (when `request_text_response` is set) — transcodes the binary response back to text JSON via `binary_json::transcode_to_text` (`decode` → `serde_json::to_vec`). The wire stays binary in both directions. @@ -68,7 +68,7 @@ The Rust SDK keeps a typed fast path: `serialize_item_body` encodes `T: Serializ ### Negotiation header -When binary is enabled, item operations set: +When binary is enabled, eligible item, query, and document read-feed operations set: ``` x-ms-cosmos-supported-serialization-formats: JsonText,CosmosBinary @@ -76,6 +76,8 @@ x-ms-cosmos-supported-serialization-formats: JsonText,CosmosBinary The value matches the .NET reference (`string.Join(",", JsonText, CosmosBinary)` — no space). The request `Content-Type` stays `application/json`; the service detects the binary body from its first byte. The **driver** sets this header whenever `binary_encoding.enabled` is set — including under `request_text_response`, where the wire stays binary and the driver transcodes the response (see above). +The `/queryplan` request is the deliberate exception: it remains text and does not advertise binary because that endpoint ignores binary negotiation. Query page bodies are normalized from binary to text once at pipeline ingest so DISTINCT and streaming ORDER BY continue to operate on their existing text representation. Rebuilt pages are restored to the negotiated format at the plan boundary. `ResponseBody::Items` encodes each item independently, including its own `0x80` preamble, so every slice remains a standalone binary document. + --- ## Flow diagrams @@ -394,9 +396,8 @@ opts.binary_encoding_request_text_response = 2; /* 2 = true */ ## Deferred work -* **Query binary negotiation** — the request-body encoding + negotiation for query pages is still deferred. Binary encoding now lives on `OperationOptions.binary_encoding`, so a driver-minted page operation *can* carry it; what remains is confirming query-body semantics (`application/query+json`) and the native cross-partition query engine's handling of binary item bytes. The query *response* decode already works via the shared choke point. -* **Binary feed responses** — the feed splitter scans **text** JSON, so binary `Documents` envelopes cannot be sliced yet; making it binary-aware is a prerequisite for any feed/query binary negotiation. * **`patch`** — excluded from binary encoding for now (the driver's request-side encode intentionally skips patch); transactional `batch` / `bulk` are deferred by spec. +* **Change feed** — remains text because its continuation and resume semantics differ from ordinary document read feeds. * **Cross-implementation vectors** — validate against captured real .NET / Java binary output. --- diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index a5d82935dcd..fe99a004fed 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2601,28 +2601,6 @@ impl CosmosDriver { .await; } - // Resolve binary encoding through the same layered view as every other - // option, and only honor it for point **item** operations (the resource - // must be a `Document`; query/feed/batch and every control-plane - // resource are deferred per the binary-encoding spec). - let binary = - if Self::binary_encoding_applies(operation.resource_type(), operation.operation_type()) - { - self.operation_options_view(&options) - .binary_encoding() - .cloned() - .unwrap_or_default() - } else { - crate::options::BinaryEncodingOptions::default() - }; - let operation = if binary.enabled { - Self::apply_request_binary_encoding(operation)? - } else { - operation - }; - - let transcode_response_to_text = binary.enabled && binary.request_text_response; - // TODO: This boxing is a temporary fix to avoid a large future. // We need to do some refactoring here to shrink the future size and avoid this heap allocation if possible. let response = Box::pin(async { @@ -2634,29 +2612,16 @@ impl CosmosDriver { }) .await?; - // Driver-side transcoding: convert the binary response body to text - // when the caller asked for a text payload over a binary wire. - if transcode_response_to_text { - if let Some(mut response) = response { - response.transcode_body_to_text()?; - return Ok(Some(response)); - } - } Ok(response) } /// Whether binary encoding applies to an operation. /// - /// Honored only for point item operations: the resource must be a - /// [`ResourceType::Document`] and the operation one of create/read/replace/ - /// upsert. Control-plane resources share those operation types but must - /// never be binary encoded (some carry JSON bodies). - fn binary_encoding_applies( - resource_type: crate::models::ResourceType, - operation_type: crate::models::OperationType, - ) -> bool { - resource_type == crate::models::ResourceType::Document - && operation_type.supports_binary_encoding() + /// Honored only for document operations with JSON bodies or feed pages. + fn binary_encoding_applies(operation: &CosmosOperation) -> bool { + operation.resource_type() == crate::models::ResourceType::Document + && operation.operation_type().supports_binary_encoding() + && !operation.is_change_feed() } /// Applies request-side binary encoding to an operation: transcodes a text @@ -2694,6 +2659,16 @@ impl CosmosDriver { Ok(operation.with_supported_serialization_formats(BINARY_NEGOTIATION_FORMATS)) } + fn query_plan_request_body(operation: &CosmosOperation) -> crate::error::Result> { + crate::binary_json::transcode_to_text(operation.body().unwrap_or_default()).map_err(|e| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERIALIZATION_REQUEST_BODY_INVALID) + .with_message("failed to prepare text query-plan request body") + .with_source(e) + .build() + }) + } + /// Executes a singleton operation (operations which return only a single result). /// /// This is a convenience method around [`execute_operation`](CosmosDriver::execute_operation) that asserts at debug-time that the operation @@ -2940,7 +2915,22 @@ impl CosmosDriver { topology.as_mut().map(|t| t as &mut dyn TopologyProvider), ); - plan.pipeline.next_page(&mut context).await + let binary_enabled = plan.binary_encoding().enabled; + let request_text_response = plan.binary_encoding().request_text_response; + let mut response = plan.pipeline.next_page(&mut context).await?; + if binary_enabled { + if let Some(response) = response.as_mut() { + if request_text_response { + response.transcode_body_to_text()?; + } else if matches!( + plan.operation().operation_type(), + crate::models::OperationType::Query | crate::models::OperationType::ReadFeed + ) { + response.transcode_body_to_binary()?; + } + } + } + Ok(response) } async fn execute_operation_direct( @@ -3238,6 +3228,19 @@ impl CosmosDriver { // reference, so it issues no additional network calls and does not change // the request flow. operation.validate_addressing()?; + let binary = if Self::binary_encoding_applies(&operation) { + self.operation_options_view(options) + .binary_encoding() + .cloned() + .unwrap_or_default() + } else { + crate::options::BinaryEncodingOptions::default() + }; + let operation = if binary.enabled { + Self::apply_request_binary_encoding(operation)? + } else { + operation + }; // Planning holds the whole pipeline-builder state across several await // points, which makes it one of the largest futures in the driver — @@ -3245,7 +3248,11 @@ impl CosmosDriver { // here so every caller awaits a pointer-sized future instead of having // to pin at its own call site and rediscover this each time the state // grows. - Box::pin(self.plan_operation_inner(operation, options, continuation, plan_options)).await + let mut plan = + Box::pin(self.plan_operation_inner(operation, options, continuation, plan_options)) + .await?; + plan.set_binary_encoding(binary); + Ok(plan) } async fn plan_operation_inner( @@ -3409,7 +3416,7 @@ impl CosmosDriver { container.clone(), std::borrow::Cow::Borrowed(crate::query::SUPPORTED_QUERY_FEATURES), ) - .with_body(operation.body().unwrap_or_default().to_vec()); + .with_body(Self::query_plan_request_body(operation)?); let response = self .execute_operation_direct( @@ -3463,9 +3470,9 @@ impl CosmosDriver { .map(|props| props.query_engine_configuration.clone()) .unwrap_or_default(); - let native_result = operation - .body() - .and_then(|b| std::str::from_utf8(b).ok()) + let query_plan_body = Self::query_plan_request_body(operation)?; + let native_result = std::str::from_utf8(&query_plan_body) + .ok() .map(|json| self.try_native_query_plan(json, container, &query_engine_config)); match native_result { @@ -6272,28 +6279,38 @@ mod tests { fn binary_encoding_applies_only_to_document_item_ops() { use crate::models::{OperationType, ResourceType}; - // Point item ops on `Document` are the only combinations that qualify. + let container = epk_test_container(r#"{"paths":["/pk"],"version":2}"#); for op in [ OperationType::Create, OperationType::Read, OperationType::Replace, OperationType::Upsert, + OperationType::Query, + OperationType::ReadFeed, ] { + let operation = CosmosOperation::new( + op, + crate::models::CosmosResourceReference::from(container.clone()) + .with_resource_type(ResourceType::Document) + .into_feed_reference(), + Some(crate::models::FeedRange::full()), + ); assert!( - CosmosDriver::binary_encoding_applies(ResourceType::Document, op), + CosmosDriver::binary_encoding_applies(&operation), "Document + {op:?} should be binary-encodable", ); } - // Non-item operation types on `Document` are excluded (query/feed/delete/patch). - for op in [ - OperationType::Delete, - OperationType::Query, - OperationType::ReadFeed, - OperationType::Patch, - ] { + for op in [OperationType::Delete, OperationType::Patch] { + let operation = CosmosOperation::new( + op, + crate::models::CosmosResourceReference::from(container.clone()) + .with_resource_type(ResourceType::Document) + .into_feed_reference(), + Some(crate::models::FeedRange::full()), + ); assert!( - !CosmosDriver::binary_encoding_applies(ResourceType::Document, op), + !CosmosDriver::binary_encoding_applies(&operation), "Document + {op:?} must not be binary-encoded", ); } @@ -6314,12 +6331,23 @@ mod tests { OperationType::Replace, OperationType::Upsert, ] { + let operation = CosmosOperation::new( + op, + crate::models::CosmosResourceReference::from(container.clone()) + .with_resource_type(rt) + .into_feed_reference(), + None, + ); assert!( - !CosmosDriver::binary_encoding_applies(rt, op), + !CosmosDriver::binary_encoding_applies(&operation), "{rt:?} + {op:?} must not be binary-encoded (control plane)", ); } } + + let change_feed = + CosmosOperation::change_feed(container, Some(crate::models::FeedRange::full())); + assert!(!CosmosDriver::binary_encoding_applies(&change_feed)); } fn binary_encoding_test_operation(body: Vec) -> CosmosOperation { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs index 7ec6322fb77..3e634aa02be 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs @@ -115,6 +115,8 @@ mod seeds { pub(super) const FALSE: u128 = seed(0xc1be517fe893b40c, 0xe9fc8a4c531cd0dd); pub(super) const TRUE: u128 = seed(0xf86d4abf9a412e74, 0x788488365c8a985d); pub(super) const STRING: u128 = seed(0x61f53f0a44204cfb, 0x09481be8ef4b56dd); + /// A single number seed keeps DISTINCT encoding-independent: text integers + /// and integral binary doubles flow through the same canonical number bytes. pub(super) const NUMBER: u128 = seed(0x2400e8b894ce9c2a, 0x790be1eabd7b9481); pub(super) const ARRAY: u128 = seed(0xfa573b014c4dc18e, 0xa014512c858eb115); pub(super) const OBJECT: u128 = seed(0x77b285ac511aef30, 0x3dcf187245822449); @@ -337,6 +339,17 @@ mod tests { assert_eq!(h(json!(-7)), h(json!(-7.0))); } + #[test] + fn text_integer_and_binary_double_have_the_same_hash() { + let text: Value = serde_json::from_slice(b"1").unwrap(); + let binary = crate::binary_json::encode(&json!(1.0)); + let binary: Value = crate::binary_json::from_slice(&binary).unwrap(); + + assert!(text.as_u64().is_some()); + assert!(binary.as_f64().is_some()); + assert_eq!(h(text), h(binary)); + } + /// .NET normalizes `-0.0` to `0.0` in `CosmosNumberHasher`, but neither /// peer has a test for it. #[test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs index c582c2d13be..110e95022bd 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use crate::models::{ContinuationToken, CosmosOperation, CosmosResponse}; +use crate::options::BinaryEncodingOptions; use super::context::PipelineContext; use super::node::{PageResult, PipelineNode}; @@ -98,6 +99,7 @@ impl Pipeline { pub struct OperationPlan { pub(crate) pipeline: Pipeline, operation: Arc, + binary_encoding: BinaryEncodingOptions, } impl OperationPlan { @@ -106,9 +108,22 @@ impl OperationPlan { Self { pipeline, operation, + binary_encoding: BinaryEncodingOptions::default(), } } + pub(crate) fn operation(&self) -> &CosmosOperation { + &self.operation + } + + pub(crate) fn binary_encoding(&self) -> &BinaryEncodingOptions { + &self.binary_encoding + } + + pub(crate) fn set_binary_encoding(&mut self, binary_encoding: BinaryEncodingOptions) { + self.binary_encoding = binary_encoding; + } + /// Snapshots this plan into a [`ContinuationToken`] suitable for cross-process /// resumption. /// diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs index 114687846e1..bb296ce1227 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs @@ -202,8 +202,11 @@ pub(crate) fn rewrite_query_body( "query".to_owned(), serde_json::Value::String(rewritten_query.to_owned()), ); - serde_json::to_vec(&value) - .map_err(|e| body_error("failed to serialize rewritten query body", e)) + serialize_query_body( + original_body, + &value, + "failed to serialize rewritten query body", + ) } /// Inserts the .NET-compatible structured `"resumeFilter"` field into an @@ -236,8 +239,11 @@ pub(crate) fn with_resume_filter( "resumeFilter".to_owned(), order_by::resume_filter_json(resume_values, rid, exclude), ); - serde_json::to_vec(&value) - .map_err(|e| body_error("failed to serialize query body with resume filter", e)) + serialize_query_body( + body, + &value, + "failed to serialize query body with resume filter", + ) } /// Parses a query operation's JSON body, erroring on a missing body or @@ -246,10 +252,27 @@ fn parse_query_body(body: Option<&[u8]>) -> crate::error::Result, + value: &serde_json::Value, + error_message: &'static str, +) -> crate::error::Result> { + let text = serde_json::to_vec(value).map_err(|e| body_error(error_message, e))?; + if original_body.is_some_and(crate::binary_json::is_binary) { + crate::binary_json::transcode_to_binary(&text) + .map_err(|e| binary_body_error("failed to transcode rewritten query body to binary", e)) + } else { + Ok(text) + } +} + /// One row parsed from a rewritten-envelope backend page, ready for the /// merge heap. `payload` retains the item's exact original JSON bytes /// (via [`RawValue`]) rather than a re-serialized value, so the emitted @@ -298,7 +321,7 @@ pub(crate) fn parse_envelope_page( ) -> crate::error::Result> { let bytes = match body { ResponseBody::NoPayload => return Ok(Vec::new()), - ResponseBody::Bytes(b) => b.clone(), + ResponseBody::Bytes(b) => normalize_page_body(b)?, ResponseBody::Items(_) => { return Err(envelope_error( "rewritten-query backend page returned an already-split `Items` body; \ @@ -536,9 +559,12 @@ pub(crate) fn retain_documents(documents: &[Box], keep: &[usize]) -> ( pub(crate) fn parse_document_page(body: &ResponseBody) -> crate::error::Result>> { match body { ResponseBody::NoPayload => Ok(Vec::new()), - ResponseBody::Bytes(bytes) => serde_json::from_slice::(bytes) - .map(|feed| feed.documents) - .map_err(|e| body_error("failed to parse backend page as a feed body", e)), + ResponseBody::Bytes(bytes) => { + let bytes = normalize_page_body(bytes)?; + serde_json::from_slice::(&bytes) + .map(|feed| feed.documents) + .map_err(|e| body_error("failed to parse backend page as a feed body", e)) + } ResponseBody::Items(_) => Err(envelope_error( "backend page returned an already-split `Items` body; expected a raw \ `Documents`-array feed body", @@ -546,6 +572,21 @@ pub(crate) fn parse_document_page(body: &ResponseBody) -> crate::error::Result crate::error::Result { + if !crate::binary_json::is_binary(bytes) { + return Ok(bytes.clone()); + } + crate::binary_json::transcode_to_text(bytes) + .map(bytes::Bytes::from) + .map_err(|source| { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) + .with_message("failed to transcode binary query page to text JSON") + .with_source(source) + .build() + }) +} + /// A fresh, empty [`DiagnosticsContext`] for a page needing no new backend /// fetch. Uses a new activity ID since it corresponds to no real request. fn empty_diagnostics() -> Arc { @@ -572,6 +613,17 @@ fn body_error(message: &'static str, source: serde_json::Error) -> crate::error: .build() } +fn binary_body_error( + message: &'static str, + source: crate::binary_json::BinaryError, +) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) + .with_message(message) + .with_source(source) + .build() +} + fn body_error_msg(message: &'static str) -> crate::error::CosmosError { crate::error::CosmosError::builder() .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) @@ -802,6 +854,52 @@ mod tests { assert_eq!(rows[1].keys, vec![OrderByItem::Undefined]); } + #[test] + fn parse_envelope_page_accepts_binary_feed_body() { + let value = serde_json::json!({ + "_rid": "abc", + "Documents": [{ + "_rid": "r1", + "orderByItems": [{"item": 1.0}], + "payload": {"id": "d1"} + }], + "_count": 1 + }); + let body = ResponseBody::from_bytes(crate::binary_json::encode(&value)); + let rows = parse_envelope_page(&body, 1).unwrap(); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].rid, "r1"); + assert_eq!(rows[0].keys, vec![OrderByItem::Number(1.0.into())]); + assert_eq!(rows[0].payload.get(), r#"{"id":"d1"}"#); + } + + #[test] + fn parse_document_page_accepts_binary_feed_body() { + let value = serde_json::json!({ + "_rid": "abc", + "Documents": [{"id": "d1"}, {"id": "d2"}], + "_count": 2 + }); + let body = ResponseBody::from_bytes(crate::binary_json::encode(&value)); + let documents = parse_document_page(&body).unwrap(); + + assert_eq!(documents.len(), 2); + assert_eq!(documents[0].get(), r#"{"id":"d1"}"#); + assert_eq!(documents[1].get(), r#"{"id":"d2"}"#); + } + + #[test] + fn malformed_binary_page_is_a_response_serialization_error() { + let body = ResponseBody::from_bytes(vec![crate::binary_json::PREAMBLE]); + let err = parse_document_page(&body).unwrap_err(); + + assert_eq!( + err.status(), + CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID + ); + } + #[test] fn parse_envelope_page_empty_body_yields_no_rows() { let rows = parse_envelope_page(&ResponseBody::NoPayload, 1).unwrap(); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs index 8e0a648628a..66cd3907456 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs @@ -2329,6 +2329,7 @@ fn success_feed_response( items: Vec, page_options: FeedPageOptions<'_>, feed_headers: FeedResponseHeaders, + binary: bool, start: Instant, ) -> AsyncRawResponse { let (page, next) = match paginate_values( @@ -2342,9 +2343,10 @@ fn success_feed_response( }; let item_count = page.len() as u32; let body = feed_to_json(envelope_name, page, rid); - let mut builder = success_response( + let mut builder = success_response_with_format( StatusCode::Ok, &body, + binary, 1.0, &feed_headers.session_token, start, @@ -2371,6 +2373,7 @@ fn success_document_feed_response( items: Vec, page_options: FeedPageOptions<'_>, feed_headers: FeedResponseHeaders, + binary: bool, start: Instant, ) -> AsyncRawResponse { let (page, next) = match paginate_document_feed_items( @@ -2384,9 +2387,10 @@ fn success_document_feed_response( }; let item_count = page.len() as u32; let body = feed_to_json(envelope_name, page, rid); - let mut builder = success_response( + let mut builder = success_response_with_format( StatusCode::Ok, &body, + binary, 1.0, &feed_headers.session_token, start, @@ -2424,7 +2428,13 @@ fn parse_query_spec( request_body: &[u8], start: Instant, ) -> Result<(String, Vec<(String, serde_json::Value)>), AsyncRawResponse> { - let spec: QuerySpec = serde_json::from_slice(request_body).map_err(|e| { + let spec: QuerySpec = if crate::binary_json::is_binary(request_body) { + crate::binary_json::from_slice(request_body) + .map_err(|e| format!("invalid binary query body: {e}")) + } else { + serde_json::from_slice(request_body).map_err(|e| format!("invalid text query body: {e}")) + } + .map_err(|e| { error_response( StatusCode::BadRequest, None, @@ -2490,6 +2500,7 @@ fn execute_query_feed( results, FeedPageOptions::from_request(parsed), feed_headers, + parsed.binary_response, start, ) } @@ -2514,6 +2525,7 @@ fn execute_document_query_feed( results, FeedPageOptions::from_request(parsed), feed_headers, + parsed.binary_response, start, ), Ok(None) => { @@ -2539,6 +2551,7 @@ fn execute_document_query_feed( results, FeedPageOptions::from_request(parsed), feed_headers, + parsed.binary_response, start, ) } @@ -2717,6 +2730,7 @@ fn handle_read_feed_databases( databases, FeedPageOptions::from_request(parsed), FeedResponseHeaders::none(), + false, start, ) } @@ -2782,6 +2796,7 @@ fn handle_read_feed_containers( containers, FeedPageOptions::from_request(parsed), FeedResponseHeaders::none(), + false, start, ) } @@ -2843,6 +2858,7 @@ fn handle_read_feed_offers( offers, FeedPageOptions::from_request(parsed), FeedResponseHeaders::none(), + false, start, ) } @@ -3184,6 +3200,7 @@ fn handle_read_feed_items( docs, FeedPageOptions::from_request(parsed), headers, + parsed.binary_response && parsed.a_im.is_none(), start, ) } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_response.rs index 0caa1154462..c2bc26e1169 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_response.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_response.rs @@ -53,6 +53,12 @@ impl CosmosResponsePayload { self.body = body.transcode_to_text()?; Ok(()) } + + fn transcode_body_to_binary(&mut self) -> crate::error::Result<()> { + let body = std::mem::take(&mut self.body); + self.body = body.transcode_to_binary()?; + Ok(()) + } } /// Result of a Cosmos DB operation. /// @@ -153,6 +159,10 @@ impl CosmosResponse { self.payload.transcode_body_to_text() } + pub(crate) fn transcode_body_to_binary(&mut self) -> crate::error::Result<()> { + self.payload.transcode_body_to_binary() + } + /// Returns a reference to the extracted headers. pub fn headers(&self) -> &CosmosResponseHeaders { self.payload.headers() diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs index 60c0777cc12..634ee5f6c8e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs @@ -638,10 +638,9 @@ impl OperationType { ) } - /// True for the point item ops (create/read/replace/upsert) eligible for - /// binary encoding. Necessary but not sufficient: the full gate also - /// requires [`ResourceType::Document`] (see - /// `CosmosDriver::binary_encoding_applies`). + /// True for item and document-feed operations eligible for binary encoding. + /// Necessary but not sufficient: the full gate also requires + /// [`ResourceType::Document`] and excludes change feed. pub(crate) fn supports_binary_encoding(self) -> bool { matches!( self, @@ -649,6 +648,8 @@ impl OperationType { | OperationType::Read | OperationType::Replace | OperationType::Upsert + | OperationType::Query + | OperationType::ReadFeed ) } @@ -894,23 +895,20 @@ mod tests { use serde::{Deserialize, Serialize}; #[test] - fn supports_binary_encoding_covers_only_bodied_point_ops() { - // Matches the binary-encoding spec §2 scope table: create/read/replace/ - // upsert. `delete` is excluded (no request or response body); query, - // feed, batch, and stored-procedure paths are deferred. + fn supports_binary_encoding_covers_item_and_feed_ops() { for op in [ OperationType::Create, OperationType::Read, OperationType::Replace, OperationType::Upsert, + OperationType::Query, + OperationType::ReadFeed, ] { assert!(op.supports_binary_encoding(), "{op:?} should be supported"); } for op in [ OperationType::Delete, - OperationType::Query, OperationType::SqlQuery, - OperationType::ReadFeed, OperationType::Batch, OperationType::Execute, OperationType::Patch, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs index b75648d9a95..9ad6364f060 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs @@ -145,19 +145,6 @@ impl ResponseBody { } Self::Items(items) => items .into_iter() - // NOTE: `deserialize_response` auto-detects binary per slice via - // the `0x80` preamble, but the feed pipeline that produces - // `Self::Items` splits the `Documents` array by scanning **text** - // JSON — it is not binary-aware. A single-preamble binary feed - // envelope would therefore be sliced into sub-documents *without* - // preambles, which `is_binary` would then route to the text path. - // This is inert today because query/feed binary negotiation is - // deferred (the service does not emit binary feeds without the - // negotiation header), so the binary `Items` branch is only - // exercised by hand-prefixed synthetic tests. When feed/query - // binary negotiation is added, the feed splitter must be made - // binary-aware (or each slice re-prefixed) before this path can - // decode real binary feeds. .map(|b| deserialize_response(&b, "failed to deserialize feed item")) .collect(), } @@ -206,6 +193,31 @@ impl ResponseBody { } } } + + /// Transcodes text JSON payloads to standalone Cosmos binary JSON buffers. + /// + /// Each [`Items`](Self::Items) slice is encoded independently, so every + /// resulting item carries the `0x80` preamble required for auto-detection. + pub(crate) fn transcode_to_binary(self) -> crate::error::Result { + fn convert(bytes: &Bytes) -> crate::error::Result { + if crate::binary_json::is_binary(bytes) { + return Ok(bytes.clone()); + } + crate::binary_json::transcode_to_binary(bytes) + .map(Bytes::from) + .map_err(|e| invalid_body_error("failed to transcode response body to binary", e)) + } + + match self { + Self::NoPayload => Ok(Self::NoPayload), + Self::Bytes(bytes) => convert(&bytes).map(Self::Bytes), + Self::Items(items) => items + .iter() + .map(convert) + .collect::>>() + .map(Self::Items), + } + } } impl From for ResponseBody { @@ -287,6 +299,28 @@ mod tests { assert!(items.is_empty()); } + #[test] + fn transcode_items_to_binary_prefixes_every_item() { + let body = ResponseBody::from_items(vec![ + Bytes::from_static(br#"{"id":"1"}"#), + Bytes::from_static(br#"{"id":"2"}"#), + ]); + let ResponseBody::Items(items) = body.transcode_to_binary().unwrap() else { + panic!("expected item body"); + }; + + assert_eq!(items.len(), 2); + assert!(items.iter().all(|item| crate::binary_json::is_binary(item))); + let decoded: Vec = ResponseBody::from_items(items).into_items().unwrap(); + assert_eq!( + decoded, + vec![ + serde_json::json!({"id": "1"}), + serde_json::json!({"id": "2"}) + ] + ); + } + #[test] fn no_payload_into_item_errors() { // No bytes to deserialize. diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index fbc592fc800..124730e2fde 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -26,7 +26,9 @@ use azure_data_cosmos_driver::models::{ ContainerReference, CosmosOperation, FeedRange, ItemReference, MaxItemCountHint, PartitionKey, PartitionKeyDefinition, }; -use azure_data_cosmos_driver::options::{DriverOptions, OperationOptions, PlanOptions}; +use azure_data_cosmos_driver::options::{ + BinaryEncodingOptions, DriverOptions, OperationOptions, OperationOptionsBuilder, PlanOptions, +}; const GATEWAY_URL: &str = "https://eastus.emulator.local"; @@ -168,10 +170,50 @@ fn documents_of( let Ok(bytes) = response.into_body().single() else { return Vec::new(); }; - let value: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + let value: serde_json::Value = if azure_data_cosmos_driver::binary_json::is_binary(&bytes) { + azure_data_cosmos_driver::binary_json::from_slice(&bytes).unwrap() + } else { + serde_json::from_slice(&bytes).unwrap() + }; value["Documents"].as_array().cloned().unwrap_or_default() } +async fn drain_query_with_options( + driver: &CosmosDriver, + container: &ContainerReference, + query: &QuerySpec, + planning_options: OperationOptions, + execution_options: OperationOptions, +) -> (Vec, Vec) { + let mut plan = Box::pin(driver.plan_operation( + query_operation(container, query, 1), + &planning_options, + None, + &PlanOptions::default(), + )) + .await + .unwrap(); + let mut values = Vec::new(); + let mut formats = Vec::new(); + + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + execution_options.clone(), + ) + .await + .unwrap() + { + let bytes = response.body().clone().single().unwrap(); + let is_binary = azure_data_cosmos_driver::binary_json::is_binary(&bytes); + formats.push(is_binary); + values.extend(documents_of(response)); + } + + (values, formats) +} + /// Sorts values by their serialized form so an unordered result can be /// compared deterministically. fn sorted(mut values: Vec) -> Vec { @@ -524,3 +566,125 @@ async fn split_mid_drain_does_not_reemit_deduplicated_values() { assert_matches_expected(scenario, all); } } + +#[tokio::test] +async fn text_and_binary_query_pages_have_pipeline_parity() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .unwrap(); + let documents: Vec = [ + ("a", "pk1", 1), + ("b", "pk2", 1), + ("c", "pk3", 2), + ("d", "pk4", 2), + ("e", "pk5", 3), + ] + .into_iter() + .map(|(id, pk, value)| serde_json::json!({"id": id, "pk": pk, "value": value})) + .collect(); + seed(&driver, &container, &documents).await; + + for query in [ + QuerySpec { + text: "SELECT * FROM c".to_owned(), + parameters: Vec::new(), + distinct_type: "None".to_owned(), + }, + QuerySpec { + text: "SELECT DISTINCT VALUE c.value FROM c".to_owned(), + parameters: Vec::new(), + distinct_type: "Unordered".to_owned(), + }, + QuerySpec { + text: "SELECT DISTINCT VALUE c.value FROM c ORDER BY c.value".to_owned(), + parameters: Vec::new(), + distinct_type: "Ordered".to_owned(), + }, + ] { + let (text, text_formats) = drain_query_with_options( + &driver, + &container, + &query, + OperationOptions::default(), + OperationOptions::default(), + ) + .await; + let binary_options = OperationOptionsBuilder::new() + .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)) + .build(); + let (binary, binary_formats) = drain_query_with_options( + &driver, + &container, + &query, + binary_options.clone(), + binary_options.clone(), + ) + .await; + let binary_as_text_options = OperationOptionsBuilder::new() + .with_binary_encoding( + BinaryEncodingOptions::new() + .with_enabled(true) + .with_request_text_response(true), + ) + .build(); + let (binary_as_text, binary_as_text_formats) = drain_query_with_options( + &driver, + &container, + &query, + binary_as_text_options.clone(), + binary_as_text_options, + ) + .await; + let (_, planned_binary_formats) = drain_query_with_options( + &driver, + &container, + &query, + binary_options.clone(), + OperationOptions::default(), + ) + .await; + let (_, planned_text_formats) = drain_query_with_options( + &driver, + &container, + &query, + OperationOptions::default(), + binary_options, + ) + .await; + + if query.distinct_type == "Ordered" { + assert_eq!(binary, text, "{}", query.text); + assert_eq!(binary_as_text, text, "{}", query.text); + } else { + assert_eq!(sorted(binary), sorted(text.clone()), "{}", query.text); + assert_eq!(sorted(binary_as_text), sorted(text), "{}", query.text); + } + assert!( + text_formats.iter().all(|is_binary| !is_binary), + "{}", + query.text + ); + assert!( + binary_formats.iter().all(|is_binary| *is_binary), + "{}", + query.text + ); + assert!( + binary_as_text_formats.iter().all(|is_binary| !is_binary), + "{}", + query.text + ); + assert!( + planned_binary_formats.iter().all(|is_binary| *is_binary), + "planning must own response encoding even when execution options differ: {}", + query.text + ); + assert!( + planned_text_formats.iter().all(|is_binary| !is_binary), + "execution options must not enable binary after text planning: {}", + query.text + ); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/read_feed.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/read_feed.rs index 1670c33bfa5..fa3c153ea23 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/read_feed.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/read_feed.rs @@ -99,6 +99,40 @@ async fn read_document_feed_paginates_documents() { assert_eq!(documents[0]["id"], "item3"); } +#[tokio::test] +async fn read_document_feed_honors_binary_negotiation() { + let ctx = setup_single_region().await; + let body = serde_json::json!({"id": "item1", "pk": "pk1"}); + let response = ctx + .emulator + .execute_request(&create_item_request( + &ctx.gateway_url, + "testdb", + "testcoll", + &body, + r#"["pk1"]"#, + false, + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::Created); + + let url = format!("{}/dbs/testdb/colls/testcoll/docs", ctx.gateway_url); + let mut request = Request::new(Url::parse(&url).unwrap(), Method::Get); + request.headers_mut().insert( + SUPPORTED_SERIALIZATION_FORMATS.clone(), + HeaderValue::from_static("JsonText,CosmosBinary"), + ); + let response = ctx.emulator.execute_request(&request).await.unwrap(); + let (status, _, body) = collect_raw_response(response).await; + + assert_eq!(status, StatusCode::Ok); + assert!(azure_data_cosmos_driver::binary_json::is_binary(&body)); + let decoded: serde_json::Value = + azure_data_cosmos_driver::binary_json::from_slice(&body).unwrap(); + assert_eq!(decoded["Documents"][0]["id"], "item1"); +} + #[tokio::test] async fn read_document_feed_rejects_invalid_continuation() { let ctx = setup_single_region().await; From d02eac37c57a5d3a55274f39d2b92c08753fef18 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 09:11:29 -0700 Subject: [PATCH 2/8] Fix binary query token portability Capture streaming ORDER BY fingerprints before request encoding so text and binary requests share continuation tokens without invalidating existing text tokens. Add SqlQuery negotiation consistency, quiet empty query-plan handling, service-style Float64 DISTINCT coverage, and emulator assertions for request negotiation and cross-mode resume. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../src/driver/cosmos_driver.rs | 68 ++++-- .../src/driver/dataflow/distinct.rs | 40 ++++ .../src/driver/dataflow/planner.rs | 65 +++-- .../dataflow/streaming_ordered_merge.rs | 69 ++++-- .../src/models/mod.rs | 3 +- .../in_memory_emulator_tests/distinct.rs | 223 +++++++++++++++++- 6 files changed, 407 insertions(+), 61 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index fe99a004fed..2a5f64769fa 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2659,14 +2659,21 @@ impl CosmosDriver { Ok(operation.with_supported_serialization_formats(BINARY_NEGOTIATION_FORMATS)) } - fn query_plan_request_body(operation: &CosmosOperation) -> crate::error::Result> { - crate::binary_json::transcode_to_text(operation.body().unwrap_or_default()).map_err(|e| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::SERIALIZATION_REQUEST_BODY_INVALID) - .with_message("failed to prepare text query-plan request body") - .with_source(e) - .build() - }) + fn query_plan_request_body( + operation: &CosmosOperation, + ) -> crate::error::Result>> { + let Some(body) = operation.body().filter(|body| !body.is_empty()) else { + return Ok(None); + }; + crate::binary_json::transcode_to_text(body) + .map(Some) + .map_err(|e| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERIALIZATION_REQUEST_BODY_INVALID) + .with_message("failed to prepare text query-plan request body") + .with_source(e) + .build() + }) } /// Executes a singleton operation (operations which return only a single result). @@ -3228,6 +3235,12 @@ impl CosmosDriver { // reference, so it issues no additional network calls and does not change // the request flow. operation.validate_addressing()?; + let query_fingerprint = matches!( + operation.operation_type(), + crate::models::OperationType::Query | crate::models::OperationType::SqlQuery + ) + .then(|| planner::streaming_query_fingerprint(&operation)) + .transpose()?; let binary = if Self::binary_encoding_applies(&operation) { self.operation_options_view(options) .binary_encoding() @@ -3248,9 +3261,14 @@ impl CosmosDriver { // here so every caller awaits a pointer-sized future instead of having // to pin at its own call site and rediscover this each time the state // grows. - let mut plan = - Box::pin(self.plan_operation_inner(operation, options, continuation, plan_options)) - .await?; + let mut plan = Box::pin(self.plan_operation_inner( + operation, + options, + continuation, + plan_options, + query_fingerprint, + )) + .await?; plan.set_binary_encoding(binary); Ok(plan) } @@ -3261,6 +3279,7 @@ impl CosmosDriver { options: &OperationOptions, continuation: Option<&ContinuationToken>, plan_options: &PlanOptions, + query_fingerprint: Option, ) -> crate::error::Result { if !self.initialized.load(Ordering::Acquire) { let endpoint = AccountEndpoint::from(self.options.account()); @@ -3385,11 +3404,14 @@ impl CosmosDriver { .as_ref() .is_some_and(planner::is_streaming_order_by) { - let pipeline = planner::build_streaming_ordered_merge( + let query_fingerprint = query_fingerprint + .expect("query operations capture a fingerprint before request encoding"); + let pipeline = planner::build_streaming_ordered_merge_with_fingerprint( &query_plan, &mut topology, &operation, resume_state, + query_fingerprint, ) .await?; return planner::finalize_plan(pipeline, operation, is_fresh, plan_options); @@ -3416,7 +3438,7 @@ impl CosmosDriver { container.clone(), std::borrow::Cow::Borrowed(crate::query::SUPPORTED_QUERY_FEATURES), ) - .with_body(Self::query_plan_request_body(operation)?); + .with_body(Self::query_plan_request_body(operation)?.unwrap_or_default()); let response = self .execute_operation_direct( @@ -3471,8 +3493,9 @@ impl CosmosDriver { .unwrap_or_default(); let query_plan_body = Self::query_plan_request_body(operation)?; - let native_result = std::str::from_utf8(&query_plan_body) - .ok() + let native_result = query_plan_body + .as_deref() + .and_then(|body| std::str::from_utf8(body).ok()) .map(|json| self.try_native_query_plan(json, container, &query_engine_config)); match native_result { @@ -3853,6 +3876,21 @@ mod tests { ) } + #[test] + fn empty_query_body_skips_native_query_plan_input() { + let resource = crate::models::CosmosResourceReference::from(test_account()); + let missing = + CosmosOperation::new(crate::models::OperationType::Query, resource.clone(), None); + let empty = CosmosOperation::new(crate::models::OperationType::Query, resource, None) + .with_body(Vec::new()); + + assert_eq!( + CosmosDriver::query_plan_request_body(&missing).unwrap(), + None + ); + assert_eq!(CosmosDriver::query_plan_request_body(&empty).unwrap(), None); + } + #[cfg(feature = "preview_dtx")] fn dtx_response( status_code: azure_core::http::StatusCode, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs index 6ad6b2e1cd1..7438652f80f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs @@ -330,6 +330,7 @@ mod tests { use super::*; use crate::driver::dataflow::mocks::*; use crate::models::ResponseBody; + use serde::Serialize; /// Builds a query-page body from a list of raw JSON document texts. fn page_body(documents: &[&str]) -> Vec { @@ -348,6 +349,34 @@ mod tests { }) } + fn binary_float_page(values: &[f64], is_terminal: bool) -> crate::error::Result { + #[derive(Serialize)] + struct FeedBody<'a> { + #[serde(rename = "_rid")] + rid: &'static str, + #[serde(rename = "Documents")] + documents: &'a [f64], + #[serde(rename = "_count")] + count: usize, + } + + let body = crate::binary_json::to_vec(&FeedBody { + rid: "", + documents: values, + count: values.len(), + }) + .unwrap(); + assert_eq!( + crate::binary_json::to_vec(&1.0f64).unwrap()[1], + crate::binary_json::markers::NUMBER_DOUBLE, + "the fixture must use the service's Float64 number form" + ); + Ok(PageResult::Page { + response: response(&body), + is_terminal, + }) + } + fn charged_page( documents: &[&str], is_terminal: bool, @@ -498,6 +527,17 @@ mod tests { assert_eq!(drain(&mut node).await.len(), 1); } + #[tokio::test] + async fn unordered_collapses_text_integer_and_binary_float64_pages() { + let child = MockLeaf::with_pages(vec![ + page(&["1"], false), + binary_float_page(&[1.0], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(strings(&drain(&mut node).await), vec!["1"]); + } + // ── Ordered map ────────────────────────────────────────────────────── #[tokio::test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs index cc4cb319c90..488e93fc79b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs @@ -277,27 +277,64 @@ pub(crate) fn is_streaming_order_by(info: &QueryInfo) -> bool { !info.order_by.is_empty() && !info.has_non_streaming_order_by } -/// Builds a [`streaming_ordered_merge::StreamingOrderedMerge`] pipeline -/// from a backend query plan whose `queryInfo.orderBy` is non-empty. -/// Mirrors [`build_sequential_drain`]'s shape, but snapshots every -/// still-active range explicitly since global ordering means any range -/// may still have unemitted rows. +/// Fingerprints the caller's query body and feed scope before request encoding. +pub(crate) fn streaming_query_fingerprint( + operation: &CosmosOperation, +) -> crate::error::Result { + streaming_ordered_merge::query_fingerprint(operation.body(), operation.target()).map_err( + |source| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERIALIZATION_REQUEST_BODY_INVALID) + .with_message( + "failed to normalize query body for continuation-token fingerprinting", + ) + .with_source(source) + .build() + }, + ) +} + +/// Builds a streaming `ORDER BY` pipeline using the operation's current body. /// -/// `resume` re-resolves each saved range against current topology and -/// rebuilds it via [`streaming_ordered_merge::build_children`], the same -/// path a live split uses. +/// Unit tests call this before any request encoding. Production planning uses +/// [`build_streaming_ordered_merge_with_fingerprint`] with the fingerprint +/// captured before binary conversion. +#[cfg(test)] pub(crate) async fn build_streaming_ordered_merge( query_plan: &QueryPlan, topology_provider: &mut dyn TopologyProvider, operation: &Arc, resume: Option, +) -> crate::error::Result { + let query_fingerprint = streaming_query_fingerprint(operation)?; + build_streaming_ordered_merge_with_fingerprint( + query_plan, + topology_provider, + operation, + resume, + query_fingerprint, + ) + .await +} + +pub(crate) async fn build_streaming_ordered_merge_with_fingerprint( + query_plan: &QueryPlan, + topology_provider: &mut dyn TopologyProvider, + operation: &Arc, + resume: Option, + query_fingerprint: String, ) -> crate::error::Result { let distinct_type = plan_distinct_type(query_plan); let (inner_resume, last_hash) = peel_distinct_resume(resume, distinct_type)?; let resumed_drained = matches!(inner_resume, Some(PipelineNodeState::Drained)); - let pipeline = - build_streaming_ordered_merge_inner(query_plan, topology_provider, operation, inner_resume) - .await?; + let pipeline = build_streaming_ordered_merge_inner( + query_plan, + topology_provider, + operation, + inner_resume, + query_fingerprint, + ) + .await?; Ok(apply_distinct( pipeline, distinct_type, @@ -311,6 +348,7 @@ async fn build_streaming_ordered_merge_inner( topology_provider: &mut dyn TopologyProvider, operation: &Arc, resume: Option, + query_fingerprint: String, ) -> crate::error::Result { validate_query_plan_for_streaming_order_by(query_plan)?; let info = query_plan @@ -339,13 +377,10 @@ async fn build_streaming_ordered_merge_inner( let plain_operation = Arc::new((**operation).clone().with_body(plain_body)); let is_resume = resume.is_some(); - // The feed scope is folded into the fingerprint (see - // `streaming_ordered_merge::query_fingerprint`) because nothing else binds + // The feed scope is folded into the fingerprint because nothing else binds // a token to it: a resumed node treats its saved ranges as authoritative, // and `is_valid_for_operation` checks only the operation kind and RID. let scope_range = operation.target(); - let query_fingerprint = - streaming_ordered_merge::query_fingerprint(operation.body(), scope_range); let saved_ranges = match resume { None => None, Some(PipelineNodeState::Drained) => { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs index bc37d9774a1..2dda9a7382d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs @@ -784,12 +784,13 @@ impl PipelineNode for StreamingOrderedMerge { /// the Gateway's rewritten query so a service-side rewrite change does not /// invalidate in-flight tokens. /// -/// Because this hashes the *serialized* body, the query body's serialization -/// shape (serde field order, optional-field emission) is a compatibility -/// surface: changing it invalidates in-flight tokens with a hard -/// `CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID` rather than silently -/// resuming the wrong query. -pub(super) fn query_fingerprint(body: Option<&[u8]>, scope: Option<&FeedRange>) -> String { +/// Text bodies retain their exact serialized form for compatibility with +/// existing tokens. The planner normally calls this before request encoding; +/// restoring an already-binary caller body to text is a best-effort fallback. +pub(super) fn query_fingerprint( + body: Option<&[u8]>, + scope: Option<&FeedRange>, +) -> crate::binary_json::Result { // The body hash is rendered fixed-width first so the two components can // never run together; EPK hex is `[0-9A-F]*`, so neither separator can // occur inside a bound. Bounds render canonically (trailing zero bytes @@ -797,7 +798,16 @@ pub(super) fn query_fingerprint(body: Option<&[u8]>, scope: Option<&FeedRange>) // backend and other SDKs may hand back a bound with that padding trimmed. // An absent scope hashes as empty, which stays distinct from the // full-container range (`-FF`). - let body_hash = crate::models::murmur_hash::murmurhash3_128(body.unwrap_or_default(), 0); + let normalized_body; + let body = match body { + Some(body) if crate::binary_json::is_binary(body) => { + normalized_body = crate::binary_json::transcode_to_text(body)?; + normalized_body.as_slice() + } + Some(body) => body, + None => &[], + }; + let body_hash = crate::models::murmur_hash::murmurhash3_128(body, 0); let scope = match scope { Some(range) => format!( "{}-{}", @@ -806,13 +816,13 @@ pub(super) fn query_fingerprint(body: Option<&[u8]>, scope: Option<&FeedRange>) ), None => String::new(), }; - format!( + Ok(format!( "{:032x}", crate::models::murmur_hash::murmurhash3_128( format!("{body_hash:032x}:{scope}").as_bytes(), 0 ) - ) + )) } /// Builds the child streams needed to cover `scope`, given its topology @@ -2920,10 +2930,10 @@ mod tests { #[test] fn query_fingerprint_distinguishes_feed_scope() { let body = br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#; - let full = query_fingerprint(Some(body), Some(&FeedRange::full())); - let left = query_fingerprint(Some(body), Some(&range("", "80"))); - let right = query_fingerprint(Some(body), Some(&range("80", "FF"))); - let unscoped = query_fingerprint(Some(body), None); + let full = query_fingerprint(Some(body), Some(&FeedRange::full())).unwrap(); + let left = query_fingerprint(Some(body), Some(&range("", "80"))).unwrap(); + let right = query_fingerprint(Some(body), Some(&range("80", "FF"))).unwrap(); + let unscoped = query_fingerprint(Some(body), None).unwrap(); assert_ne!(full, left); assert_ne!(full, right); @@ -2938,8 +2948,8 @@ mod tests { fn query_fingerprint_separators_cannot_collide() { // Both scopes render as `408080` once the bound separator is dropped. assert_ne!( - query_fingerprint(None, Some(&range("40", "8080"))), - query_fingerprint(None, Some(&range("4080", "80"))), + query_fingerprint(None, Some(&range("40", "8080"))).unwrap(), + query_fingerprint(None, Some(&range("4080", "80"))).unwrap(), ); } @@ -2950,22 +2960,33 @@ mod tests { #[test] fn query_fingerprint_ignores_trailing_zero_padding_in_scope() { assert_eq!( - query_fingerprint(None, Some(&range("", "80"))), - query_fingerprint(None, Some(&range("", "8000"))), + query_fingerprint(None, Some(&range("", "80"))).unwrap(), + query_fingerprint(None, Some(&range("", "8000"))).unwrap(), ); assert_eq!( - query_fingerprint(None, Some(&range("40", "80"))), - query_fingerprint(None, Some(&range("400000", "8000"))), + query_fingerprint(None, Some(&range("40", "80"))).unwrap(), + query_fingerprint(None, Some(&range("400000", "8000"))).unwrap(), ); } - /// Same body and same scope is stable, so an unchanged query resumes. + /// Preserve the exact fingerprint minted before binary query encoding so + /// in-flight text tokens remain valid. #[test] - fn query_fingerprint_is_stable_for_identical_inputs() { - let body = br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#; + fn parameterized_text_query_fingerprint_matches_historical_value() { + let body = br#"{"query":"SELECT * FROM c WHERE c.rank >= @min ORDER BY c.rank","parameters":[{"name":"@min","value":1}]}"#; + assert_eq!( + query_fingerprint(Some(body), Some(&range("", "80"))).unwrap(), + "b84b9c269862dcd73781038d90add3be", + ); + } + + #[test] + fn already_binary_query_fingerprint_is_stable() { + let text = br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#; + let binary = crate::binary_json::transcode_to_binary(text).unwrap(); assert_eq!( - query_fingerprint(Some(body), Some(&range("", "80"))), - query_fingerprint(Some(body), Some(&range("", "80"))), + query_fingerprint(Some(&binary), Some(&range("", "80"))).unwrap(), + query_fingerprint(Some(&binary), Some(&range("", "80"))).unwrap(), ); } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs index 634ee5f6c8e..6cf303e6340 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs @@ -649,6 +649,7 @@ impl OperationType { | OperationType::Replace | OperationType::Upsert | OperationType::Query + | OperationType::SqlQuery | OperationType::ReadFeed ) } @@ -902,13 +903,13 @@ mod tests { OperationType::Replace, OperationType::Upsert, OperationType::Query, + OperationType::SqlQuery, OperationType::ReadFeed, ] { assert!(op.supports_binary_encoding(), "{op:?} should be supported"); } for op in [ OperationType::Delete, - OperationType::SqlQuery, OperationType::Batch, OperationType::Execute, OperationType::Patch, diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index 124730e2fde..0fb50ed1d33 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -11,16 +11,16 @@ //! Scenarios come from `tests/fixtures/distinct_scenarios.json`, the same //! source-attributed catalog every other layer reads. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; -use azure_core::http::Url; +use azure_core::http::{headers::HeaderName, Request, Url}; use serde::Deserialize; use azure_data_cosmos_driver::driver::CosmosDriver; use azure_data_cosmos_driver::in_memory_emulator::{ - ConsistencyLevel, ContainerConfig, InMemoryEmulatorHttpClient, VirtualAccountConfig, - VirtualRegion, + ConsistencyLevel, ContainerConfig, InMemoryEmulatorHttpClient, RequestObserver, + VirtualAccountConfig, VirtualRegion, }; use azure_data_cosmos_driver::models::{ ContainerReference, CosmosOperation, FeedRange, ItemReference, MaxItemCountHint, PartitionKey, @@ -31,6 +31,10 @@ use azure_data_cosmos_driver::options::{ }; const GATEWAY_URL: &str = "https://eastus.emulator.local"; +static IS_QUERY: HeaderName = HeaderName::from_static("x-ms-documentdb-isquery"); +static IS_QUERY_PLAN: HeaderName = HeaderName::from_static("x-ms-cosmos-is-query-plan-request"); +static SUPPORTED_FORMATS: HeaderName = + HeaderName::from_static("x-ms-cosmos-supported-serialization-formats"); const CATALOG_JSON: &str = include_str!("../fixtures/distinct_scenarios.json"); @@ -77,9 +81,48 @@ fn catalog() -> Catalog { serde_json::from_str(CATALOG_JSON).expect("catalog must parse") } +#[derive(Debug, Default)] +struct QueryRequestRecorder { + binary_modes: Mutex>, +} + +impl QueryRequestRecorder { + fn take(&self) -> Vec { + std::mem::take(&mut *self.binary_modes.lock().unwrap()) + } +} + +impl RequestObserver for QueryRequestRecorder { + fn on_request(&self, request: &Request) { + let headers = request.headers(); + let is_true = |name: &HeaderName| { + headers + .get_optional_str(name) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) + }; + if !is_true(&IS_QUERY) || is_true(&IS_QUERY_PLAN) { + return; + } + let binary = headers + .get_optional_str(&SUPPORTED_FORMATS) + .is_some_and(|value| { + value + .split(',') + .any(|format| format.trim().eq_ignore_ascii_case("CosmosBinary")) + }); + self.binary_modes.lock().unwrap().push(binary); + } +} + /// Builds a two-physical-partition in-memory emulator container and a driver /// wired to it. async fn setup() -> (Arc, Arc) { + setup_with_observer(None).await +} + +async fn setup_with_observer( + observer: Option>, +) -> (Arc, Arc) { let config = VirtualAccountConfig::new(vec![VirtualRegion::new( "East US", Url::parse(GATEWAY_URL).unwrap(), @@ -87,7 +130,12 @@ async fn setup() -> (Arc, Arc) { .unwrap() .with_consistency(ConsistencyLevel::Session); - let emulator = Arc::new(InMemoryEmulatorHttpClient::new(config)); + let emulator = InMemoryEmulatorHttpClient::new(config); + let emulator = match observer { + Some(observer) => emulator.with_request_observer(observer), + None => emulator, + }; + let emulator = Arc::new(emulator); let store = emulator.store(); store.create_database("testdb"); let container_config = ContainerConfig::new() @@ -117,6 +165,17 @@ async fn setup() -> (Arc, Arc) { (emulator, driver) } +async fn setup_with_query_recorder() -> ( + Arc, + Arc, + Arc, +) { + let recorder = Arc::new(QueryRequestRecorder::default()); + let observer: Arc = recorder.clone(); + let (emulator, driver) = setup_with_observer(Some(observer)).await; + (emulator, driver, recorder) +} + async fn seed( driver: &CosmosDriver, container: &ContainerReference, @@ -214,6 +273,51 @@ async fn drain_query_with_options( (values, formats) } +async fn drain_with_cross_encoding_resume( + driver: &CosmosDriver, + container: &ContainerReference, + query: &QuerySpec, + mint_options: OperationOptions, + resume_options: OperationOptions, +) -> Vec { + let mut plan = Box::pin(driver.plan_operation( + query_operation(container, query, 1), + &mint_options, + None, + &PlanOptions::default(), + )) + .await + .unwrap(); + let first = driver + .execute_plan(&mut plan, Some(container.clone()), mint_options) + .await + .unwrap() + .expect("ordered query must emit a first page"); + let mut values = documents_of(first); + let token = plan.to_continuation_token().unwrap(); + + let mut resumed = Box::pin(driver.plan_operation( + query_operation(container, query, 1), + &resume_options, + Some(&token), + &PlanOptions::default(), + )) + .await + .unwrap(); + while let Some(response) = driver + .execute_plan( + &mut resumed, + Some(container.clone()), + resume_options.clone(), + ) + .await + .unwrap() + { + values.extend(documents_of(response)); + } + values +} + /// Sorts values by their serialized form so an unordered result can be /// compared deterministically. fn sorted(mut values: Vec) -> Vec { @@ -569,7 +673,7 @@ async fn split_mid_drain_does_not_reemit_deduplicated_values() { #[tokio::test] async fn text_and_binary_query_pages_have_pipeline_parity() { - let (_emulator, driver) = setup().await; + let (_emulator, driver, recorder) = setup_with_query_recorder().await; let container = driver .resolve_container("testdb", "testcoll") .await @@ -611,6 +715,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { OperationOptions::default(), ) .await; + let text_request_modes = recorder.take(); let binary_options = OperationOptionsBuilder::new() .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)) .build(); @@ -622,6 +727,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { binary_options.clone(), ) .await; + let binary_request_modes = recorder.take(); let binary_as_text_options = OperationOptionsBuilder::new() .with_binary_encoding( BinaryEncodingOptions::new() @@ -637,6 +743,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { binary_as_text_options, ) .await; + let binary_as_text_request_modes = recorder.take(); let (_, planned_binary_formats) = drain_query_with_options( &driver, &container, @@ -645,6 +752,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { OperationOptions::default(), ) .await; + let planned_binary_request_modes = recorder.take(); let (_, planned_text_formats) = drain_query_with_options( &driver, &container, @@ -653,6 +761,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { binary_options, ) .await; + let planned_text_request_modes = recorder.take(); if query.distinct_type == "Ordered" { assert_eq!(binary, text, "{}", query.text); @@ -671,20 +780,122 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { "{}", query.text ); + assert!( + !binary_request_modes.is_empty() + && binary_request_modes.iter().all(|is_binary| *is_binary), + "binary negotiation must reach per-partition requests: {}", + query.text + ); + assert!( + !text_request_modes.is_empty() && text_request_modes.iter().all(|is_binary| !is_binary), + "text requests must not advertise binary: {}", + query.text + ); assert!( binary_as_text_formats.iter().all(|is_binary| !is_binary), "{}", query.text ); + assert!( + !binary_as_text_request_modes.is_empty() + && binary_as_text_request_modes + .iter() + .all(|is_binary| *is_binary), + "text-response mode must still negotiate a binary wire: {}", + query.text + ); assert!( planned_binary_formats.iter().all(|is_binary| *is_binary), "planning must own response encoding even when execution options differ: {}", query.text ); + assert!( + !planned_binary_request_modes.is_empty() + && planned_binary_request_modes + .iter() + .all(|is_binary| *is_binary), + "binary planning must own wire negotiation: {}", + query.text + ); assert!( planned_text_formats.iter().all(|is_binary| !is_binary), "execution options must not enable binary after text planning: {}", query.text ); + assert!( + !planned_text_request_modes.is_empty() + && planned_text_request_modes + .iter() + .all(|is_binary| !is_binary), + "binary execution options must not change text-planned negotiation: {}", + query.text + ); + } +} + +#[tokio::test] +async fn ordered_distinct_tokens_resume_across_binary_modes() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .unwrap(); + let documents: Vec = [("a", "pk1", 1), ("b", "pk2", 2), ("c", "pk3", 3)] + .into_iter() + .map(|(id, pk, value)| serde_json::json!({"id": id, "pk": pk, "value": value})) + .collect(); + seed(&driver, &container, &documents).await; + let binary = OperationOptionsBuilder::new() + .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)) + .build(); + for (query, expected) in [ + ( + QuerySpec { + text: "SELECT DISTINCT VALUE c.value FROM c ORDER BY c.value".to_owned(), + parameters: Vec::new(), + distinct_type: "Ordered".to_owned(), + }, + vec![ + serde_json::json!(1), + serde_json::json!(2), + serde_json::json!(3), + ], + ), + ( + QuerySpec { + text: "SELECT DISTINCT VALUE c.value FROM c WHERE c.value >= @min ORDER BY c.value" + .to_owned(), + parameters: vec![serde_json::json!({"name": "@min", "value": 2})], + distinct_type: "Ordered".to_owned(), + }, + vec![serde_json::json!(2), serde_json::json!(3)], + ), + ] { + assert_eq!( + drain_with_cross_encoding_resume( + &driver, + &container, + &query, + OperationOptions::default(), + binary.clone(), + ) + .await, + expected, + "text-minted token must resume in binary mode: {}", + query.text + ); + assert_eq!( + drain_with_cross_encoding_resume( + &driver, + &container, + &query, + binary.clone(), + OperationOptions::default(), + ) + .await, + expected, + "binary-minted token must resume in text mode: {}", + query.text + ); } } From bf5395eea53f55624e0374325737ca24a0fb82e5 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 09:26:39 -0700 Subject: [PATCH 3/8] Strengthen binary token regression Serialize emulator query fixtures in the same field order as the public SDK so parameterized text and binary fingerprints actually diverge when capture happens too late. Remove the vacuous already-binary fingerprint self-comparison. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../dataflow/streaming_ordered_merge.rs | 10 ---------- .../in_memory_emulator_tests/distinct.rs | 19 ++++++++++++++----- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs index 2dda9a7382d..7111a3b2172 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs @@ -2979,14 +2979,4 @@ mod tests { "b84b9c269862dcd73781038d90add3be", ); } - - #[test] - fn already_binary_query_fingerprint_is_stable() { - let text = br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#; - let binary = crate::binary_json::transcode_to_binary(text).unwrap(); - assert_eq!( - query_fingerprint(Some(&binary), Some(&range("", "80"))).unwrap(), - query_fingerprint(Some(&binary), Some(&range("", "80"))).unwrap(), - ); - } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index 0fb50ed1d33..bc6fdb49025 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -15,7 +15,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use azure_core::http::{headers::HeaderName, Request, Url}; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use azure_data_cosmos_driver::driver::CosmosDriver; use azure_data_cosmos_driver::in_memory_emulator::{ @@ -203,10 +203,19 @@ async fn seed( } fn query_body(query: &QuerySpec) -> Vec { - serde_json::to_vec(&serde_json::json!({ - "query": query.text, - "parameters": query.parameters, - })) + #[derive(Serialize)] + struct QueryBody<'a> { + query: &'a str, + #[serde(skip_serializing_if = "Vec::is_empty")] + parameters: Vec, + } + + // Match azure_data_cosmos::Query field order because fingerprints hash bytes; + // a json! map sorts these keys and masks cross-encoding token regressions. + serde_json::to_vec(&QueryBody { + query: &query.text, + parameters: query.parameters.clone(), + }) .unwrap() } From 4ca61e0e4f1aa2da95055dc11fcc121ea750bd63 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 09:36:17 -0700 Subject: [PATCH 4/8] Return fingerprint invariant errors Return a typed CosmosError when streaming ORDER BY planning lacks its pre-encoding fingerprint instead of panicking across an FFI host boundary. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../src/driver/cosmos_driver.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 2a5f64769fa..189535fa564 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -3404,8 +3404,15 @@ impl CosmosDriver { .as_ref() .is_some_and(planner::is_streaming_order_by) { - let query_fingerprint = query_fingerprint - .expect("query operations capture a fingerprint before request encoding"); + let query_fingerprint = query_fingerprint.ok_or_else(|| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE) + .with_message( + "internal error: streaming ORDER BY query is missing its pre-encoding \ + request fingerprint", + ) + .build() + })?; let pipeline = planner::build_streaming_ordered_merge_with_fingerprint( &query_plan, &mut topology, From e64c4784c25f6e173ad49b1a52ead9930eb94514 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 09:59:14 -0700 Subject: [PATCH 5/8] Add query fingerprint invariant status Return a dedicated client-side status when streaming ORDER BY planning is missing its pre-encoding fingerprint. Export the same status through the native C ABI so FFI hosts receive a named diagnostic instead of a magic number. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../src/driver/cosmos_driver.rs | 4 +++- .../src/error/cosmos_status.rs | 24 +++++++++++++++++++ .../include/azurecosmosdriver.h | 4 ++++ .../src/error.rs | 2 ++ 4 files changed, 33 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 189535fa564..a49cfd0b6d7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -3406,7 +3406,9 @@ impl CosmosDriver { { let query_fingerprint = query_fingerprint.ok_or_else(|| { crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE) + .with_status( + crate::error::CosmosStatus::CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING, + ) .with_message( "internal error: streaming ORDER BY query is missing its pre-encoding \ request fingerprint", diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs index b583754b2b6..56e63e4d3c5 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs @@ -524,6 +524,7 @@ impl SubStatusCode { 20213 => Some("ClientContinuationTokenSavedRangeUnhonored"), 20214 => Some("ClientContinuationTokenOrderByStateInvalid"), 20215 => Some("ClientStreamingMergeSplitReplacementInvalid"), + 20216 => Some("ClientStreamingOrderByFingerprintMissing"), 20300 => Some("ClientNoOverlappingFeedRangesForSessionToken"), 20301 => Some("ClientNoThroughputOfferForResource"), 20302 => Some("ClientQueryPlanProducedEmptyRanges"), @@ -1512,6 +1513,10 @@ impl SubStatusCode { pub const CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID: SubStatusCode = SubStatusCode(20215); + /// A streaming `ORDER BY` plan is missing the request fingerprint captured + /// before wire encoding (20216). Indicates an internal planner invariant violation. + pub const CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING: SubStatusCode = SubStatusCode(20216); + // ----- 20300-20349: SDK-detected service contract violations ----- /// The supplied session-token feed ranges contain no overlap with @@ -2500,6 +2505,13 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID), }; + /// 500 / 20216 — streaming `ORDER BY` planning is missing its + /// pre-encoding request fingerprint. + pub const CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING), + }; + // SDK-detected service contract violations (HTTP varies, sub-status 20300-20349) /// 410 / 20300 — the supplied session-token feed ranges contain no @@ -2670,6 +2682,18 @@ mod tests { ); } + #[test] + fn streaming_order_by_fingerprint_missing_status_is_searchable() { + assert_eq!( + CosmosStatus::CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING.name(), + Some("ClientStreamingOrderByFingerprintMissing") + ); + assert_eq!( + CosmosStatus::CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING.sub_status(), + Some(SubStatusCode::CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING) + ); + } + #[test] fn with_sub_status_unambiguous() { let status = CosmosStatus::new(StatusCode::TooManyRequests).with_sub_status(3200); diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h index a8970d39c23..07bb427338d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h +++ b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h @@ -699,6 +699,10 @@ enum cosmos_sub_status_t * `CLIENT_CONTINUATION_TOKEN_SAVED_RANGE_UNHONORED` (20213). */ COSMOS_SUB_STATUS_CLIENT_CONTINUATION_TOKEN_SAVED_RANGE_UNHONORED = 20213, + /** + * `CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING` (20216). + */ + COSMOS_SUB_STATUS_CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING = 20216, /** * `CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE` (20301). */ diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/error.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/error.rs index 1fb8c81b630..8456f8b8250 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/error.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/error.rs @@ -184,6 +184,8 @@ pub enum CosmosSubStatus { CosmosSubStatusClientSingletonOperationReturnedEmptyPage = 20210, /// `CLIENT_CONTINUATION_TOKEN_SAVED_RANGE_UNHONORED` (20213). CosmosSubStatusClientContinuationTokenSavedRangeUnhonored = 20213, + /// `CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING` (20216). + CosmosSubStatusClientStreamingOrderByFingerprintMissing = 20216, /// `CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE` (20301). CosmosSubStatusClientNoThroughputOfferForResource = 20301, /// `CLIENT_QUERY_PLAN_PRODUCED_EMPTY_RANGES` (20302). From 40f089b6ca116624f97df183573c2dabaef2229f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 10:05:56 -0700 Subject: [PATCH 6/8] Align binary feed response handling Derive pipeline-boundary binary restoration from the request-side eligibility gate so Query, SqlQuery, and ReadFeed cannot drift apart. Exhaustively pin request and restore behavior for every operation type. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../src/driver/cosmos_driver.rs | 68 +++++++++++-------- 1 file changed, 41 insertions(+), 27 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index a49cfd0b6d7..0c6ef56bd24 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2624,6 +2624,12 @@ impl CosmosDriver { && !operation.is_change_feed() } + /// Whether a feed pipeline negotiated binary and must restore that wire + /// format after its text-based processing stages complete. + fn binary_feed_response_restore_applies(operation: &CosmosOperation) -> bool { + Self::binary_encoding_applies(operation) && operation.operation_type().is_feed() + } + /// Applies request-side binary encoding to an operation: transcodes a text /// request body to Cosmos binary JSON (an already-binary or empty body is /// passed through) and advertises binary responses via the @@ -2929,10 +2935,7 @@ impl CosmosDriver { if let Some(response) = response.as_mut() { if request_text_response { response.transcode_body_to_text()?; - } else if matches!( - plan.operation().operation_type(), - crate::models::OperationType::Query | crate::models::OperationType::ReadFeed - ) { + } else if Self::binary_feed_response_restore_applies(plan.operation()) { response.transcode_body_to_binary()?; } } @@ -6323,17 +6326,29 @@ mod tests { // ── apply_request_binary_encoding (schema-agnostic request-side encode) ── #[test] - fn binary_encoding_applies_only_to_document_item_ops() { + fn binary_encoding_request_and_feed_restore_predicates_agree() { use crate::models::{OperationType, ResourceType}; let container = epk_test_container(r#"{"paths":["/pk"],"version":2}"#); - for op in [ - OperationType::Create, - OperationType::Read, - OperationType::Replace, - OperationType::Upsert, - OperationType::Query, - OperationType::ReadFeed, + for (op, request_applies, feed_restore_applies) in [ + (OperationType::Create, true, false), + (OperationType::Read, true, false), + (OperationType::ReadFeed, true, true), + (OperationType::Replace, true, false), + (OperationType::Delete, false, false), + (OperationType::Upsert, true, false), + (OperationType::Query, true, true), + (OperationType::SqlQuery, true, true), + (OperationType::QueryPlan, false, false), + (OperationType::Batch, false, false), + (OperationType::Head, false, false), + (OperationType::HeadFeed, false, false), + (OperationType::Execute, false, false), + (OperationType::Patch, false, false), + #[cfg(feature = "preview_dtx")] + (OperationType::CommitDistributedTransaction, false, false), + #[cfg(feature = "preview_dtx")] + (OperationType::ReadDistributedTransaction, false, false), ] { let operation = CosmosOperation::new( op, @@ -6342,23 +6357,15 @@ mod tests { .into_feed_reference(), Some(crate::models::FeedRange::full()), ); - assert!( + assert_eq!( CosmosDriver::binary_encoding_applies(&operation), - "Document + {op:?} should be binary-encodable", - ); - } - - for op in [OperationType::Delete, OperationType::Patch] { - let operation = CosmosOperation::new( - op, - crate::models::CosmosResourceReference::from(container.clone()) - .with_resource_type(ResourceType::Document) - .into_feed_reference(), - Some(crate::models::FeedRange::full()), + request_applies, + "unexpected request-side binary eligibility for Document + {op:?}", ); - assert!( - !CosmosDriver::binary_encoding_applies(&operation), - "Document + {op:?} must not be binary-encoded", + assert_eq!( + CosmosDriver::binary_feed_response_restore_applies(&operation), + feed_restore_applies, + "request/restore binary eligibility diverged for Document + {op:?}", ); } @@ -6389,12 +6396,19 @@ mod tests { !CosmosDriver::binary_encoding_applies(&operation), "{rt:?} + {op:?} must not be binary-encoded (control plane)", ); + assert!( + !CosmosDriver::binary_feed_response_restore_applies(&operation), + "{rt:?} + {op:?} must not restore binary feed responses", + ); } } let change_feed = CosmosOperation::change_feed(container, Some(crate::models::FeedRange::full())); assert!(!CosmosDriver::binary_encoding_applies(&change_feed)); + assert!(!CosmosDriver::binary_feed_response_restore_applies( + &change_feed + )); } fn binary_encoding_test_operation(body: Vec) -> CosmosOperation { From 6b02c99bd1952c9c874cd40e406020b84e109c73 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 12 Aug 2026 12:44:04 -0700 Subject: [PATCH 7/8] Refine binary query response handling Keep SQL query request bodies textual while negotiating binary response pages, and normalize integral doubles when binary responses are explicitly transcoded to text.\n\nModel service number behavior in the emulator, strengthen query fuzzer and split coverage, and document the measured Gateway behavior and defensive invariants.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../tests/binary_roundtrip_fuzzer.rs | 98 +++++++++++++++++-- .../cosmos_query_distinct_split.rs | 42 +++++--- .../docs/BINARY_ENCODING_HLD.md | 32 +++--- .../src/binary_json/mod.rs | 70 ++++++++++++- .../src/driver/cosmos_driver.rs | 59 ++++++++--- .../src/driver/dataflow/planner.rs | 25 +---- .../src/driver/dataflow/query_response.rs | 37 +------ .../dataflow/streaming_ordered_merge.rs | 46 ++++----- .../src/error/cosmos_status.rs | 4 +- .../src/in_memory_emulator/operations.rs | 10 +- .../src/in_memory_emulator/response.rs | 29 +++++- .../src/models/response_body.rs | 12 +-- .../binary_response_format.rs | 4 +- .../in_memory_emulator_tests/distinct.rs | 28 ++++++ 14 files changed, 343 insertions(+), 153 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs index b45695623d1..b7f7bfff010 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs @@ -50,6 +50,7 @@ use azure_data_cosmos::{ }; use azure_data_cosmos_driver::models::ConnectionString; use futures::TryStreamExt; +use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use serde_json::{Map, Number, Value}; use sha2::{Digest, Sha256}; @@ -1882,6 +1883,11 @@ struct IntProbe { wide: u64, } +#[derive(Deserialize, Debug, PartialEq)] +struct QueryIntProbe { + int: i64, +} + /// Round-trips an [`IntProbe`] so `deserialize_integer` (the production change) /// is exercised live, then asserts the typed values survived. Only meaningful on /// the pure-binary config (see the call site). @@ -1972,12 +1978,12 @@ where } } -async fn query_values( +async fn query_values( container: &ContainerClient, sql: &str, run_id: &str, context: &str, -) -> Result, Box> { +) -> Result, Box> { let mut attempt = 0; loop { attempt += 1; @@ -2016,7 +2022,7 @@ fn canonical_query_results(values: Vec, ordered: bool) -> Vec { } value => value, }; - canonicalize(&value) + structural_query_key(&value) }) .collect(); if !ordered { @@ -2025,6 +2031,49 @@ fn canonical_query_results(values: Vec, ordered: bool) -> Vec { canonical } +fn structural_query_key(value: &Value) -> String { + fn write(value: &Value, output: &mut String) { + match value { + Value::Null => output.push('n'), + Value::Bool(value) => output.push(if *value { 't' } else { 'f' }), + Value::Number(number) => { + if let Some(value) = number.as_i64() { + output.push_str(&format!("i:{value};")); + } else if let Some(value) = number.as_u64() { + output.push_str(&format!("u:{value};")); + } else { + let value = number.as_f64().expect("finite JSON number"); + output.push_str(&format!("d:{:016x};", value.to_bits())); + } + } + Value::String(value) => { + output.push_str("s:"); + output.push_str(&serde_json::to_string(value).expect("string serializes")); + output.push(';'); + } + Value::Array(values) => { + output.push('['); + values.iter().for_each(|value| write(value, output)); + output.push(']'); + } + Value::Object(values) => { + output.push('{'); + let mut entries: Vec<_> = values.iter().collect(); + entries.sort_unstable_by_key(|(key, _)| *key); + for (key, value) in entries { + write(&Value::String(key.clone()), output); + write(value, output); + } + output.push('}'); + } + } + } + + let mut output = String::new(); + write(value, &mut output); + output +} + // ───────────────────────────────────────────────────────────────────────────── // The fuzzer test // ───────────────────────────────────────────────────────────────────────────── @@ -2108,7 +2157,9 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { // AZURE_COSMOS_FUZZ_SEED reproduces the exact document *and* its // canonical form (a random `Uuid` here would defeat that promise). let id = format!("{run_id}-{config_idx}"); - let pk = format!("{run_id}-pk-{partition_bucket}-{config_idx}"); + // Reuse 16 logical partitions per encoding config so point + // operations cover multiple items sharing the same partition key. + let pk = format!("fuzz-pk-{partition_bucket}-{config_idx}"); let mut doc = base_doc.clone(); doc.insert("id".to_string(), Value::String(id.clone())); doc.insert("pk".to_string(), Value::String(pk.clone())); @@ -2238,9 +2289,8 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { // The four ops above decode into `serde_json::Value` (→ // `deserialize_any`), so they do NOT cover the native typed-integer // path (`deserialize_integer`) this PR ships. A typed probe covers it - // live on the pure-binary config — the only mode that returns binary - // for an integer field (text modes return text, which `serde_json` - // rejects into an integer). + // live on the pure-binary config — the only mode that exercises the + // binary deserializer's integral-Double coercion directly. if *label == "binary" { assert_typed_integer_probe(&container, &pk, iter, cfg.seed, &context).await?; checked += 1; @@ -2285,17 +2335,41 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { } else { expected = Some(actual); } + checked += 1; } } + for (label, client) in &clients { + let container = client + .database_client(&database_name) + .container_client(&container_name) + .await?; + let context = format!( + "iter={iter} config={label} query=typed-integer seed={}", + cfg.seed + ); + let values: Vec = query_values( + &container, + "SELECT VALUE {\"int\": 7} FROM c WHERE c.fuzzRun = @run", + &run_id, + &context, + ) + .await?; + assert!( + !values.is_empty() && values.iter().all(|value| value.int == 7), + "{context}: typed integer query returned unexpected values: {values:?}" + ); + checked += 1; + } + if (iter + 1) % 100 == 0 { println!("... {} iterations, {checked} round-trips OK", iter + 1); } } println!( - "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops + 3 queries/config = {checked} canonical comparisons, all equal (seed={})", + "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops + 4 queries/config = {checked} canonical comparisons, all equal (seed={})", cfg.iterations, configs.len(), cfg.seed @@ -2545,6 +2619,14 @@ mod tests { assert_eq!(canon(&serde_json::json!(-0.0)), "0"); } + #[test] + fn query_comparison_preserves_integral_float_representation() { + assert_ne!( + structural_query_key(&serde_json::json!(1)), + structural_query_key(&serde_json::json!(1.0)) + ); + } + #[test] fn canonicalize_keeps_non_integral_floats() { assert_eq!(canon(&serde_json::json!(3.5)), "3.5"); diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs index a37a4bca540..8c3957fe470 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -//! Live-only binary-encoding split coverage for cross-partition `DISTINCT`. +//! Live split coverage for text and binary cross-partition `DISTINCT`. //! //! Neither .NET nor Java tests `DISTINCT` against a real partition split //! (.NET's `FullPipelineTests.TestMerge` covers `ORDER BY` only), so this is @@ -119,12 +119,8 @@ where /// Part 2 captures an ordered `DISTINCT` continuation token before the split /// (already taken, since the split happened in part 1) and resumes it against /// the post-split topology. -#[tokio::test] -#[cfg_attr( - not(test_category = "split"), - ignore = "requires test_category 'split'" -)] -pub async fn binary_distinct_query_across_split_returns_each_value_once( +async fn run_distinct_query_across_split_returns_each_value_once( + binary: bool, ) -> Result<(), Box> { TestClient::run_with_unique_db( async |run_context, db_client| { @@ -272,11 +268,33 @@ pub async fn binary_distinct_query_across_split_returns_each_value_once( }, // A real split takes minutes; the 80s default would abort mid-poll. // Matches the other split tests in this directory. - Some( - TestOptions::new() - .with_timeout(Duration::from_secs(40 * 60)) - .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)), - ), + Some({ + let options = TestOptions::new().with_timeout(Duration::from_secs(40 * 60)); + if binary { + options.with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)) + } else { + options + } + }), ) .await } + +#[tokio::test] +#[cfg_attr( + not(test_category = "split"), + ignore = "requires test_category 'split'" +)] +pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), Box> { + run_distinct_query_across_split_returns_each_value_once(false).await +} + +#[tokio::test] +#[cfg_attr( + not(test_category = "split"), + ignore = "requires test_category 'split'" +)] +pub async fn binary_distinct_query_across_split_returns_each_value_once( +) -> Result<(), Box> { + run_distinct_query_across_split_returns_each_value_once(true).await +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md b/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md index 9e29b85ebe7..093ac1da421 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md +++ b/sdk/cosmos/azure_data_cosmos_driver/docs/BINARY_ENCODING_HLD.md @@ -15,10 +15,10 @@ Adds first-class support for **Cosmos binary JSON** to the Rust SDK and driver. This design delivers a **complete decoder** and a **native serde codec**, and makes binary encoding a **driver capability** on `OperationOptions.binary_encoding` so it is shared by every consumer of the driver — the Rust SDK **and** any FFI-based SDK (.NET, Java, Go, …). It is opt-in (`CosmosClientBuilder::with_binary_encoding_options`, with an `AZURE_COSMOS_BINARY_ENCODING_ENABLED` environment-variable fallback). When the option is off, every request and response is **byte-for-byte unchanged** — the binary code is inert. -Because the option lives on the driver and is schema-agnostic, the driver performs the byte-level transcoding **both ways** when needed: +Because the option lives on the driver and is schema-agnostic, the driver performs byte-level transcoding when needed: -* an opt-in **text-response** mode (`BinaryEncodingOptions::request_text_response`) keeps the wire binary in both directions (efficient RUs and bandwidth) while the driver transcodes the binary **response** back to text JSON; -* a caller that deals only in **text** (most importantly an FFI host) can enable binary and the driver transcodes its text **request** body to binary — so it gets a fully binary wire **without encoding anything itself**. +* an opt-in **text-response** mode (`BinaryEncodingOptions::request_text_response`) keeps responses binary on the wire (efficient RUs and bandwidth) while the driver transcodes them back to text JSON; +* for item writes, a caller that deals only in **text** (most importantly an FFI host) can enable binary and let the driver transcode its text request body to binary. Query specifications deliberately remain text. A self-contained, in-tree **end-to-end validation loop** is included via the in-memory emulator (no Docker, no live account, no external test vectors). @@ -57,8 +57,8 @@ FFI hosts set the equivalent flat fields on the C ABI `cosmos_operation_options_ When `binary_encoding.enabled` is set, `CosmosDriver::plan_operation` and `execute_plan` own the wire format both ways: -* **Request** (`apply_request_binary_encoding`) — transcodes a **text** request body to Cosmos binary JSON via `binary_json::transcode_to_binary` (`serde_json::from_slice` → `encode`) and advertises `JsonText,CosmosBinary`. An **already-binary** or empty body passes through unchanged, so a caller that pre-encoded pays nothing. -* **Response** (when `request_text_response` is set) — transcodes the binary response back to text JSON via `binary_json::transcode_to_text` (`decode` → `serde_json::to_vec`). The wire stays binary in both directions. +* **Request** (`apply_request_binary_encoding`) — transcodes a **text item** request body to Cosmos binary JSON via `binary_json::transcode_to_binary` (`serde_json::from_slice` → `encode`) and advertises `CosmosBinary`. An **already-binary** or empty item body passes through unchanged. Query and `SqlQuery` bodies remain text and only advertise response support. +* **Response** (when `request_text_response` is set) — transcodes the binary response back to text JSON via `binary_json::transcode_to_text` (`decode` → integral-double normalization → `serde_json::to_vec`). This matches the service's text rendering of stored integral doubles. This keeps transcoding in the **driver** (not the backend) and, because it is schema-agnostic, lets a text-only FFI host get an efficient binary wire without any encoding on its side. @@ -71,12 +71,20 @@ The Rust SDK keeps a typed fast path: `serialize_item_body` encodes `T: Serializ When binary is enabled, eligible item, query, and document read-feed operations set: ``` -x-ms-cosmos-supported-serialization-formats: JsonText,CosmosBinary +x-ms-cosmos-supported-serialization-formats: CosmosBinary ``` -The value matches the .NET reference (`string.Join(",", JsonText, CosmosBinary)` — no space). The request `Content-Type` stays `application/json`; the service detects the binary body from its first byte. The **driver** sets this header whenever `binary_encoding.enabled` is set — including under `request_text_response`, where the wire stays binary and the driver transcodes the response (see above). +The request `Content-Type` stays `application/json`; binary responses retain that content type and are distinguished by the `0x80` preamble. The **driver** sets this header whenever `binary_encoding.enabled` is set — including under `request_text_response`, where the response wire stays binary and the driver transcodes it (see above). -The `/queryplan` request is the deliberate exception: it remains text and does not advertise binary because that endpoint ignores binary negotiation. Query page bodies are normalized from binary to text once at pipeline ingest so DISTINCT and streaming ORDER BY continue to operate on their existing text representation. Rebuilt pages are restored to the negotiated format at the plan boundary. `ResponseBody::Items` encodes each item independently, including its own `0x80` preamble, so every slice remains a standalone binary document. +The `/queryplan` request is the deliberate exception: it remains text and does not advertise binary because that endpoint ignores binary negotiation. Query page bodies are normalized from binary to text once at pipeline ingest so DISTINCT and streaming ORDER BY continue to operate on their existing text representation. Rebuilt pages are restored to the requested format at the plan boundary, including mixed-rollout text fallback pages. `ResponseBody::Items` binary conversion remains defensive test coverage; production feed paths currently use a single envelope buffer. + +Live gateway probes established the query contract: + +* Binary query request bodies are accepted, both with and without the `CosmosBinary` header. +* A text query body plus the `CosmosBinary` header already returns a binary page, so binary request encoding adds cost without changing the response. +* Text query results render a stored integral double such as `3.0` as `3`. The binary-to-text transcode therefore normalizes integral doubles to integer JSON syntax. + +The driver deliberately sends query specifications as text, matching other SDKs, and uses the header only for query response negotiation. --- @@ -183,7 +191,7 @@ flowchart TD REQ -->|"no (text, e.g. FFI)"| RT["transcode_to_binary
(from_slice → encode)"] REQ -->|"yes (SDK pre-encoded)"| PASS_BIN["pass through"] RT --> HDR - PASS_BIN --> HDR["advertise JsonText,CosmosBinary"] + PASS_BIN --> HDR["advertise CosmosBinary"] HDR --> WIRE["binary body on the wire"] WIRE --> SVC[("Cosmos DB")] SVC --> RESP["binary response (0x80)"] @@ -240,7 +248,7 @@ sequenceDiagram Cod-->>SDK: 0x80-prefixed binary bytes SDK->>DRV: operation + OperationOptions.binary_encoding Note over DRV: request body already binary → pass through - DRV->>Svc: binary body + JsonText,CosmosBinary + DRV->>Svc: binary body + CosmosBinary Svc-->>DRV: binary response (0x80) opt request_text_response DRV->>Cod: transcode_to_text (decode → to_vec) @@ -274,7 +282,7 @@ sequenceDiagram Note over DRV: request body is TEXT → transcode to binary DRV->>Cod: transcode_to_binary (from_slice → encode) Cod-->>DRV: 0x80-prefixed binary body - DRV->>Svc: binary body + JsonText,CosmosBinary + DRV->>Svc: binary body + CosmosBinary Svc-->>DRV: binary response (0x80) DRV->>Cod: transcode_to_text (decode → to_vec) Cod-->>DRV: TEXT json body @@ -301,7 +309,7 @@ When the option is **off**, both paths collapse to the existing text behavior | `deserialize_response` / `ResponseBody::transcode_to_text` | `models/response_body.rs` | Decode choke point for `into_single` / `into_items`; in-place binary→text conversion | | `BinaryEncodingOptions` | `azure_data_cosmos_driver/src/options/binary_encoding.rs` | **Driver-owned** options (`enabled`, `request_text_response`), on `OperationOptions.binary_encoding`; the SDK re-exports it | | `OperationOptions.binary_encoding` | `driver/options/operation_options.rs` | Layered option carrying binary encoding to every consumer (SDK + FFI) | -| `execute_operation` / `apply_request_binary_encoding` | `driver/cosmos_driver.rs` | Driver owns the wire: transcodes text→binary request, advertises `CosmosBinary`, transcodes binary→text response | +| `execute_operation` / `apply_request_binary_encoding` | `driver/cosmos_driver.rs` | Driver owns the wire: transcodes text→binary item requests, leaves query specifications text, advertises `CosmosBinary`, and transcodes binary→text responses | | `CosmosResponse::transcode_body_to_text` | `models/cosmos_response.rs` | Applies driver-side response transcoding | | `resolve_binary_encoding` / `with_binary_encoding` | `azure_data_cosmos/src/clients/{mod,container_client}.rs` | SDK: resolve enablement once; set the option on `OperationOptions` per item op | | `serialize_item_body` | `azure_data_cosmos/src/clients/container_client.rs` | SDK typed fast path: pre-encode `T` to binary (driver passes it through) | diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/binary_json/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/binary_json/mod.rs index b3607cd6a7b..d9c87705088 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/binary_json/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/binary_json/mod.rs @@ -80,9 +80,9 @@ pub fn is_binary(buffer: &[u8]) -> bool { /// Transcodes a Cosmos binary JSON buffer to UTF-8 **text** JSON. /// /// This is the driver-side conversion used when an upstream SDK/app wants to -/// deal only with text JSON while still keeping the wire binary (efficient RUs -/// and network bandwidth): the request and the service response stay binary, -/// and the driver converts the binary response to text before handing it back. +/// deal only with text JSON while still keeping the response wire binary +/// (efficient RUs and network bandwidth): the driver converts the binary +/// response to text before handing it back. /// /// Behavior: /// @@ -102,11 +102,43 @@ pub fn transcode_to_text(buffer: &[u8]) -> Result> { // Already text (or empty): nothing to convert. return Ok(buffer.to_vec()); } - let value = decode(buffer)?; + let mut value = decode(buffer)?; + normalize_integral_floats(&mut value); serde_json::to_vec(&value) .map_err(|e| BinaryError::Custom(format!("failed to re-serialize decoded value: {e}"))) } +pub(crate) fn normalize_integral_floats(value: &mut serde_json::Value) { + const U64_EXCLUSIVE_UPPER_BOUND: f64 = 18_446_744_073_709_551_616.0; + + match value { + serde_json::Value::Number(number) if number.is_f64() => { + let Some(float) = number.as_f64() else { + return; + }; + if float.fract() != 0.0 { + return; + } + + // Cosmos stores JSON numbers as doubles but renders integral values + // as integers in text mode. Use an exclusive 2^64 upper bound so a + // cast cannot saturate u64::MAX; larger doubles stay floating point. + if float >= 0.0 && float < U64_EXCLUSIVE_UPPER_BOUND { + *number = serde_json::Number::from(float as u64); + } else if float >= i64::MIN as f64 && float < 0.0 { + *number = serde_json::Number::from(float as i64); + } + } + serde_json::Value::Array(values) => { + values.iter_mut().for_each(normalize_integral_floats); + } + serde_json::Value::Object(values) => { + values.values_mut().for_each(normalize_integral_floats); + } + _ => {} + } +} + /// Transcodes a UTF-8 **text** JSON buffer to Cosmos **binary** JSON. /// /// This is the mirror of [`transcode_to_text`] for the **request** path: when a @@ -185,6 +217,36 @@ mod tests { assert_eq!(reparsed, value); } + #[test] + fn transcode_binary_integral_float_to_text_integer() { + let binary = encode(&serde_json::json!({ "n": 1.0 })); + + assert_eq!(transcode_to_text(&binary).unwrap(), br#"{"n":1}"#); + } + + #[test] + fn transcode_integral_float_respects_json_integer_bounds() { + let largest_u64_float = f64::from_bits((u64::MAX as f64).to_bits() - 1); + let below_i64_min = f64::from_bits((i64::MIN as f64).to_bits() + 1); + let binary = encode(&serde_json::json!({ + "i64_min": i64::MIN as f64, + "largest_u64_float": largest_u64_float, + "u64_exclusive_bound": 18_446_744_073_709_551_616.0, + "below_i64_min": below_i64_min, + })); + + let text = transcode_to_text(&binary).unwrap(); + let value: serde_json::Value = serde_json::from_slice(&text).unwrap(); + + assert_eq!(value["i64_min"].as_i64(), Some(i64::MIN)); + assert_eq!( + value["largest_u64_float"].as_u64(), + Some(18_446_744_073_709_549_568) + ); + assert!(value["u64_exclusive_bound"].is_f64()); + assert!(value["below_i64_min"].is_f64()); + } + #[test] fn transcode_passes_text_through_unchanged() { // A text buffer is returned byte-for-byte unchanged. diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 0c6ef56bd24..5e024f605eb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2624,27 +2624,26 @@ impl CosmosDriver { && !operation.is_change_feed() } - /// Whether a feed pipeline negotiated binary and must restore that wire - /// format after its text-based processing stages complete. + /// Whether a feed pipeline requested binary output after text normalization. + /// Text fallback pages are encoded too, preserving one format across a mixed rollout. fn binary_feed_response_restore_applies(operation: &CosmosOperation) -> bool { Self::binary_encoding_applies(operation) && operation.operation_type().is_feed() } - /// Applies request-side binary encoding to an operation: transcodes a text - /// request body to Cosmos binary JSON (an already-binary or empty body is - /// passed through) and advertises binary responses via the - /// `x-ms-cosmos-supported-serialization-formats` header. + /// Applies request-side binary encoding to an operation and advertises + /// binary responses via the supported-serialization-formats header. /// - /// This is schema-agnostic — it operates on the raw body bytes — so a - /// caller that deals only in text JSON gets a binary wire without encoding - /// anything itself. + /// Query specifications remain text because the negotiation header alone + /// produces binary query pages; item bodies retain the schema-agnostic + /// text-to-binary transcode. fn apply_request_binary_encoding( operation: CosmosOperation, ) -> crate::error::Result { - // Transcode a non-empty *text* body to binary. A body that is already - // binary (the SDK's typed fast path) or empty is left in place — no - // clone — so only genuinely text bodies pay the conversion. - let transcoded = match operation.body() { + let encode_body = !matches!( + operation.operation_type(), + crate::models::OperationType::Query | crate::models::OperationType::SqlQuery + ); + let transcoded = match operation.body().filter(|_| encode_body) { Some(body) if !body.is_empty() && !crate::binary_json::is_binary(body) => { Some(crate::binary_json::transcode_to_binary(body).map_err(|e| { crate::error::CosmosError::builder() @@ -3242,8 +3241,7 @@ impl CosmosDriver { operation.operation_type(), crate::models::OperationType::Query | crate::models::OperationType::SqlQuery ) - .then(|| planner::streaming_query_fingerprint(&operation)) - .transpose()?; + .then(|| planner::streaming_query_fingerprint(&operation)); let binary = if Self::binary_encoding_applies(&operation) { self.operation_options_view(options) .binary_encoding() @@ -6466,6 +6464,37 @@ mod tests { ); } + #[test] + fn apply_request_binary_encoding_keeps_query_body_text() { + use crate::models::{OperationType, ResourceType}; + + let container = epk_test_container(r#"{"paths":["/pk"],"version":2}"#); + let text = + br#"{"query":"SELECT * FROM c WHERE c.id = @p","parameters":[{"name":"@p","value":1}]}"# + .to_vec(); + for operation_type in [OperationType::Query, OperationType::SqlQuery] { + let operation = CosmosOperation::new( + operation_type, + crate::models::CosmosResourceReference::from(container.clone()) + .with_resource_type(ResourceType::Document) + .into_feed_reference(), + Some(crate::models::FeedRange::full()), + ) + .with_body(text.clone()); + + let operation = CosmosDriver::apply_request_binary_encoding(operation).unwrap(); + + assert_eq!(operation.body(), Some(text.as_slice())); + assert_eq!( + operation + .request_headers() + .supported_serialization_formats + .as_deref(), + Some("CosmosBinary"), + ); + } + } + #[test] fn apply_request_binary_encoding_errors_on_invalid_text_body() { // A body that is neither binary nor valid JSON surfaces as a diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs index 488e93fc79b..8796cbe2997 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs @@ -278,27 +278,10 @@ pub(crate) fn is_streaming_order_by(info: &QueryInfo) -> bool { } /// Fingerprints the caller's query body and feed scope before request encoding. -pub(crate) fn streaming_query_fingerprint( - operation: &CosmosOperation, -) -> crate::error::Result { - streaming_ordered_merge::query_fingerprint(operation.body(), operation.target()).map_err( - |source| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::SERIALIZATION_REQUEST_BODY_INVALID) - .with_message( - "failed to normalize query body for continuation-token fingerprinting", - ) - .with_source(source) - .build() - }, - ) +pub(crate) fn streaming_query_fingerprint(operation: &CosmosOperation) -> String { + streaming_ordered_merge::query_fingerprint(operation.body(), operation.target()) } -/// Builds a streaming `ORDER BY` pipeline using the operation's current body. -/// -/// Unit tests call this before any request encoding. Production planning uses -/// [`build_streaming_ordered_merge_with_fingerprint`] with the fingerprint -/// captured before binary conversion. #[cfg(test)] pub(crate) async fn build_streaming_ordered_merge( query_plan: &QueryPlan, @@ -306,7 +289,7 @@ pub(crate) async fn build_streaming_ordered_merge( operation: &Arc, resume: Option, ) -> crate::error::Result { - let query_fingerprint = streaming_query_fingerprint(operation)?; + let query_fingerprint = streaming_query_fingerprint(operation); build_streaming_ordered_merge_with_fingerprint( query_plan, topology_provider, @@ -317,6 +300,8 @@ pub(crate) async fn build_streaming_ordered_merge( .await } +/// Builds a streaming `ORDER BY` pipeline with the query fingerprint captured +/// before any optional request encoding can alter the caller's body bytes. pub(crate) async fn build_streaming_ordered_merge_with_fingerprint( query_plan: &QueryPlan, topology_provider: &mut dyn TopologyProvider, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs index bb296ce1227..93e361981dc 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs @@ -202,11 +202,7 @@ pub(crate) fn rewrite_query_body( "query".to_owned(), serde_json::Value::String(rewritten_query.to_owned()), ); - serialize_query_body( - original_body, - &value, - "failed to serialize rewritten query body", - ) + serialize_query_body(&value, "failed to serialize rewritten query body") } /// Inserts the .NET-compatible structured `"resumeFilter"` field into an @@ -239,11 +235,7 @@ pub(crate) fn with_resume_filter( "resumeFilter".to_owned(), order_by::resume_filter_json(resume_values, rid, exclude), ); - serialize_query_body( - body, - &value, - "failed to serialize query body with resume filter", - ) + serialize_query_body(&value, "failed to serialize query body with resume filter") } /// Parses a query operation's JSON body, erroring on a missing body or @@ -252,25 +244,15 @@ fn parse_query_body(body: Option<&[u8]>) -> crate::error::Result, value: &serde_json::Value, error_message: &'static str, ) -> crate::error::Result> { - let text = serde_json::to_vec(value).map_err(|e| body_error(error_message, e))?; - if original_body.is_some_and(crate::binary_json::is_binary) { - crate::binary_json::transcode_to_binary(&text) - .map_err(|e| binary_body_error("failed to transcode rewritten query body to binary", e)) - } else { - Ok(text) - } + serde_json::to_vec(value).map_err(|e| body_error(error_message, e)) } /// One row parsed from a rewritten-envelope backend page, ready for the @@ -613,17 +595,6 @@ fn body_error(message: &'static str, source: serde_json::Error) -> crate::error: .build() } -fn binary_body_error( - message: &'static str, - source: crate::binary_json::BinaryError, -) -> crate::error::CosmosError { - crate::error::CosmosError::builder() - .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) - .with_message(message) - .with_source(source) - .build() -} - fn body_error_msg(message: &'static str) -> crate::error::CosmosError { crate::error::CosmosError::builder() .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs index 7111a3b2172..b96bb85de90 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs @@ -784,13 +784,9 @@ impl PipelineNode for StreamingOrderedMerge { /// the Gateway's rewritten query so a service-side rewrite change does not /// invalidate in-flight tokens. /// -/// Text bodies retain their exact serialized form for compatibility with -/// existing tokens. The planner normally calls this before request encoding; -/// restoring an already-binary caller body to text is a best-effort fallback. -pub(super) fn query_fingerprint( - body: Option<&[u8]>, - scope: Option<&FeedRange>, -) -> crate::binary_json::Result { +/// Bodies retain their exact serialized form for compatibility with existing +/// tokens. The planner calls this before any optional request encoding. +pub(super) fn query_fingerprint(body: Option<&[u8]>, scope: Option<&FeedRange>) -> String { // The body hash is rendered fixed-width first so the two components can // never run together; EPK hex is `[0-9A-F]*`, so neither separator can // occur inside a bound. Bounds render canonically (trailing zero bytes @@ -798,15 +794,7 @@ pub(super) fn query_fingerprint( // backend and other SDKs may hand back a bound with that padding trimmed. // An absent scope hashes as empty, which stays distinct from the // full-container range (`-FF`). - let normalized_body; - let body = match body { - Some(body) if crate::binary_json::is_binary(body) => { - normalized_body = crate::binary_json::transcode_to_text(body)?; - normalized_body.as_slice() - } - Some(body) => body, - None => &[], - }; + let body = body.unwrap_or_default(); let body_hash = crate::models::murmur_hash::murmurhash3_128(body, 0); let scope = match scope { Some(range) => format!( @@ -816,13 +804,13 @@ pub(super) fn query_fingerprint( ), None => String::new(), }; - Ok(format!( + format!( "{:032x}", crate::models::murmur_hash::murmurhash3_128( format!("{body_hash:032x}:{scope}").as_bytes(), 0 ) - )) + ) } /// Builds the child streams needed to cover `scope`, given its topology @@ -2930,10 +2918,10 @@ mod tests { #[test] fn query_fingerprint_distinguishes_feed_scope() { let body = br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#; - let full = query_fingerprint(Some(body), Some(&FeedRange::full())).unwrap(); - let left = query_fingerprint(Some(body), Some(&range("", "80"))).unwrap(); - let right = query_fingerprint(Some(body), Some(&range("80", "FF"))).unwrap(); - let unscoped = query_fingerprint(Some(body), None).unwrap(); + let full = query_fingerprint(Some(body), Some(&FeedRange::full())); + let left = query_fingerprint(Some(body), Some(&range("", "80"))); + let right = query_fingerprint(Some(body), Some(&range("80", "FF"))); + let unscoped = query_fingerprint(Some(body), None); assert_ne!(full, left); assert_ne!(full, right); @@ -2948,8 +2936,8 @@ mod tests { fn query_fingerprint_separators_cannot_collide() { // Both scopes render as `408080` once the bound separator is dropped. assert_ne!( - query_fingerprint(None, Some(&range("40", "8080"))).unwrap(), - query_fingerprint(None, Some(&range("4080", "80"))).unwrap(), + query_fingerprint(None, Some(&range("40", "8080"))), + query_fingerprint(None, Some(&range("4080", "80"))), ); } @@ -2960,12 +2948,12 @@ mod tests { #[test] fn query_fingerprint_ignores_trailing_zero_padding_in_scope() { assert_eq!( - query_fingerprint(None, Some(&range("", "80"))).unwrap(), - query_fingerprint(None, Some(&range("", "8000"))).unwrap(), + query_fingerprint(None, Some(&range("", "80"))), + query_fingerprint(None, Some(&range("", "8000"))), ); assert_eq!( - query_fingerprint(None, Some(&range("40", "80"))).unwrap(), - query_fingerprint(None, Some(&range("400000", "8000"))).unwrap(), + query_fingerprint(None, Some(&range("40", "80"))), + query_fingerprint(None, Some(&range("400000", "8000"))), ); } @@ -2975,7 +2963,7 @@ mod tests { fn parameterized_text_query_fingerprint_matches_historical_value() { let body = br#"{"query":"SELECT * FROM c WHERE c.rank >= @min ORDER BY c.rank","parameters":[{"name":"@min","value":1}]}"#; assert_eq!( - query_fingerprint(Some(body), Some(&range("", "80"))).unwrap(), + query_fingerprint(Some(body), Some(&range("", "80"))), "b84b9c269862dcd73781038d90add3be", ); } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs index b7867747fb4..5e1bdbf2aa9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs @@ -1522,8 +1522,8 @@ impl SubStatusCode { pub const CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID: SubStatusCode = SubStatusCode(20215); - /// A streaming `ORDER BY` plan is missing the request fingerprint captured - /// before wire encoding (20216). Indicates an internal planner invariant violation. + /// A streaming `ORDER BY` plan is missing its request fingerprint (20216). + /// Unreachable by construction; retained to avoid an FFI-boundary panic if that invariant regresses. pub const CLIENT_STREAMING_ORDER_BY_FINGERPRINT_MISSING: SubStatusCode = SubStatusCode(20216); // ----- 20300-20349: SDK-detected service contract violations ----- diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs index 66cd3907456..14846ec5e00 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs @@ -2428,18 +2428,12 @@ fn parse_query_spec( request_body: &[u8], start: Instant, ) -> Result<(String, Vec<(String, serde_json::Value)>), AsyncRawResponse> { - let spec: QuerySpec = if crate::binary_json::is_binary(request_body) { - crate::binary_json::from_slice(request_body) - .map_err(|e| format!("invalid binary query body: {e}")) - } else { - serde_json::from_slice(request_body).map_err(|e| format!("invalid text query body: {e}")) - } - .map_err(|e| { + let spec: QuerySpec = serde_json::from_slice(request_body).map_err(|e| { error_response( StatusCode::BadRequest, None, "BadRequest", - &format!("Invalid query JSON body: {e}"), + &format!("Invalid text query JSON body: {e}"), 0.0, "", start, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/response.rs index b5b91799c65..dbf5707bece 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/response.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/response.rs @@ -227,13 +227,16 @@ impl ResponseBuilder { /// auto-detects from the first byte, so the `Content-Type` stays /// `application/json` either way (mirroring the real service). pub fn with_value_body(mut self, body: &serde_json::Value, binary: bool) -> Self { + let body = service_number_model(body.clone()); self.body = if binary { - crate::binary_json::encode(body) + crate::binary_json::encode(&body) } else { // The emulator owns these `Value`s, so a serialization failure is a // bug in the emulator — fail loudly rather than emit an empty body // that would mask the defect downstream. - serde_json::to_vec(body).expect("emulator response body must serialize to JSON") + let mut body = body; + crate::binary_json::normalize_integral_floats(&mut body); + serde_json::to_vec(&body).expect("emulator response body must serialize to JSON") }; self } @@ -249,6 +252,28 @@ impl ResponseBuilder { } } +fn service_number_model(value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Number(number) => serde_json::Number::from_f64( + number + .as_f64() + .expect("JSON numbers are representable as finite f64"), + ) + .map(serde_json::Value::Number) + .expect("finite f64 is a valid JSON number"), + serde_json::Value::Array(values) => { + serde_json::Value::Array(values.into_iter().map(service_number_model).collect()) + } + serde_json::Value::Object(values) => serde_json::Value::Object( + values + .into_iter() + .map(|(key, value)| (key, service_number_model(value))) + .collect(), + ), + value => value, + } +} + /// Creates a success response with a JSON body. pub(crate) fn success_response( status: StatusCode, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs index 9ad6364f060..6a19fb20de4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/response_body.rs @@ -199,20 +199,20 @@ impl ResponseBody { /// Each [`Items`](Self::Items) slice is encoded independently, so every /// resulting item carries the `0x80` preamble required for auto-detection. pub(crate) fn transcode_to_binary(self) -> crate::error::Result { - fn convert(bytes: &Bytes) -> crate::error::Result { - if crate::binary_json::is_binary(bytes) { - return Ok(bytes.clone()); + fn convert(bytes: Bytes) -> crate::error::Result { + if crate::binary_json::is_binary(&bytes) { + return Ok(bytes); } - crate::binary_json::transcode_to_binary(bytes) + crate::binary_json::transcode_to_binary(&bytes) .map(Bytes::from) .map_err(|e| invalid_body_error("failed to transcode response body to binary", e)) } match self { Self::NoPayload => Ok(Self::NoPayload), - Self::Bytes(bytes) => convert(&bytes).map(Self::Bytes), + Self::Bytes(bytes) => convert(bytes).map(Self::Bytes), Self::Items(items) => items - .iter() + .into_iter() .map(convert) .collect::>>() .map(Self::Items), diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/binary_response_format.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/binary_response_format.rs index 2fb78ad8d55..5b6ab4c507d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/binary_response_format.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/binary_response_format.rs @@ -94,10 +94,10 @@ async fn enabled_default_negotiation_yields_binary_response() { "response body must be detected as binary", ); - // And it decodes back to the stored document. + // Binary responses preserve the service's double-only number representation. let decoded: serde_json::Value = binary_json::decode(&raw).unwrap(); assert_eq!(decoded["id"], "bin-1"); - assert_eq!(decoded["value"], 42); + assert_eq!(decoded["value"], 42.0); } /// Explicit `JsonText`-only negotiation: the response body is **text**, even diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index bc6fdb49025..41c60df7b48 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -335,6 +335,33 @@ fn sorted(mut values: Vec) -> Vec { text } +fn normalize_integral_floats(value: &mut serde_json::Value) { + match value { + serde_json::Value::Number(number) if number.is_f64() => { + let float = number.as_f64().expect("finite JSON number"); + if float.fract() == 0.0 { + if float >= i64::MIN as f64 && float < 0.0 { + *number = serde_json::Number::from(float as i64); + } else if float >= 0.0 && float < 18_446_744_073_709_551_616.0 { + *number = serde_json::Number::from(float as u64); + } + } + } + serde_json::Value::Array(values) => { + values.iter_mut().for_each(normalize_integral_floats); + } + serde_json::Value::Object(values) => { + values.values_mut().for_each(normalize_integral_floats); + } + _ => {} + } +} + +fn normalize_binary_values(mut values: Vec) -> Vec { + values.iter_mut().for_each(normalize_integral_floats); + values +} + fn assert_matches_expected(scenario: &Scenario, actual: Vec) { if !scenario.expected_ids.is_empty() { let mut ids: Vec = actual @@ -772,6 +799,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { .await; let planned_text_request_modes = recorder.take(); + let binary = normalize_binary_values(binary); if query.distinct_type == "Ordered" { assert_eq!(binary, text, "{}", query.text); assert_eq!(binary_as_text, text, "{}", query.text); From 8b3e95e0f6a4020a31437101ae3ac9b05e30f07c Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 13 Aug 2026 10:53:45 -0700 Subject: [PATCH 8/8] Harden binary feed page splitting Pin binary envelope normalization at the shared feed-splitting boundary and prevent format parity assertions from passing without a classifiable page. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 17d20364-4ec8-4cef-9aa2-b9bc87c18430 --- .../src/driver/dataflow/skip_take_page.rs | 22 +++++++++++++++---- .../in_memory_emulator_tests/distinct.rs | 17 ++++++++++++-- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/skip_take_page.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/skip_take_page.rs index 9e187c968b6..6c8d8b06315 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/skip_take_page.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/skip_take_page.rs @@ -12,8 +12,9 @@ //! Rather than re-serialize a trimmed envelope back to bytes (which the calling //! SDK would then have to re-parse), these helpers split the `Documents` array //! into a list of per-document [`Bytes`] — each a zero-copy -//! [`slice_ref`](bytes::Bytes::slice_ref) of the original page buffer — and -//! trim that list. [`SkipTake`] emits the surviving slices directly as a +//! [`slice_ref`](bytes::Bytes::slice_ref) of the text page buffer — and trim +//! that list. Binary pages are first normalized into one text buffer. +//! [`SkipTake`] emits the surviving slices directly as a //! [`ResponseBody::Items`](crate::models::ResponseBody::Items) body. //! //! [`SkipTake`]: super::SkipTake @@ -29,7 +30,7 @@ pub(crate) struct SkipTakeOutcome { /// Number of documents kept (equal to `items.len()`). pub emitted: u64, /// The surviving per-document payloads, each an unmodified slice of the - /// original page bytes. + /// normalized page bytes. pub items: Vec, } @@ -43,7 +44,8 @@ struct RawQueryPage<'a> { } /// Splits a backend query-page envelope into a list of per-document payloads, -/// each a zero-copy [`slice_ref`](bytes::Bytes::slice_ref) of `body`. +/// each a zero-copy [`slice_ref`](bytes::Bytes::slice_ref) of the normalized +/// text buffer. Text pages retain their original allocation. /// /// An empty (`NoPayload`) body is treated as a zero-document page. pub(crate) fn split_feed_envelope(body: &Bytes) -> crate::error::Result> { @@ -167,6 +169,18 @@ mod tests { assert!(out.items.is_empty()); } + #[test] + fn binary_envelope_is_normalized_before_splitting() { + let body = Bytes::from(crate::binary_json::encode(&serde_json::json!({ + "Documents": [{"id": 1.0}, {"id": 2.0}], + "_count": 2.0, + }))); + + let items = split_feed_envelope(&body).unwrap(); + + assert_eq!(raw(&items), vec![r#"{"id":1}"#, r#"{"id":2}"#]); + } + #[test] fn preserves_numeric_precision() { // A high-precision number must survive the split unchanged, byte-for-byte. diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index 75815119f94..fe2fcf3a55d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -286,8 +286,8 @@ async fn drain_query_with_options( { // The pipeline emits pre-split `Items`, each slice carrying its own // `0x80` preamble; a single-partition passthrough page is one `Bytes` - // envelope. An empty page has no bytes to classify, so it inherits the - // text answer and is filtered out of the all/none assertions below. + // envelope. An empty page has no bytes to classify, so it has no format + // and is filtered out of the all/none assertions below. let is_binary = match response.body() { ResponseBody::Bytes(bytes) => azure_data_cosmos_driver::binary_json::is_binary(bytes), ResponseBody::Items(items) => items @@ -831,6 +831,19 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { assert_eq!(sorted(binary), sorted(text.clone()), "{}", query.text); assert_eq!(sorted(binary_as_text), sorted(text), "{}", query.text); } + assert!( + [ + &text_formats, + &binary_formats, + &binary_as_text_formats, + &planned_binary_formats, + &planned_text_formats, + ] + .iter() + .all(|formats| !formats.is_empty()), + "each mode must emit at least one classifiable page: {}", + query.text + ); assert!( text_formats.iter().all(|is_binary| !is_binary), "{}",