From 7593a4b50bed28e6c4500f55c05e7547eb6a6bf6 Mon Sep 17 00:00:00 2001 From: andrijast-db Date: Mon, 10 Aug 2026 13:47:32 +0000 Subject: [PATCH 1/3] initial --- rust/NEXT_CHANGELOG.md | 27 ++ rust/README.md | 39 +++ rust/examples/README.md | 1 + rust/examples/proto/Cargo.toml | 4 + rust/examples/proto/README.md | 53 +++ rust/examples/proto/dynamic/from_uc.rs | 92 +++++ rust/sdk/src/dynamic_proto.rs | 6 +- rust/sdk/src/errors.rs | 12 + rust/sdk/src/lib.rs | 1 + rust/sdk/src/sdk.rs | 43 +++ rust/sdk/src/uc_schema.rs | 444 +++++++++++++++++++++++++ rust/tests/Cargo.toml | 4 + rust/tests/src/mock_uc.rs | 268 +++++++++++++++ rust/tests/src/uc_schema_tests.rs | 394 ++++++++++++++++++++++ 14 files changed, 1386 insertions(+), 2 deletions(-) create mode 100644 rust/examples/proto/dynamic/from_uc.rs create mode 100644 rust/sdk/src/uc_schema.rs create mode 100644 rust/tests/src/mock_uc.rs create mode 100644 rust/tests/src/uc_schema_tests.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 8fde35b3..56390e06 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -6,6 +6,29 @@ ### New Features and Improvements +- Dynamic protobuf streams can resolve their schema from Unity Catalog, so a + runtime descriptor no longer has to be assembled by hand. Fetch the descriptor + from the live table metadata and pass it to the existing `.dynamic_proto(...)` + selector: + + ```rust + let descriptor = sdk + .fetch_message_descriptor("catalog.schema.table", client_id, client_secret) + .await?; + + let stream = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth(client_id, client_secret) + .dynamic_proto(descriptor) + .build() + .await?; + ``` + + `ZerobusSdk::fetch_message_descriptor()` uses the SDK's configured + `unity_catalog_url`; the underlying `uc_schema` module takes the endpoint + directly. The fetch needs OAuth credentials able to read the table's metadata. + ### Bug Fixes - Arrow Flight acknowledgment deadlines are pending-relative: no timer runs @@ -37,3 +60,7 @@ ### Deprecations ### API Changes + +- Added `ZerobusSdk::fetch_message_descriptor()`, the `uc_schema` module + (`fetch_message_descriptor`, `fetch_table_schema`), and the + `ZerobusError::SchemaFetchError { message, retryable }` variant. All additive. diff --git a/rust/README.md b/rust/README.md index f5f2fc0e..35e8f578 100644 --- a/rust/README.md +++ b/rust/README.md @@ -573,6 +573,43 @@ stream.flush().await?; // wait once for all pending acknowledgments On the wire this is identical to `.compiled_proto(...)`; the difference is that records are built dynamically rather than from a generated struct. See the [`dynamic_proto`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/dynamic_proto/) module and the `proto_dynamic_single` example for details. +##### Fetching the Schema from Unity Catalog + +Rather than assembling the columns yourself, let the SDK read the table's schema from Unity Catalog. `fetch_message_descriptor` resolves the descriptor from the live table metadata; pass it to `.dynamic_proto(...)` as usual: + +```rust +// Fetch the descriptor once from Unity Catalog (uses the SDK's unity_catalog_url). +let descriptor = sdk + .fetch_message_descriptor("catalog.schema.orders", client_id, client_secret) + .await?; + +// Inspect it if the columns are unknown to the program... +for field in descriptor.fields() { + println!("{} ({:?})", field.name(), field.kind()); +} + +// ...then plug it into the ordinary dynamic-proto selector. Cloning a descriptor +// is cheap (Arc-backed), so one fetch can serve many streams. +let mut stream = sdk + .stream_builder().table("catalog.schema.orders") + .oauth(client_id, client_secret) + .dynamic_proto(descriptor) + .build() + .await?; + +// Records are built exactly as above — `new_record()` uses the fetched schema. +for i in 0..100_000i64 { + let mut record = stream.new_record()?; + record.set("id", i)?.set("customer_name", "Alice Smith")?; + let _offset = stream.ingest_record_offset(ProtoBytes(record.encode()?)).await?; // queue only +} +stream.flush().await?; // wait once for all pending acknowledgments +``` + +The fetch needs OAuth credentials able to read the table's metadata (they are presented to the Unity Catalog REST API) and `unity_catalog_url` on the SDK builder. For direct control over the endpoint — outside an `SDK`, or against a different workspace — call `uc_schema::fetch_message_descriptor(unity_catalog_url, table, client_id, client_secret)`. + +The fetched schema is a snapshot. If the table changes afterwards, the server rejects stream creation with `ZerobusError::InvalidSchema`; re-fetch and rebuild to pick the change up. A failed fetch surfaces as `ZerobusError::SchemaFetchError`, whose `retryable` flag is set for transport errors and 5xx/429 responses. Note that `DATE` and `TIMESTAMP` columns map to integers (days and microseconds since the Unix epoch) — see the [`schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/schema/) module for the full type mapping, and [`uc_schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/uc_schema/) for the fetch API. + Setters can be called in any order. The builder validates at `build()` time that both authentication and format have been configured. ### 5. Ingest Data @@ -1000,6 +1037,7 @@ The `examples/` directory contains four working examples covering different seri | `proto/compiled/batch.rs` | Protocol Buffers | Batch | `cargo run -p rust-examples-proto --example proto_compiled_batch` | | `proto/dynamic/single.rs` | Protocol Buffers (runtime schema) | Single-record | `cargo run -p rust-examples-proto --example proto_dynamic_single` | | `proto/dynamic/batch.rs` | Protocol Buffers (runtime schema) | Batch | `cargo run -p rust-examples-proto --example proto_dynamic_batch` | +| `proto/dynamic/from_uc.rs` | Protocol Buffers (schema fetched from Unity Catalog) | Single-record | `cargo run -p rust-examples-proto --example proto_dynamic_from_uc` | Check [`examples/README.md`](https://github.com/databricks/zerobus-sdk/blob/main/rust/examples/README.md) for setup instructions and detailed comparisons. @@ -1273,6 +1311,7 @@ cargo run -p rust-examples-proto --example proto_compiled_batch # Build and run Protocol Buffers dynamic-schema examples cargo run -p rust-examples-proto --example proto_dynamic_single cargo run -p rust-examples-proto --example proto_dynamic_batch +cargo run -p rust-examples-proto --example proto_dynamic_from_uc ``` ## Community and Contributing diff --git a/rust/examples/README.md b/rust/examples/README.md index 5c796f57..bfd69e5e 100644 --- a/rust/examples/README.md +++ b/rust/examples/README.md @@ -42,6 +42,7 @@ The SDK supports two serialization formats and two ingestion methods: | [Proto Compiled Batch](proto/README.md#compiled-batch-example) | Protocol Buffers | Batch | `cargo run -p rust-examples-proto --example proto_compiled_batch` | | [Proto Dynamic](proto/README.md#dynamic-schema-example) | Protocol Buffers | Single-record (runtime schema) | `cargo run -p rust-examples-proto --example proto_dynamic_single` | | [Proto Dynamic Batch](proto/README.md#dynamic-batch) | Protocol Buffers | Batch (runtime schema) | `cargo run -p rust-examples-proto --example proto_dynamic_batch` | +| [Proto Dynamic from UC](proto/README.md#dynamic-schema-from-unity-catalog) | Protocol Buffers | Single-record (schema fetched from Unity Catalog) | `cargo run -p rust-examples-proto --example proto_dynamic_from_uc` | | [Arrow](arrow/README.md) | Arrow Flight (Beta) | `RecordBatch` | `cargo run -p example_arrow` | ## Prerequisites diff --git a/rust/examples/proto/Cargo.toml b/rust/examples/proto/Cargo.toml index 8e2e28f3..b629073e 100644 --- a/rust/examples/proto/Cargo.toml +++ b/rust/examples/proto/Cargo.toml @@ -16,6 +16,10 @@ path = "compiled/single.rs" name = "proto_dynamic_batch" path = "dynamic/batch.rs" +[[example]] +name = "proto_dynamic_from_uc" +path = "dynamic/from_uc.rs" + [[example]] name = "proto_dynamic_single" path = "dynamic/single.rs" diff --git a/rust/examples/proto/README.md b/rust/examples/proto/README.md index 163823d5..38261343 100644 --- a/rust/examples/proto/README.md +++ b/rust/examples/proto/README.md @@ -17,6 +17,7 @@ This directory contains examples demonstrating Protocol Buffers-based data inges - [Running the Example](#running-the-example-2) - [Code Highlights](#code-highlights-2) - [Dynamic Batch](#dynamic-batch) + - [Dynamic Schema from Unity Catalog](#dynamic-schema-from-unity-catalog) - [Adapting for Your Custom Table](#adapting-for-your-custom-table) - [Generate Schema Files](#generate-schema-files) - [Update main.rs](#update-mainrs) @@ -43,6 +44,7 @@ The examples are grouped by how the protobuf schema is obtained: are built field-by-field with `DynamicRecord`. - **`dynamic/single.rs`** - Build the descriptor in code and ingest dynamic records one at a time - **`dynamic/batch.rs`** - Ingest multiple dynamic records at once using `ingest_records_offset()` + - **`dynamic/from_uc.rs`** - Fetch the schema from Unity Catalog with `fetch_message_descriptor` and feed it to `.dynamic_proto(...)`, so no columns are hardcoded ## Three Ways to Pass Data @@ -280,6 +282,57 @@ if let Some(offset) = stream.ingest_records_offset(batch).await? { stream.flush().await?; ``` +### Dynamic Schema from Unity Catalog + +`dynamic/from_uc.rs` goes one step further: instead of assembling the columns in code, +`fetch_message_descriptor` reads the table's schema from Unity Catalog and the resolved +descriptor is handed to the ordinary `.dynamic_proto(...)` selector. Nothing about the +schema is hardcoded, so the same program works against any table the credentials can read: + +```bash +cargo run -p rust-examples-proto --example proto_dynamic_from_uc +``` + +```rust +// `unity_catalog_url` is required — it is where the schema is fetched from. +let sdk = ZerobusSdk::builder() + .endpoint(SERVER_ENDPOINT) + .unity_catalog_url(DATABRICKS_WORKSPACE_URL) + .build()?; + +// Fetch the descriptor from the live table metadata. +let descriptor = sdk + .fetch_message_descriptor(TABLE_NAME, DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .await?; + +// The fetched schema can be inspected when the columns are genuinely unknown. +for field in descriptor.fields() { + println!(" {} ({:?})", field.name(), field.kind()); +} + +// Plug it into the same builder used for a hand-built descriptor. +let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .dynamic_proto(descriptor) + .build() + .await?; + +// Records are built exactly as in the examples above. +for i in 0..1_000i64 { + let mut record = stream.new_record()?; + record.set("id", i)?.set("customer_name", "Alice Smith")?; + stream.ingest_record_offset(ProtoBytes(record.encode()?)).await?; // queue only +} +stream.flush().await?; // wait once for all pending acks +``` + +This needs `.oauth(...)` credentials (they are presented to the Unity Catalog REST +API) and `unity_catalog_url` on the SDK builder. `sdk.fetch_message_descriptor(...)` +uses that configured URL; for direct control over the endpoint, call +`uc_schema::fetch_message_descriptor(unity_catalog_url, table, client_id, client_secret)`. + ## Adapting for Your Custom Table To use your own table, you need to generate schema files and update the example code. diff --git a/rust/examples/proto/dynamic/from_uc.rs b/rust/examples/proto/dynamic/from_uc.rs new file mode 100644 index 00000000..7448c4b6 --- /dev/null +++ b/rust/examples/proto/dynamic/from_uc.rs @@ -0,0 +1,92 @@ +//! Dynamic protobuf ingestion with the schema fetched from Unity Catalog. +//! +//! Unlike `dynamic/single.rs`, which builds the descriptor in code, this fetches +//! it with `fetch_message_descriptor` and feeds it to the usual `.dynamic_proto(...)`. +//! +//! Throughput: ingest in a loop, then `flush()` once — never wait per record. + +use std::error::Error; + +use databricks_zerobus_ingest_sdk::{ProtoBytes, ZerobusSdk, ZerobusStream}; + +// Change constants to match your data. +const TABLE_NAME: &str = ""; +const DATABRICKS_CLIENT_ID: &str = ""; +const DATABRICKS_CLIENT_SECRET: &str = ""; + +// Uncomment the appropriate lines for your cloud. + +// For AWS: +const DATABRICKS_WORKSPACE_URL: &str = "https://.cloud.databricks.com"; +const SERVER_ENDPOINT: &str = "https://.zerobus..cloud.databricks.com"; + +// For Azure: +// const DATABRICKS_WORKSPACE_URL: &str = "https://.azuredatabricks.net"; +// const SERVER_ENDPOINT: &str = "https://.zerobus..azuredatabricks.net"; + +#[tokio::main] +async fn main() -> Result<(), Box> { + // `unity_catalog_url` is required: it is where the schema is fetched from. + let sdk_handle = ZerobusSdk::builder() + .endpoint(SERVER_ENDPOINT) + .unity_catalog_url(DATABRICKS_WORKSPACE_URL) + .build()?; + + // Descriptor from live table metadata — no columns or `.proto` needed up front. + let descriptor = sdk_handle + .fetch_message_descriptor(TABLE_NAME, DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .await?; + + // The fetched schema can be inspected when the columns are unknown to the program. + println!("Fetched schema '{}':", descriptor.name()); + for field in descriptor.fields() { + println!(" {} ({:?})", field.name(), field.kind()); + } + + let mut stream = sdk_handle + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .dynamic_proto(descriptor) + .max_inflight_requests(100) + .build() + .await?; + + ingest_records(&mut stream).await?; + + stream.close().await?; + println!("Stream closed successfully"); + + Ok(()) +} + +async fn ingest_records(stream: &mut ZerobusStream) -> Result<(), Box> { + // (customer_name, quantity, price) — adjust the field names below to your table. + let orders = [ + ("Alice Smith", 2i32, 25.99f64), + ("Bob Johnson", 1, 89.99), + ("Carol Williams", 3, 45.00), + ]; + + for (i, (customer_name, quantity, price)) in orders.iter().enumerate() { + // set()'s value must match the column's proto type (BIGINT -> i64, INT -> i32). + let mut record = stream.new_record()?; + record + .set("id", i as i64)? + .set("customer_name", *customer_name)? + .set("quantity", *quantity)? + .set("price", *price)?; + + // encode() enforces proto2 required fields; queues without waiting for the ack. + let offset_id = stream + .ingest_record_offset(ProtoBytes(record.encode()?)) + .await?; + println!("Record {i} queued with offset ID: {offset_id}"); + } + + // Wait once for all pending acks — not after each ingest. + stream.flush().await?; + println!("All records acknowledged"); + + Ok(()) +} diff --git a/rust/sdk/src/dynamic_proto.rs b/rust/sdk/src/dynamic_proto.rs index 051a49d7..808f564e 100644 --- a/rust/sdk/src/dynamic_proto.rs +++ b/rust/sdk/src/dynamic_proto.rs @@ -9,8 +9,10 @@ //! //! Obtain the [`MessageDescriptor`] from [`message_descriptor`], which resolves a //! [`prost_types::DescriptorProto`] (built with -//! [`crate::schema::descriptor_from_uc_columns`] or fetched from Unity Catalog), -//! or from your own [`prost_reflect::DescriptorPool`]. +//! [`crate::schema::descriptor_from_uc_columns`]), or from your own +//! [`prost_reflect::DescriptorPool`]. To read the schema straight from Unity +//! Catalog instead, see [`crate::uc_schema`] (or +//! [`ZerobusSdk::fetch_message_descriptor`](crate::ZerobusSdk::fetch_message_descriptor)). //! //! Ingest in a loop, then `flush()` once — never wait per record. //! diff --git a/rust/sdk/src/errors.rs b/rust/sdk/src/errors.rs index 79def003..3e528ade 100644 --- a/rust/sdk/src/errors.rs +++ b/rust/sdk/src/errors.rs @@ -131,6 +131,15 @@ pub enum ZerobusError { /// Returned when OAuth token fetching fails due to network or server errors. #[error("Token fetch failed: {0}")] TokenFetchError(String), + /// Returned when resolving a table's schema from Unity Catalog failed (see + /// [`crate::uc_schema`]). `retryable` is set for transport errors and 5xx/429 + /// responses, unset for a rejected request or an unusable body. + /// + /// `#[non_exhaustive]` so further detail (e.g. the HTTP status) can be added + /// later without a breaking change. + #[error("Failed to fetch table schema from Unity Catalog: {message}.")] + #[non_exhaustive] + SchemaFetchError { message: String, retryable: bool }, } /// List of gRPC status codes that indicate unretriable errors. @@ -210,6 +219,9 @@ impl ZerobusError { ZerobusError::InvalidStateError(_) => false, ZerobusError::ConnectionTimeout(_) => true, ZerobusError::TokenFetchError(_) => true, + // Classified at the fetch site, where the HTTP status and transport + // error are known. + ZerobusError::SchemaFetchError { retryable, .. } => *retryable, } } diff --git a/rust/sdk/src/lib.rs b/rust/sdk/src/lib.rs index f9d2bfdf..75bf0056 100644 --- a/rust/sdk/src/lib.rs +++ b/rust/sdk/src/lib.rs @@ -57,6 +57,7 @@ mod stream_configuration; pub mod stream_options; mod tls_config; mod token_cache; +pub mod uc_schema; pub use builder::{StreamBuilder, ZerobusSdkBuilder}; pub use callbacks::AckCallback; diff --git a/rust/sdk/src/sdk.rs b/rust/sdk/src/sdk.rs index 9cfeb723..75694f0a 100644 --- a/rust/sdk/src/sdk.rs +++ b/rust/sdk/src/sdk.rs @@ -117,6 +117,49 @@ impl ZerobusSdk { StreamBuilder::new(self) } + /// Fetch `table_name`'s schema from Unity Catalog and resolve it to a + /// [`MessageDescriptor`](crate::MessageDescriptor), using this SDK's configured + /// `unity_catalog_url`. Pass the result to + /// [`dynamic_proto`](StreamBuilder::dynamic_proto). For direct control over the + /// endpoint, use [`uc_schema::fetch_message_descriptor`](crate::uc_schema::fetch_message_descriptor). + /// + /// # Examples + /// + /// ```no_run + /// # use databricks_zerobus_ingest_sdk::ZerobusSdk; + /// # async fn example(sdk: &ZerobusSdk) -> Result<(), Box> { + /// let descriptor = sdk + /// .fetch_message_descriptor("catalog.schema.table", "client-id", "client-secret") + /// .await?; + /// let stream = sdk + /// .stream_builder() + /// .table("catalog.schema.table") + /// .oauth("client-id", "client-secret") + /// .dynamic_proto(descriptor) + /// .build() + /// .await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// See [`uc_schema::fetch_message_descriptor`](crate::uc_schema::fetch_message_descriptor). + pub async fn fetch_message_descriptor( + &self, + table_name: &str, + client_id: &str, + client_secret: &str, + ) -> ZerobusResult { + crate::uc_schema::fetch_message_descriptor( + &self.unity_catalog_url, + table_name, + client_id, + client_secret, + ) + .await + } + /// Creates a new SDK instance with explicit configuration. /// /// This is used internally by the builder pattern. `sdk_identifier` is the diff --git a/rust/sdk/src/uc_schema.rs b/rust/sdk/src/uc_schema.rs new file mode 100644 index 00000000..3d9300cc --- /dev/null +++ b/rust/sdk/src/uc_schema.rs @@ -0,0 +1,444 @@ +//! Fetch a table's schema from Unity Catalog and resolve it to a protobuf +//! [`MessageDescriptor`], for [`dynamic_proto`](crate::StreamBuilder::dynamic_proto) +//! when the schema is only known at runtime. The runtime counterpart to +//! [`crate::schema`]: reads `GET /api/2.1/unity-catalog/tables/{full_name}` and +//! converts it via [`descriptor_from_uc_schema`]. +//! +//! Fetching is a separate step, so the descriptor can be inspected and reused +//! across streams (cloning it is cheap — Arc-backed). +//! [`ZerobusSdk::fetch_message_descriptor`](crate::ZerobusSdk::fetch_message_descriptor) +//! wraps [`fetch_message_descriptor`] with the SDK's `unity_catalog_url`: +//! +//! ```no_run +//! # use databricks_zerobus_ingest_sdk::ZerobusSdk; +//! # async fn example(sdk: &ZerobusSdk) -> Result<(), Box> { +//! let descriptor = sdk +//! .fetch_message_descriptor("catalog.schema.table", "client-id", "client-secret") +//! .await?; +//! let stream = sdk +//! .stream_builder() +//! .table("catalog.schema.table") +//! .oauth("client-id", "client-secret") +//! .dynamic_proto(descriptor) +//! .build() +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Columns map per [`crate::schema`] (note `DATE`/`TIMESTAMP` become integers, +//! not `google.protobuf.Timestamp`). The descriptor is a snapshot: if the table +//! changes afterwards the server rejects stream creation with +//! [`ZerobusError::InvalidSchema`], so re-fetch and rebuild. + +use std::time::Duration; + +use prost_reflect::MessageDescriptor; +use reqwest::StatusCode; +use tracing::{debug, warn}; + +use crate::dynamic_proto::message_descriptor; +use crate::schema::{descriptor_from_uc_schema, UcTableSchema}; +use crate::{ZerobusError, ZerobusResult}; + +/// Deadline for a single fetch (token mint plus schema read). +const FETCH_TIMEOUT: Duration = Duration::from_secs(30); + +/// Cap on a buffered response body. Both responses are small (a token, a column +/// list); the bound guards against an unexpected reply. Rejected outright, not +/// truncated, so we never act on a partial schema. +const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; + +/// Fetch `table_name`'s schema from Unity Catalog and resolve it to a +/// [`MessageDescriptor`] for [`dynamic_proto`](crate::StreamBuilder::dynamic_proto). +/// +/// `unity_catalog_url` is the workspace URL; `table_name` is `catalog.schema.table`. +/// The OAuth credentials must be able to read the table's metadata. Makes two HTTP +/// requests (token mint + table read), so reuse the result — cloning it is cheap +/// (Arc-backed). +/// +/// # Errors +/// +/// - [`ZerobusError::InvalidTableName`] if `table_name` is not `catalog.schema.table`. +/// - [`ZerobusError::InvalidUCEndpointError`] if `unity_catalog_url` is unusable. +/// - [`ZerobusError::SchemaFetchError`] if the token mint or table read fails +/// (`retryable` set for transport errors and 5xx/429). +/// - [`ZerobusError::InvalidArgument`] if the schema has no protobuf +/// representation (e.g. an unsupported column type). +pub async fn fetch_message_descriptor( + unity_catalog_url: &str, + table_name: &str, + client_id: &str, + client_secret: &str, +) -> ZerobusResult { + let schema = + fetch_table_schema(unity_catalog_url, table_name, client_id, client_secret).await?; + descriptor_from_schema(&schema) +} + +/// Fetch `table_name`'s raw Unity Catalog schema, without converting it to a +/// protobuf descriptor. Useful to inspect the columns directly; most callers +/// want [`fetch_message_descriptor`]. +/// +/// # Errors +/// +/// The same as [`fetch_message_descriptor`], minus the descriptor conversion. +pub async fn fetch_table_schema( + unity_catalog_url: &str, + table_name: &str, + client_id: &str, + client_secret: &str, +) -> ZerobusResult { + validate_table_name(table_name)?; + let base = normalize_endpoint(unity_catalog_url)?; + + let client = reqwest::Client::builder() + .timeout(FETCH_TIMEOUT) + .build() + .map_err(|e| fetch_error(format!("failed to build HTTP client: {e}"), false))?; + + debug!(table = %table_name, "fetching UC table schema"); + let token = mint_metadata_token(&client, &base, client_id, client_secret).await?; + let schema = get_table(&client, &base, &token, table_name).await?; + + if schema.columns.is_empty() { + return Err(fetch_error( + format!("Unity Catalog returned no columns for table '{table_name}'"), + false, + )); + } + Ok(schema) +} + +/// Convert a fetched [`UcTableSchema`] into a [`MessageDescriptor`], mapping a +/// conversion failure to [`ZerobusError::InvalidArgument`]. +fn descriptor_from_schema(schema: &UcTableSchema) -> ZerobusResult { + let descriptor = descriptor_from_uc_schema(schema).map_err(|e| { + ZerobusError::InvalidArgument(format!( + "cannot convert Unity Catalog schema for table '{}' to a protobuf descriptor: {e}", + schema.name + )) + })?; + message_descriptor(&descriptor) +} + +/// Mint an OAuth token for reading table metadata. +/// +/// Separate from [`crate::DefaultTokenFactory`], which mints an ingestion token +/// (`zerobusDirectWriteApi`/`zerobuswrite`) the UC REST API rejects; this +/// requests plain `all-apis` client credentials. +async fn mint_metadata_token( + client: &reqwest::Client, + base: &reqwest::Url, + client_id: &str, + client_secret: &str, +) -> ZerobusResult { + let url = join_path(base, ["oidc", "v1", "token"]); + let params = [("grant_type", "client_credentials"), ("scope", "all-apis")]; + + let response = client + .post(url) + .basic_auth(client_id, Some(client_secret)) + .form(¶ms) + .send() + .await + .map_err(|e| transport_error("token request", &e))?; + + let body = read_body("token request", response).await?; + let body: serde_json::Value = serde_json::from_slice(&body) + .map_err(|e| fetch_error(format!("could not parse token response: {e}"), false))?; + + let token = body["access_token"] + .as_str() + .ok_or_else(|| fetch_error("token response has no access_token".to_string(), false))?; + + // Reject a token that can't be a header value here, not opaquely on the next request. + if token.is_empty() || !token.bytes().all(|b| b >= 0x20 && b != 0x7f) { + return Err(fetch_error( + "token response contains an unusable access_token".to_string(), + false, + )); + } + Ok(token.to_string()) +} + +/// Read one table's metadata from the Unity Catalog REST API. +async fn get_table( + client: &reqwest::Client, + base: &reqwest::Url, + token: &str, + table_name: &str, +) -> ZerobusResult { + // `join_path` percent-encodes the segment, so the table name can't alter the path. + let url = join_path(base, ["api", "2.1", "unity-catalog", "tables", table_name]); + + let response = client + .get(url) + .bearer_auth(token) + .header(reqwest::header::ACCEPT, "application/json") + .send() + .await + .map_err(|e| transport_error("schema request", &e))?; + + let body = read_body("schema request", response).await?; + serde_json::from_slice(&body).map_err(|e| { + fetch_error( + format!("could not parse Unity Catalog response for table '{table_name}': {e}"), + false, + ) + }) +} + +/// Read a response body, failing on a non-success status or an oversized body. +async fn read_body(operation: &str, response: reqwest::Response) -> ZerobusResult> { + let status = response.status(); + + // Early reject via Content-Length; the streamed read bounds an absent/understated one. + if response + .content_length() + .is_some_and(|len| len > MAX_RESPONSE_BYTES as u64) + { + return Err(oversized_body_error(operation, status)); + } + + let mut response = response; + let mut body = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .map_err(|e| transport_error(operation, &e))? + { + if body.len() + chunk.len() > MAX_RESPONSE_BYTES { + return Err(oversized_body_error(operation, status)); + } + body.extend_from_slice(&chunk); + } + + if !status.is_success() { + // Truncate the server's error body so an HTML page can't swamp the log. + let detail = String::from_utf8_lossy(&body); + let detail: String = detail.trim().chars().take(512).collect(); + let retryable = status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS; + let message = if detail.is_empty() { + format!("{operation} failed with HTTP {status}") + } else { + format!("{operation} failed with HTTP {status}: {detail}") + }; + warn!(%status, retryable, "{operation} to Unity Catalog failed"); + return Err(fetch_error(message, retryable)); + } + + Ok(body) +} + +fn oversized_body_error(operation: &str, status: StatusCode) -> ZerobusError { + fetch_error( + format!( + "{operation} returned a response larger than {MAX_RESPONSE_BYTES} bytes (HTTP {status})" + ), + false, + ) +} + +/// Classify a `reqwest` transport failure: timeouts, connection, and incomplete +/// request/response errors are transient; anything else terminal. +fn transport_error(operation: &str, error: &reqwest::Error) -> ZerobusError { + let retryable = error.is_timeout() || error.is_connect() || error.is_request(); + fetch_error(format!("{operation} failed: {error}"), retryable) +} + +fn fetch_error(message: String, retryable: bool) -> ZerobusError { + ZerobusError::SchemaFetchError { message, retryable } +} + +/// Parse the workspace URL, defaulting a missing scheme to `https` (matching +/// [`ZerobusSdkBuilder::endpoint`](crate::ZerobusSdkBuilder::endpoint)). +fn normalize_endpoint(unity_catalog_url: &str) -> ZerobusResult { + let trimmed = unity_catalog_url.trim(); + if trimmed.is_empty() { + return Err(ZerobusError::InvalidUCEndpointError( + "unity_catalog_url is required to fetch a schema from Unity Catalog; set it on the SDK builder".to_string(), + )); + } + + let candidate = if trimmed.contains("://") { + trimmed.to_string() + } else { + format!("https://{trimmed}") + }; + + let url = reqwest::Url::parse(&candidate) + .map_err(|e| ZerobusError::InvalidUCEndpointError(format!("{unity_catalog_url}: {e}")))?; + if !matches!(url.scheme(), "http" | "https") { + return Err(ZerobusError::InvalidUCEndpointError(format!( + "{unity_catalog_url}: expected an http or https URL" + ))); + } + if !url.has_host() { + return Err(ZerobusError::InvalidUCEndpointError(format!( + "{unity_catalog_url}: URL has no host" + ))); + } + // Reject embedded credentials so a secret can't leak into a quoted-URL error. + if !url.username().is_empty() || url.password().is_some() { + return Err(ZerobusError::InvalidUCEndpointError( + "unity_catalog_url must not embed credentials".to_string(), + )); + } + Ok(url) +} + +/// Append `segments` to `base`'s path, percent-encoding each one. +fn join_path<'a>(base: &reqwest::Url, segments: impl IntoIterator) -> reqwest::Url { + let mut url = base.clone(); + { + // `base` has a host (validated), so it's never a cannot-be-a-base URL. + let mut path = url + .path_segments_mut() + .expect("validated endpoint always has a host"); + // pop_if_empty drops the empty segment a trailing slash would leave. + path.pop_if_empty().extend(segments); + } + url +} + +/// Reject a table name that is not `catalog.schema.table`, before any network call. +fn validate_table_name(table_name: &str) -> ZerobusResult<()> { + let parts: Vec<&str> = table_name.split('.').collect(); + if parts.len() != 3 || parts.iter().any(|p| p.trim().is_empty()) { + return Err(ZerobusError::InvalidTableName(format!( + "expected 'catalog.schema.table', got '{table_name}'" + ))); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema::UcColumn; + + #[test] + fn validate_table_name_requires_three_nonempty_parts() { + assert!(validate_table_name("cat.sch.tbl").is_ok()); + for bad in ["cat.sch", "cat.sch.tbl.extra", "", ".sch.tbl", "cat. .tbl"] { + assert!( + matches!( + validate_table_name(bad), + Err(ZerobusError::InvalidTableName(_)) + ), + "expected {bad:?} to be rejected" + ); + } + } + + #[test] + fn normalize_endpoint_defaults_to_https_and_rejects_bad_input() { + assert_eq!( + normalize_endpoint("workspace.cloud.databricks.com") + .unwrap() + .as_str(), + "https://workspace.cloud.databricks.com/" + ); + assert_eq!( + normalize_endpoint(" http://localhost:8080 ") + .unwrap() + .as_str(), + "http://localhost:8080/" + ); + for bad in [ + "", + " ", + "ftp://example.com", + "not a url", + // Credentials in the URL would end up in error messages. + "https://user:secret@workspace.cloud.databricks.com", + ] { + assert!( + matches!( + normalize_endpoint(bad), + Err(ZerobusError::InvalidUCEndpointError(_)) + ), + "expected {bad:?} to be rejected" + ); + } + } + + #[test] + fn join_path_percent_encodes_and_handles_trailing_slash() { + let base = normalize_endpoint("https://workspace.cloud.databricks.com/").unwrap(); + let url = join_path(&base, ["api", "2.1", "unity-catalog", "tables", "c.s.t"]); + assert_eq!( + url.as_str(), + "https://workspace.cloud.databricks.com/api/2.1/unity-catalog/tables/c.s.t" + ); + + // A name needing escaping must not escape its path segment. + let url = join_path(&base, ["tables", "c.s.odd name/../x"]); + assert_eq!( + url.as_str(), + "https://workspace.cloud.databricks.com/tables/c.s.odd%20name%2F..%2Fx" + ); + } + + #[test] + fn descriptor_from_schema_converts_columns() { + let schema = UcTableSchema { + name: "orders".to_string(), + catalog_name: "main".to_string(), + schema_name: "sales".to_string(), + columns: vec![ + UcColumn { + name: "id".to_string(), + type_name: "BIGINT".to_string(), + type_text: "bigint".to_string(), + type_json: String::new(), + nullable: false, + position: 0, + }, + UcColumn { + name: "customer".to_string(), + type_name: "STRING".to_string(), + type_text: "string".to_string(), + type_json: String::new(), + nullable: true, + position: 1, + }, + ], + }; + + let md = descriptor_from_schema(&schema).unwrap(); + // Message name comes from `descriptor_from_uc_schema`: _. + assert_eq!(md.name(), "SalesOrders"); + assert_eq!(md.get_field_by_name("id").unwrap().number(), 1); + assert_eq!(md.get_field_by_name("customer").unwrap().number(), 2); + } + + #[test] + fn descriptor_from_schema_rejects_unsupported_column() { + let schema = UcTableSchema { + name: "t".to_string(), + catalog_name: "c".to_string(), + schema_name: "s".to_string(), + columns: vec![UcColumn { + name: "weird".to_string(), + type_name: "INTERVAL".to_string(), + type_text: String::new(), + type_json: String::new(), + nullable: true, + position: 0, + }], + }; + + match descriptor_from_schema(&schema) { + Err(ZerobusError::InvalidArgument(msg)) => assert!(msg.contains("INTERVAL"), "{msg}"), + other => panic!("expected InvalidArgument, got {other:?}"), + } + } + + #[test] + fn schema_fetch_error_retryability_is_carried() { + assert!(fetch_error("boom".to_string(), true).is_retryable()); + assert!(!fetch_error("boom".to_string(), false).is_retryable()); + } +} diff --git a/rust/tests/Cargo.toml b/rust/tests/Cargo.toml index f45c75b0..1ca3f9df 100644 --- a/rust/tests/Cargo.toml +++ b/rust/tests/Cargo.toml @@ -20,6 +20,10 @@ path = "src/multiplexed_stream_tests.rs" name = "arrow_tests" path = "src/arrow_tests.rs" +[[test]] +name = "uc_schema_tests" +path = "src/uc_schema_tests.rs" + [dependencies] async-trait.workspace = true prost.workspace = true diff --git a/rust/tests/src/mock_uc.rs b/rust/tests/src/mock_uc.rs new file mode 100644 index 00000000..f4fa8f3a --- /dev/null +++ b/rust/tests/src/mock_uc.rs @@ -0,0 +1,268 @@ +//! A minimal HTTP mock of the Unity Catalog endpoints the SDK's schema fetch +//! uses: `POST /oidc/v1/token` and `GET /api/2.1/unity-catalog/tables/{name}`. +//! +//! Hand-rolled over a `TcpListener` rather than pulling in an HTTP mock crate: +//! the fetch path only needs these two routes, and the tests assert on the raw +//! request line and headers. + +#![allow(dead_code)] + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +/// How a route should respond. +#[derive(Clone, Debug)] +pub enum MockReply { + /// `200` with this JSON body. + Json(String), + /// This status with this body (used for error classification tests). + Status(u16, String), + /// A `200` JSON body, but with a `Content-Length` claiming `usize` bytes so + /// the client's size guard is exercised without sending that much data. + OverlongContentLength(usize), + /// Accept the connection, then drop it without replying. + Hangup, +} + +/// What the mock recorded about the requests it served. +#[derive(Default, Debug)] +pub struct MockRequests { + /// Request target of every schema request, in order (e.g. + /// `/api/2.1/unity-catalog/tables/c.s.t`). + pub schema_paths: Mutex>, + /// `authorization` header of every schema request, in order. + pub schema_auth: Mutex>, + /// `authorization` header of every token request, in order. + pub token_auth: Mutex>, + /// Body of every token request, in order. + pub token_bodies: Mutex>, + pub token_calls: AtomicUsize, + pub schema_calls: AtomicUsize, +} + +/// A running mock. Dropping it stops the accept loop. +pub struct MockUc { + pub url: String, + pub requests: Arc, + shutdown: tokio::sync::oneshot::Sender<()>, +} + +impl MockUc { + pub fn token_calls(&self) -> usize { + self.requests.token_calls.load(Ordering::SeqCst) + } + + pub fn schema_calls(&self) -> usize { + self.requests.schema_calls.load(Ordering::SeqCst) + } + + /// Stop serving. Not required — dropping the handle does the same. + pub fn stop(self) { + let _ = self.shutdown.send(()); + } +} + +/// Start a mock serving `token_reply` on the token route and `schema_reply` on +/// the table route, bound to an ephemeral loopback port. +pub async fn start_mock_uc(token_reply: MockReply, schema_reply: MockReply) -> MockUc { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("http://{addr}"); + let requests = Arc::new(MockRequests::default()); + let (shutdown, mut shutdown_rx) = tokio::sync::oneshot::channel(); + + let served = Arc::clone(&requests); + tokio::spawn(async move { + loop { + let accepted = tokio::select! { + accepted = listener.accept() => accepted, + _ = &mut shutdown_rx => break, + }; + let Ok((socket, _)) = accepted else { break }; + + let served = Arc::clone(&served); + let token_reply = token_reply.clone(); + let schema_reply = schema_reply.clone(); + tokio::spawn(async move { + handle_connection(socket, served, token_reply, schema_reply).await; + }); + } + }); + + MockUc { + url, + requests, + shutdown, + } +} + +/// Convenience: a mock that mints `token` and serves `columns_json` as the +/// table's `columns` array. +pub async fn start_mock_uc_serving_columns(token: &str, columns_json: &str) -> MockUc { + start_mock_uc( + MockReply::Json(format!(r#"{{"access_token":"{token}","expires_in":3600}}"#)), + MockReply::Json(table_response(columns_json)), + ) + .await +} + +/// A UC `tables/{name}` response body wrapping `columns_json`. +pub fn table_response(columns_json: &str) -> String { + format!( + r#"{{"name":"orders","catalog_name":"main","schema_name":"sales","columns":{columns_json}}}"# + ) +} + +/// Two simple columns: `id` (BIGINT, non-null) and `customer` (STRING, nullable). +pub fn simple_columns_json() -> &'static str { + r#"[ + {"name":"id","type_name":"BIGINT","type_text":"bigint","type_json":"","nullable":false,"position":0}, + {"name":"customer","type_name":"STRING","type_text":"string","type_json":"","nullable":true,"position":1} + ]"# +} + +async fn handle_connection( + mut socket: tokio::net::TcpStream, + requests: Arc, + token_reply: MockReply, + schema_reply: MockReply, +) { + let Some(request) = read_request(&mut socket).await else { + return; + }; + + let is_token = request.target.starts_with("/oidc/v1/token"); + if is_token { + requests.token_calls.fetch_add(1, Ordering::SeqCst); + requests + .token_auth + .lock() + .unwrap() + .push(request.authorization.clone()); + requests.token_bodies.lock().unwrap().push(request.body); + } else { + requests.schema_calls.fetch_add(1, Ordering::SeqCst); + requests + .schema_paths + .lock() + .unwrap() + .push(request.target.clone()); + requests + .schema_auth + .lock() + .unwrap() + .push(request.authorization.clone()); + } + + let reply = if is_token { token_reply } else { schema_reply }; + let response = match reply { + MockReply::Json(body) => http_response(200, "application/json", &body), + MockReply::Status(status, body) => http_response(status, "text/plain", &body), + MockReply::OverlongContentLength(claimed) => { + // Claim a huge body but send a tiny one: the client must reject on the + // advertised length rather than reading to completion. + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {claimed}\r\n\r\n{{}}" + ) + } + MockReply::Hangup => return, + }; + + let _ = socket.write_all(response.as_bytes()).await; + let _ = socket.flush().await; +} + +struct MockRequest { + target: String, + authorization: String, + body: String, +} + +/// Read one request: the head, then `Content-Length` bytes of body. +async fn read_request(socket: &mut tokio::net::TcpStream) -> Option { + let mut raw = Vec::new(); + let mut buf = [0u8; 1024]; + + // Read until the end of the head. + let head_end = loop { + let n = socket.read(&mut buf).await.ok()?; + if n == 0 { + return None; + } + raw.extend_from_slice(&buf[..n]); + if let Some(pos) = find_head_end(&raw) { + break pos; + } + if raw.len() > 64 * 1024 { + return None; + } + }; + + let head = String::from_utf8_lossy(&raw[..head_end]).to_string(); + let mut lines = head.lines(); + let target = lines + .next()? + .split_whitespace() + .nth(1) + .unwrap_or_default() + .to_string(); + + let header = |name: &str| -> Option { + head.lines() + .filter(|l| { + l.split(':') + .next() + .is_some_and(|k| k.trim().eq_ignore_ascii_case(name)) + }) + .map(|l| l.split_once(':').map(|(_, v)| v.trim().to_string())) + .next() + .flatten() + }; + let authorization = header("authorization").unwrap_or_default(); + let content_length: usize = header("content-length") + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + + // The body may already be (partly) buffered with the head. + let mut body = raw[head_end..].to_vec(); + while body.len() < content_length { + let n = socket.read(&mut buf).await.ok()?; + if n == 0 { + break; + } + body.extend_from_slice(&buf[..n]); + } + body.truncate(content_length); + + Some(MockRequest { + target, + authorization, + body: String::from_utf8_lossy(&body).to_string(), + }) +} + +/// Index just past the blank line ending the request head. +fn find_head_end(raw: &[u8]) -> Option { + raw.windows(4).position(|w| w == b"\r\n\r\n").map(|p| p + 4) +} + +fn http_response(status: u16, content_type: &str, body: &str) -> String { + let reason = match status { + 200 => "OK", + 400 => "Bad Request", + 401 => "Unauthorized", + 403 => "Forbidden", + 404 => "Not Found", + 429 => "Too Many Requests", + 500 => "Internal Server Error", + 503 => "Service Unavailable", + _ => "Status", + }; + format!( + "HTTP/1.1 {status} {reason}\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) +} diff --git a/rust/tests/src/uc_schema_tests.rs b/rust/tests/src/uc_schema_tests.rs new file mode 100644 index 00000000..73d97e8f --- /dev/null +++ b/rust/tests/src/uc_schema_tests.rs @@ -0,0 +1,394 @@ +//! Tests for fetching a table's schema from Unity Catalog and using it as the +//! dynamic-proto descriptor (`uc_schema`, `ZerobusSdk::fetch_message_descriptor`). + +mod mock_uc; +mod utils; + +use databricks_zerobus_ingest_sdk::uc_schema::{fetch_message_descriptor, fetch_table_schema}; +use databricks_zerobus_ingest_sdk::{ZerobusError, ZerobusSdk}; +use mock_uc::{ + simple_columns_json, start_mock_uc, start_mock_uc_serving_columns, table_response, MockReply, +}; +use utils::setup_tracing; + +const TABLE_NAME: &str = "main.sales.orders"; + +/// A token reply minting `access_token`. +fn token_ok() -> MockReply { + MockReply::Json(r#"{"access_token":"test-token","expires_in":3600}"#.to_string()) +} + +mod fetch_tests { + use super::*; + + #[tokio::test] + async fn fetches_and_resolves_descriptor() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + + let descriptor = + fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + // Message name is _
, sanitized. + assert_eq!(descriptor.name(), "SalesOrders"); + // Proto field number is the UC position + 1. + assert_eq!(descriptor.get_field_by_name("id").unwrap().number(), 1); + assert_eq!( + descriptor.get_field_by_name("customer").unwrap().number(), + 2 + ); + + assert_eq!(mock.token_calls(), 1, "should mint exactly one token"); + assert_eq!(mock.schema_calls(), 1, "should make one schema request"); + } + + #[tokio::test] + async fn sends_basic_auth_token_request_then_bearer_schema_request() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + + fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + // The token request authenticates with client credentials... + let token_auth = mock.requests.token_auth.lock().unwrap().clone(); + assert_eq!(token_auth.len(), 1); + assert!( + token_auth[0].starts_with("Basic "), + "expected Basic auth on the token request, got {:?}", + token_auth[0] + ); + + // ...and requests plain all-apis client credentials, not an ingestion + // token (which UC's REST API would reject). + let token_body = mock.requests.token_bodies.lock().unwrap().clone(); + assert!( + token_body[0].contains("grant_type=client_credentials") + && token_body[0].contains("scope=all-apis"), + "unexpected token request body: {:?}", + token_body[0] + ); + assert!( + !token_body[0].contains("authorization_details") + && !token_body[0].contains("zerobusDirectWriteApi"), + "token request must not request ingestion scopes: {:?}", + token_body[0] + ); + + // The schema request presents the minted token as a bearer token. + let schema_auth = mock.requests.schema_auth.lock().unwrap().clone(); + assert_eq!(schema_auth, vec!["Bearer test-token".to_string()]); + } + + #[tokio::test] + async fn requests_the_expected_table_path() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + + fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + let paths = mock.requests.schema_paths.lock().unwrap().clone(); + assert_eq!( + paths, + vec![format!("/api/2.1/unity-catalog/tables/{TABLE_NAME}")] + ); + } + + #[tokio::test] + async fn fetch_table_schema_returns_raw_columns() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + + let schema = fetch_table_schema(&mock.url, TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + assert_eq!(schema.name, "orders"); + assert_eq!(schema.catalog_name, "main"); + assert_eq!(schema.schema_name, "sales"); + assert_eq!(schema.columns.len(), 2); + assert_eq!(schema.columns[0].name, "id"); + assert_eq!(schema.columns[0].type_name, "BIGINT"); + assert!(!schema.columns[0].nullable); + } + + #[tokio::test] + async fn resolves_complex_columns_from_type_json() { + setup_tracing(); + let columns = r#"[ + {"name":"id","type_name":"BIGINT","type_text":"bigint","type_json":"","nullable":false,"position":0}, + {"name":"address","type_name":"STRUCT","type_text":"struct","type_json":"{\"type\":\"struct\",\"fields\":[{\"name\":\"street\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]}","nullable":true,"position":1} + ]"#; + let mock = start_mock_uc_serving_columns("test-token", columns).await; + + let descriptor = + fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + let address = descriptor.get_field_by_name("address").unwrap(); + let nested = match address.kind() { + prost_reflect::Kind::Message(m) => m, + other => panic!("expected a nested message, got {other:?}"), + }; + assert!(nested.get_field_by_name("street").is_some()); + } +} + +mod error_classification_tests { + use super::*; + + /// Assert `result` failed with a `SchemaFetchError` whose retryability + /// matches `retryable` and whose message contains `needle`. + fn assert_schema_fetch_error( + result: Result, + retryable: bool, + needle: &str, + ) { + match result { + Err(err @ ZerobusError::SchemaFetchError { .. }) => { + assert_eq!( + err.is_retryable(), + retryable, + "unexpected retryability for {err}" + ); + assert!( + err.to_string().contains(needle), + "expected {needle:?} in error, got: {err}" + ); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + } + + #[tokio::test] + async fn server_error_on_schema_request_is_retryable() { + setup_tracing(); + let mock = start_mock_uc( + token_ok(), + MockReply::Status(503, "upstream unavailable".to_string()), + ) + .await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, true, "503"); + } + + #[tokio::test] + async fn too_many_requests_is_retryable() { + setup_tracing(); + let mock = start_mock_uc(token_ok(), MockReply::Status(429, "slow down".to_string())).await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, true, "429"); + } + + #[tokio::test] + async fn client_errors_are_not_retryable() { + setup_tracing(); + for status in [400u16, 401, 403, 404] { + let mock = + start_mock_uc(token_ok(), MockReply::Status(status, "denied".to_string())).await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, &status.to_string()); + } + } + + #[tokio::test] + async fn failed_token_request_is_reported_as_such() { + setup_tracing(); + let mock = start_mock_uc( + MockReply::Status(401, "bad credentials".to_string()), + MockReply::Json(table_response(simple_columns_json())), + ) + .await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, "token request"); + assert_eq!( + mock.schema_calls(), + 0, + "must not request the schema without a token" + ); + } + + #[tokio::test] + async fn token_response_without_access_token_is_rejected() { + setup_tracing(); + let mock = start_mock_uc( + MockReply::Json(r#"{"expires_in":3600}"#.to_string()), + MockReply::Json(table_response(simple_columns_json())), + ) + .await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, "access_token"); + } + + #[tokio::test] + async fn unparseable_schema_response_is_rejected() { + setup_tracing(); + let mock = start_mock_uc(token_ok(), MockReply::Json("this is not json".to_string())).await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, "could not parse"); + } + + #[tokio::test] + async fn empty_column_list_is_rejected() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", "[]").await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, "no columns"); + } + + #[tokio::test] + async fn oversized_response_is_rejected() { + setup_tracing(); + let mock = start_mock_uc( + token_ok(), + MockReply::OverlongContentLength(64 * 1024 * 1024), + ) + .await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, false, "larger than"); + } + + #[tokio::test] + async fn dropped_connection_is_retryable() { + setup_tracing(); + let mock = start_mock_uc(MockReply::Hangup, MockReply::Hangup).await; + + let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; + assert_schema_fetch_error(result, true, "token request"); + } + + #[tokio::test] + async fn unsupported_column_type_is_invalid_argument() { + setup_tracing(); + let columns = r#"[ + {"name":"span","type_name":"INTERVAL","type_text":"interval","type_json":"","nullable":true,"position":0} + ]"#; + let mock = start_mock_uc_serving_columns("test-token", columns).await; + + // A schema the server returned but that has no protobuf representation is + // a caller-facing argument problem, not a fetch failure. + match fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await { + Err(ZerobusError::InvalidArgument(msg)) => { + assert!(msg.contains("INTERVAL"), "unexpected message: {msg}"); + } + other => panic!("expected InvalidArgument, got {other:?}"), + } + } + + #[tokio::test] + async fn invalid_table_name_fails_before_any_request() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + + for bad in ["orders", "sales.orders", "main.sales.orders.extra", ""] { + match fetch_message_descriptor(&mock.url, bad, "cid", "csec").await { + Err(ZerobusError::InvalidTableName(msg)) => { + assert!(msg.contains("catalog.schema.table"), "got: {msg}"); + } + other => panic!("expected InvalidTableName for {bad:?}, got {other:?}"), + } + } + assert_eq!(mock.token_calls(), 0, "must not hit the network"); + } + + #[tokio::test] + async fn invalid_endpoint_fails_before_any_request() { + setup_tracing(); + for bad in ["", " ", "ftp://example.com"] { + match fetch_message_descriptor(bad, TABLE_NAME, "cid", "csec").await { + Err(ZerobusError::InvalidUCEndpointError(_)) => {} + other => panic!("expected InvalidUCEndpointError for {bad:?}, got {other:?}"), + } + } + } +} + +mod sdk_convenience_tests { + use super::*; + + /// An SDK whose `unity_catalog_url` points at `uc_url`. The zerobus endpoint + /// is never dialed here — these tests only exercise the schema fetch and the + /// builder wiring, both of which precede the gRPC connection. + fn sdk_with_uc(uc_url: &str) -> ZerobusSdk { + ZerobusSdk::builder() + .endpoint("http://127.0.0.1:1") + .unity_catalog_url(uc_url) + .no_tls() + .build() + .expect("sdk should build") + } + + #[tokio::test] + async fn sdk_fetch_message_descriptor_uses_configured_uc_url() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + let sdk = sdk_with_uc(&mock.url); + + let descriptor = sdk + .fetch_message_descriptor(TABLE_NAME, "client-id", "client-secret") + .await + .expect("fetch should succeed"); + + assert_eq!(descriptor.name(), "SalesOrders"); + assert_eq!(mock.schema_calls(), 1); + } + + #[tokio::test] + async fn fetched_descriptor_plugs_into_dynamic_proto_builder() { + setup_tracing(); + let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; + let sdk = sdk_with_uc(&mock.url); + + // The whole point of the utility: fetch once, then hand the descriptor to + // the existing `.dynamic_proto()` selector. `validate()` confirms the + // builder accepts it without any further schema work. + let descriptor = sdk + .fetch_message_descriptor(TABLE_NAME, "cid", "csec") + .await + .expect("fetch should succeed"); + + let builder = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth("cid", "csec") + .dynamic_proto(descriptor); + builder.validate().expect("validation should succeed"); + } + + #[tokio::test] + async fn fetch_failure_surfaces_from_sdk_helper() { + setup_tracing(); + let mock = start_mock_uc( + token_ok(), + MockReply::Status(404, "table not found".to_string()), + ) + .await; + let sdk = sdk_with_uc(&mock.url); + + match sdk + .fetch_message_descriptor(TABLE_NAME, "cid", "csec") + .await + { + Err(err @ ZerobusError::SchemaFetchError { .. }) => { + assert!(!err.is_retryable()); + assert!(err.to_string().contains("404"), "got: {err}"); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert_eq!(mock.schema_calls(), 1); + } +} From 75ef9ea5de082ee30f1fc30605d7f004ebd2f3e7 Mon Sep 17 00:00:00 2001 From: andrijast-db Date: Mon, 10 Aug 2026 15:38:01 +0000 Subject: [PATCH 2/3] better --- rust/NEXT_CHANGELOG.md | 2 +- rust/README.md | 2 +- rust/examples/proto/dynamic/from_uc.rs | 58 ++-- rust/sdk/src/errors.rs | 16 +- rust/sdk/src/uc_schema.rs | 233 +++------------ rust/tests/Cargo.toml | 4 - rust/tests/src/mock_uc.rs | 268 ----------------- rust/tests/src/uc_schema_tests.rs | 394 ------------------------- 8 files changed, 65 insertions(+), 912 deletions(-) delete mode 100644 rust/tests/src/mock_uc.rs delete mode 100644 rust/tests/src/uc_schema_tests.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 56390e06..48d225e7 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -63,4 +63,4 @@ - Added `ZerobusSdk::fetch_message_descriptor()`, the `uc_schema` module (`fetch_message_descriptor`, `fetch_table_schema`), and the - `ZerobusError::SchemaFetchError { message, retryable }` variant. All additive. + `ZerobusError::SchemaFetchError` variant. All additive. diff --git a/rust/README.md b/rust/README.md index 35e8f578..2e59db13 100644 --- a/rust/README.md +++ b/rust/README.md @@ -608,7 +608,7 @@ stream.flush().await?; // wait once for all pending acknowledgments The fetch needs OAuth credentials able to read the table's metadata (they are presented to the Unity Catalog REST API) and `unity_catalog_url` on the SDK builder. For direct control over the endpoint — outside an `SDK`, or against a different workspace — call `uc_schema::fetch_message_descriptor(unity_catalog_url, table, client_id, client_secret)`. -The fetched schema is a snapshot. If the table changes afterwards, the server rejects stream creation with `ZerobusError::InvalidSchema`; re-fetch and rebuild to pick the change up. A failed fetch surfaces as `ZerobusError::SchemaFetchError`, whose `retryable` flag is set for transport errors and 5xx/429 responses. Note that `DATE` and `TIMESTAMP` columns map to integers (days and microseconds since the Unix epoch) — see the [`schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/schema/) module for the full type mapping, and [`uc_schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/uc_schema/) for the fetch API. +The fetched schema is a snapshot. If the table changes afterwards, the server rejects stream creation with `ZerobusError::InvalidSchema`; re-fetch and rebuild to pick the change up. A failed fetch surfaces as `ZerobusError::SchemaFetchError`. Note that `DATE` and `TIMESTAMP` columns map to integers (days and microseconds since the Unix epoch) — see the [`schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/schema/) module for the full type mapping, and [`uc_schema`](https://docs.rs/databricks-zerobus-ingest-sdk/latest/databricks_zerobus_ingest_sdk/uc_schema/) for the fetch API. Setters can be called in any order. The builder validates at `build()` time that both authentication and format have been configured. diff --git a/rust/examples/proto/dynamic/from_uc.rs b/rust/examples/proto/dynamic/from_uc.rs index 7448c4b6..53ae6f16 100644 --- a/rust/examples/proto/dynamic/from_uc.rs +++ b/rust/examples/proto/dynamic/from_uc.rs @@ -7,67 +7,45 @@ use std::error::Error; -use databricks_zerobus_ingest_sdk::{ProtoBytes, ZerobusSdk, ZerobusStream}; +use databricks_zerobus_ingest_sdk::{ProtoBytes, ZerobusSdk}; // Change constants to match your data. const TABLE_NAME: &str = ""; const DATABRICKS_CLIENT_ID: &str = ""; const DATABRICKS_CLIENT_SECRET: &str = ""; -// Uncomment the appropriate lines for your cloud. - -// For AWS: +// For AWS (for Azure, use *.azuredatabricks.net): const DATABRICKS_WORKSPACE_URL: &str = "https://.cloud.databricks.com"; const SERVER_ENDPOINT: &str = "https://.zerobus..cloud.databricks.com"; -// For Azure: -// const DATABRICKS_WORKSPACE_URL: &str = "https://.azuredatabricks.net"; -// const SERVER_ENDPOINT: &str = "https://.zerobus..azuredatabricks.net"; - #[tokio::main] async fn main() -> Result<(), Box> { // `unity_catalog_url` is required: it is where the schema is fetched from. - let sdk_handle = ZerobusSdk::builder() + let sdk = ZerobusSdk::builder() .endpoint(SERVER_ENDPOINT) .unity_catalog_url(DATABRICKS_WORKSPACE_URL) .build()?; // Descriptor from live table metadata — no columns or `.proto` needed up front. - let descriptor = sdk_handle + let descriptor = sdk .fetch_message_descriptor(TABLE_NAME, DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) .await?; + println!( + "Fetched schema '{}' with {} fields", + descriptor.name(), + descriptor.fields().count() + ); - // The fetched schema can be inspected when the columns are unknown to the program. - println!("Fetched schema '{}':", descriptor.name()); - for field in descriptor.fields() { - println!(" {} ({:?})", field.name(), field.kind()); - } - - let mut stream = sdk_handle + let mut stream = sdk .stream_builder() .table(TABLE_NAME) .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) .dynamic_proto(descriptor) - .max_inflight_requests(100) .build() .await?; - ingest_records(&mut stream).await?; - - stream.close().await?; - println!("Stream closed successfully"); - - Ok(()) -} - -async fn ingest_records(stream: &mut ZerobusStream) -> Result<(), Box> { - // (customer_name, quantity, price) — adjust the field names below to your table. - let orders = [ - ("Alice Smith", 2i32, 25.99f64), - ("Bob Johnson", 1, 89.99), - ("Carol Williams", 3, 45.00), - ]; - + // (customer_name, quantity, price) — adjust the field names to your table. + let orders = [("Alice Smith", 2i32, 25.99f64), ("Bob Johnson", 1, 89.99)]; for (i, (customer_name, quantity, price)) in orders.iter().enumerate() { // set()'s value must match the column's proto type (BIGINT -> i64, INT -> i32). let mut record = stream.new_record()?; @@ -76,17 +54,15 @@ async fn ingest_records(stream: &mut ZerobusStream) -> Result<(), Box .set("customer_name", *customer_name)? .set("quantity", *quantity)? .set("price", *price)?; - - // encode() enforces proto2 required fields; queues without waiting for the ack. - let offset_id = stream + // Queue without waiting for the ack. + stream .ingest_record_offset(ProtoBytes(record.encode()?)) .await?; - println!("Record {i} queued with offset ID: {offset_id}"); } - // Wait once for all pending acks — not after each ingest. - stream.flush().await?; - println!("All records acknowledged"); + stream.flush().await?; // wait once for all pending acks + stream.close().await?; + println!("Done"); Ok(()) } diff --git a/rust/sdk/src/errors.rs b/rust/sdk/src/errors.rs index 3e528ade..85b5b021 100644 --- a/rust/sdk/src/errors.rs +++ b/rust/sdk/src/errors.rs @@ -132,14 +132,9 @@ pub enum ZerobusError { #[error("Token fetch failed: {0}")] TokenFetchError(String), /// Returned when resolving a table's schema from Unity Catalog failed (see - /// [`crate::uc_schema`]). `retryable` is set for transport errors and 5xx/429 - /// responses, unset for a rejected request or an unusable body. - /// - /// `#[non_exhaustive]` so further detail (e.g. the HTTP status) can be added - /// later without a breaking change. - #[error("Failed to fetch table schema from Unity Catalog: {message}.")] - #[non_exhaustive] - SchemaFetchError { message: String, retryable: bool }, + /// [`crate::uc_schema`]). + #[error("Failed to fetch table schema from Unity Catalog: {0}.")] + SchemaFetchError(String), } /// List of gRPC status codes that indicate unretriable errors. @@ -219,9 +214,8 @@ impl ZerobusError { ZerobusError::InvalidStateError(_) => false, ZerobusError::ConnectionTimeout(_) => true, ZerobusError::TokenFetchError(_) => true, - // Classified at the fetch site, where the HTTP status and transport - // error are known. - ZerobusError::SchemaFetchError { retryable, .. } => *retryable, + // A schema fetch is a one-shot setup step, not part of the recovery loop. + ZerobusError::SchemaFetchError(_) => false, } } diff --git a/rust/sdk/src/uc_schema.rs b/rust/sdk/src/uc_schema.rs index 3d9300cc..fb7d3209 100644 --- a/rust/sdk/src/uc_schema.rs +++ b/rust/sdk/src/uc_schema.rs @@ -34,8 +34,7 @@ use std::time::Duration; use prost_reflect::MessageDescriptor; -use reqwest::StatusCode; -use tracing::{debug, warn}; +use tracing::debug; use crate::dynamic_proto::message_descriptor; use crate::schema::{descriptor_from_uc_schema, UcTableSchema}; @@ -44,11 +43,6 @@ use crate::{ZerobusError, ZerobusResult}; /// Deadline for a single fetch (token mint plus schema read). const FETCH_TIMEOUT: Duration = Duration::from_secs(30); -/// Cap on a buffered response body. Both responses are small (a token, a column -/// list); the bound guards against an unexpected reply. Rejected outright, not -/// truncated, so we never act on a partial schema. -const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; - /// Fetch `table_name`'s schema from Unity Catalog and resolve it to a /// [`MessageDescriptor`] for [`dynamic_proto`](crate::StreamBuilder::dynamic_proto). /// @@ -61,8 +55,7 @@ const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; /// /// - [`ZerobusError::InvalidTableName`] if `table_name` is not `catalog.schema.table`. /// - [`ZerobusError::InvalidUCEndpointError`] if `unity_catalog_url` is unusable. -/// - [`ZerobusError::SchemaFetchError`] if the token mint or table read fails -/// (`retryable` set for transport errors and 5xx/429). +/// - [`ZerobusError::SchemaFetchError`] if the token mint or table read fails. /// - [`ZerobusError::InvalidArgument`] if the schema has no protobuf /// representation (e.g. an unsupported column type). pub async fn fetch_message_descriptor( @@ -73,7 +66,12 @@ pub async fn fetch_message_descriptor( ) -> ZerobusResult { let schema = fetch_table_schema(unity_catalog_url, table_name, client_id, client_secret).await?; - descriptor_from_schema(&schema) + let descriptor = descriptor_from_uc_schema(&schema).map_err(|e| { + ZerobusError::InvalidArgument(format!( + "cannot convert Unity Catalog schema for table '{table_name}' to a protobuf descriptor: {e}" + )) + })?; + message_descriptor(&descriptor) } /// Fetch `table_name`'s raw Unity Catalog schema, without converting it to a @@ -95,33 +93,35 @@ pub async fn fetch_table_schema( let client = reqwest::Client::builder() .timeout(FETCH_TIMEOUT) .build() - .map_err(|e| fetch_error(format!("failed to build HTTP client: {e}"), false))?; + .map_err(|e| fetch_error(format!("failed to build HTTP client: {e}")))?; debug!(table = %table_name, "fetching UC table schema"); let token = mint_metadata_token(&client, &base, client_id, client_secret).await?; - let schema = get_table(&client, &base, &token, table_name).await?; + // `join_path` percent-encodes the segment, so the table name can't alter the path. + let url = join_path(&base, ["api", "2.1", "unity-catalog", "tables", table_name]); + let body = client + .get(url) + .bearer_auth(&token) + .header(reqwest::header::ACCEPT, "application/json") + .send() + .await + .and_then(reqwest::Response::error_for_status) + .map_err(|e| fetch_error(format!("schema request failed: {e}")))? + .bytes() + .await + .map_err(|e| fetch_error(format!("reading schema response failed: {e}")))?; + + let schema: UcTableSchema = serde_json::from_slice(&body) + .map_err(|e| fetch_error(format!("could not parse Unity Catalog response: {e}")))?; if schema.columns.is_empty() { - return Err(fetch_error( - format!("Unity Catalog returned no columns for table '{table_name}'"), - false, - )); + return Err(fetch_error(format!( + "Unity Catalog returned no columns for table '{table_name}'" + ))); } Ok(schema) } -/// Convert a fetched [`UcTableSchema`] into a [`MessageDescriptor`], mapping a -/// conversion failure to [`ZerobusError::InvalidArgument`]. -fn descriptor_from_schema(schema: &UcTableSchema) -> ZerobusResult { - let descriptor = descriptor_from_uc_schema(schema).map_err(|e| { - ZerobusError::InvalidArgument(format!( - "cannot convert Unity Catalog schema for table '{}' to a protobuf descriptor: {e}", - schema.name - )) - })?; - message_descriptor(&descriptor) -} - /// Mint an OAuth token for reading table metadata. /// /// Separate from [`crate::DefaultTokenFactory`], which mints an ingestion token @@ -136,119 +136,35 @@ async fn mint_metadata_token( let url = join_path(base, ["oidc", "v1", "token"]); let params = [("grant_type", "client_credentials"), ("scope", "all-apis")]; - let response = client + let body = client .post(url) .basic_auth(client_id, Some(client_secret)) .form(¶ms) .send() .await - .map_err(|e| transport_error("token request", &e))?; + .and_then(reqwest::Response::error_for_status) + .map_err(|e| fetch_error(format!("token request failed: {e}")))? + .bytes() + .await + .map_err(|e| fetch_error(format!("reading token response failed: {e}")))?; - let body = read_body("token request", response).await?; let body: serde_json::Value = serde_json::from_slice(&body) - .map_err(|e| fetch_error(format!("could not parse token response: {e}"), false))?; - + .map_err(|e| fetch_error(format!("could not parse token response: {e}")))?; let token = body["access_token"] .as_str() - .ok_or_else(|| fetch_error("token response has no access_token".to_string(), false))?; + .ok_or_else(|| fetch_error("token response has no access_token".to_string()))?; // Reject a token that can't be a header value here, not opaquely on the next request. if token.is_empty() || !token.bytes().all(|b| b >= 0x20 && b != 0x7f) { return Err(fetch_error( "token response contains an unusable access_token".to_string(), - false, )); } Ok(token.to_string()) } -/// Read one table's metadata from the Unity Catalog REST API. -async fn get_table( - client: &reqwest::Client, - base: &reqwest::Url, - token: &str, - table_name: &str, -) -> ZerobusResult { - // `join_path` percent-encodes the segment, so the table name can't alter the path. - let url = join_path(base, ["api", "2.1", "unity-catalog", "tables", table_name]); - - let response = client - .get(url) - .bearer_auth(token) - .header(reqwest::header::ACCEPT, "application/json") - .send() - .await - .map_err(|e| transport_error("schema request", &e))?; - - let body = read_body("schema request", response).await?; - serde_json::from_slice(&body).map_err(|e| { - fetch_error( - format!("could not parse Unity Catalog response for table '{table_name}': {e}"), - false, - ) - }) -} - -/// Read a response body, failing on a non-success status or an oversized body. -async fn read_body(operation: &str, response: reqwest::Response) -> ZerobusResult> { - let status = response.status(); - - // Early reject via Content-Length; the streamed read bounds an absent/understated one. - if response - .content_length() - .is_some_and(|len| len > MAX_RESPONSE_BYTES as u64) - { - return Err(oversized_body_error(operation, status)); - } - - let mut response = response; - let mut body = Vec::new(); - while let Some(chunk) = response - .chunk() - .await - .map_err(|e| transport_error(operation, &e))? - { - if body.len() + chunk.len() > MAX_RESPONSE_BYTES { - return Err(oversized_body_error(operation, status)); - } - body.extend_from_slice(&chunk); - } - - if !status.is_success() { - // Truncate the server's error body so an HTML page can't swamp the log. - let detail = String::from_utf8_lossy(&body); - let detail: String = detail.trim().chars().take(512).collect(); - let retryable = status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS; - let message = if detail.is_empty() { - format!("{operation} failed with HTTP {status}") - } else { - format!("{operation} failed with HTTP {status}: {detail}") - }; - warn!(%status, retryable, "{operation} to Unity Catalog failed"); - return Err(fetch_error(message, retryable)); - } - - Ok(body) -} - -fn oversized_body_error(operation: &str, status: StatusCode) -> ZerobusError { - fetch_error( - format!( - "{operation} returned a response larger than {MAX_RESPONSE_BYTES} bytes (HTTP {status})" - ), - false, - ) -} - -/// Classify a `reqwest` transport failure: timeouts, connection, and incomplete -/// request/response errors are transient; anything else terminal. -fn transport_error(operation: &str, error: &reqwest::Error) -> ZerobusError { - let retryable = error.is_timeout() || error.is_connect() || error.is_request(); - fetch_error(format!("{operation} failed: {error}"), retryable) -} - -fn fetch_error(message: String, retryable: bool) -> ZerobusError { - ZerobusError::SchemaFetchError { message, retryable } +fn fetch_error(message: String) -> ZerobusError { + ZerobusError::SchemaFetchError(message) } /// Parse the workspace URL, defaulting a missing scheme to `https` (matching @@ -257,7 +173,7 @@ fn normalize_endpoint(unity_catalog_url: &str) -> ZerobusResult { let trimmed = unity_catalog_url.trim(); if trimmed.is_empty() { return Err(ZerobusError::InvalidUCEndpointError( - "unity_catalog_url is required to fetch a schema from Unity Catalog; set it on the SDK builder".to_string(), + "unity_catalog_url is required; set it on the SDK builder".to_string(), )); } @@ -269,14 +185,9 @@ fn normalize_endpoint(unity_catalog_url: &str) -> ZerobusResult { let url = reqwest::Url::parse(&candidate) .map_err(|e| ZerobusError::InvalidUCEndpointError(format!("{unity_catalog_url}: {e}")))?; - if !matches!(url.scheme(), "http" | "https") { - return Err(ZerobusError::InvalidUCEndpointError(format!( - "{unity_catalog_url}: expected an http or https URL" - ))); - } - if !url.has_host() { + if !matches!(url.scheme(), "http" | "https") || !url.has_host() { return Err(ZerobusError::InvalidUCEndpointError(format!( - "{unity_catalog_url}: URL has no host" + "{unity_catalog_url}: expected an http or https URL with a host" ))); } // Reject embedded credentials so a secret can't leak into a quoted-URL error. @@ -316,7 +227,6 @@ fn validate_table_name(table_name: &str) -> ZerobusResult<()> { #[cfg(test)] mod tests { use super::*; - use crate::schema::UcColumn; #[test] fn validate_table_name_requires_three_nonempty_parts() { @@ -380,65 +290,4 @@ mod tests { "https://workspace.cloud.databricks.com/tables/c.s.odd%20name%2F..%2Fx" ); } - - #[test] - fn descriptor_from_schema_converts_columns() { - let schema = UcTableSchema { - name: "orders".to_string(), - catalog_name: "main".to_string(), - schema_name: "sales".to_string(), - columns: vec![ - UcColumn { - name: "id".to_string(), - type_name: "BIGINT".to_string(), - type_text: "bigint".to_string(), - type_json: String::new(), - nullable: false, - position: 0, - }, - UcColumn { - name: "customer".to_string(), - type_name: "STRING".to_string(), - type_text: "string".to_string(), - type_json: String::new(), - nullable: true, - position: 1, - }, - ], - }; - - let md = descriptor_from_schema(&schema).unwrap(); - // Message name comes from `descriptor_from_uc_schema`: _
. - assert_eq!(md.name(), "SalesOrders"); - assert_eq!(md.get_field_by_name("id").unwrap().number(), 1); - assert_eq!(md.get_field_by_name("customer").unwrap().number(), 2); - } - - #[test] - fn descriptor_from_schema_rejects_unsupported_column() { - let schema = UcTableSchema { - name: "t".to_string(), - catalog_name: "c".to_string(), - schema_name: "s".to_string(), - columns: vec![UcColumn { - name: "weird".to_string(), - type_name: "INTERVAL".to_string(), - type_text: String::new(), - type_json: String::new(), - nullable: true, - position: 0, - }], - }; - - match descriptor_from_schema(&schema) { - Err(ZerobusError::InvalidArgument(msg)) => assert!(msg.contains("INTERVAL"), "{msg}"), - other => panic!("expected InvalidArgument, got {other:?}"), - } - } - - #[test] - fn schema_fetch_error_retryability_is_carried() { - assert!(fetch_error("boom".to_string(), true).is_retryable()); - assert!(!fetch_error("boom".to_string(), false).is_retryable()); - } } diff --git a/rust/tests/Cargo.toml b/rust/tests/Cargo.toml index 1ca3f9df..f45c75b0 100644 --- a/rust/tests/Cargo.toml +++ b/rust/tests/Cargo.toml @@ -20,10 +20,6 @@ path = "src/multiplexed_stream_tests.rs" name = "arrow_tests" path = "src/arrow_tests.rs" -[[test]] -name = "uc_schema_tests" -path = "src/uc_schema_tests.rs" - [dependencies] async-trait.workspace = true prost.workspace = true diff --git a/rust/tests/src/mock_uc.rs b/rust/tests/src/mock_uc.rs deleted file mode 100644 index f4fa8f3a..00000000 --- a/rust/tests/src/mock_uc.rs +++ /dev/null @@ -1,268 +0,0 @@ -//! A minimal HTTP mock of the Unity Catalog endpoints the SDK's schema fetch -//! uses: `POST /oidc/v1/token` and `GET /api/2.1/unity-catalog/tables/{name}`. -//! -//! Hand-rolled over a `TcpListener` rather than pulling in an HTTP mock crate: -//! the fetch path only needs these two routes, and the tests assert on the raw -//! request line and headers. - -#![allow(dead_code)] - -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; - -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpListener; - -/// How a route should respond. -#[derive(Clone, Debug)] -pub enum MockReply { - /// `200` with this JSON body. - Json(String), - /// This status with this body (used for error classification tests). - Status(u16, String), - /// A `200` JSON body, but with a `Content-Length` claiming `usize` bytes so - /// the client's size guard is exercised without sending that much data. - OverlongContentLength(usize), - /// Accept the connection, then drop it without replying. - Hangup, -} - -/// What the mock recorded about the requests it served. -#[derive(Default, Debug)] -pub struct MockRequests { - /// Request target of every schema request, in order (e.g. - /// `/api/2.1/unity-catalog/tables/c.s.t`). - pub schema_paths: Mutex>, - /// `authorization` header of every schema request, in order. - pub schema_auth: Mutex>, - /// `authorization` header of every token request, in order. - pub token_auth: Mutex>, - /// Body of every token request, in order. - pub token_bodies: Mutex>, - pub token_calls: AtomicUsize, - pub schema_calls: AtomicUsize, -} - -/// A running mock. Dropping it stops the accept loop. -pub struct MockUc { - pub url: String, - pub requests: Arc, - shutdown: tokio::sync::oneshot::Sender<()>, -} - -impl MockUc { - pub fn token_calls(&self) -> usize { - self.requests.token_calls.load(Ordering::SeqCst) - } - - pub fn schema_calls(&self) -> usize { - self.requests.schema_calls.load(Ordering::SeqCst) - } - - /// Stop serving. Not required — dropping the handle does the same. - pub fn stop(self) { - let _ = self.shutdown.send(()); - } -} - -/// Start a mock serving `token_reply` on the token route and `schema_reply` on -/// the table route, bound to an ephemeral loopback port. -pub async fn start_mock_uc(token_reply: MockReply, schema_reply: MockReply) -> MockUc { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let url = format!("http://{addr}"); - let requests = Arc::new(MockRequests::default()); - let (shutdown, mut shutdown_rx) = tokio::sync::oneshot::channel(); - - let served = Arc::clone(&requests); - tokio::spawn(async move { - loop { - let accepted = tokio::select! { - accepted = listener.accept() => accepted, - _ = &mut shutdown_rx => break, - }; - let Ok((socket, _)) = accepted else { break }; - - let served = Arc::clone(&served); - let token_reply = token_reply.clone(); - let schema_reply = schema_reply.clone(); - tokio::spawn(async move { - handle_connection(socket, served, token_reply, schema_reply).await; - }); - } - }); - - MockUc { - url, - requests, - shutdown, - } -} - -/// Convenience: a mock that mints `token` and serves `columns_json` as the -/// table's `columns` array. -pub async fn start_mock_uc_serving_columns(token: &str, columns_json: &str) -> MockUc { - start_mock_uc( - MockReply::Json(format!(r#"{{"access_token":"{token}","expires_in":3600}}"#)), - MockReply::Json(table_response(columns_json)), - ) - .await -} - -/// A UC `tables/{name}` response body wrapping `columns_json`. -pub fn table_response(columns_json: &str) -> String { - format!( - r#"{{"name":"orders","catalog_name":"main","schema_name":"sales","columns":{columns_json}}}"# - ) -} - -/// Two simple columns: `id` (BIGINT, non-null) and `customer` (STRING, nullable). -pub fn simple_columns_json() -> &'static str { - r#"[ - {"name":"id","type_name":"BIGINT","type_text":"bigint","type_json":"","nullable":false,"position":0}, - {"name":"customer","type_name":"STRING","type_text":"string","type_json":"","nullable":true,"position":1} - ]"# -} - -async fn handle_connection( - mut socket: tokio::net::TcpStream, - requests: Arc, - token_reply: MockReply, - schema_reply: MockReply, -) { - let Some(request) = read_request(&mut socket).await else { - return; - }; - - let is_token = request.target.starts_with("/oidc/v1/token"); - if is_token { - requests.token_calls.fetch_add(1, Ordering::SeqCst); - requests - .token_auth - .lock() - .unwrap() - .push(request.authorization.clone()); - requests.token_bodies.lock().unwrap().push(request.body); - } else { - requests.schema_calls.fetch_add(1, Ordering::SeqCst); - requests - .schema_paths - .lock() - .unwrap() - .push(request.target.clone()); - requests - .schema_auth - .lock() - .unwrap() - .push(request.authorization.clone()); - } - - let reply = if is_token { token_reply } else { schema_reply }; - let response = match reply { - MockReply::Json(body) => http_response(200, "application/json", &body), - MockReply::Status(status, body) => http_response(status, "text/plain", &body), - MockReply::OverlongContentLength(claimed) => { - // Claim a huge body but send a tiny one: the client must reject on the - // advertised length rather than reading to completion. - format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {claimed}\r\n\r\n{{}}" - ) - } - MockReply::Hangup => return, - }; - - let _ = socket.write_all(response.as_bytes()).await; - let _ = socket.flush().await; -} - -struct MockRequest { - target: String, - authorization: String, - body: String, -} - -/// Read one request: the head, then `Content-Length` bytes of body. -async fn read_request(socket: &mut tokio::net::TcpStream) -> Option { - let mut raw = Vec::new(); - let mut buf = [0u8; 1024]; - - // Read until the end of the head. - let head_end = loop { - let n = socket.read(&mut buf).await.ok()?; - if n == 0 { - return None; - } - raw.extend_from_slice(&buf[..n]); - if let Some(pos) = find_head_end(&raw) { - break pos; - } - if raw.len() > 64 * 1024 { - return None; - } - }; - - let head = String::from_utf8_lossy(&raw[..head_end]).to_string(); - let mut lines = head.lines(); - let target = lines - .next()? - .split_whitespace() - .nth(1) - .unwrap_or_default() - .to_string(); - - let header = |name: &str| -> Option { - head.lines() - .filter(|l| { - l.split(':') - .next() - .is_some_and(|k| k.trim().eq_ignore_ascii_case(name)) - }) - .map(|l| l.split_once(':').map(|(_, v)| v.trim().to_string())) - .next() - .flatten() - }; - let authorization = header("authorization").unwrap_or_default(); - let content_length: usize = header("content-length") - .and_then(|v| v.parse().ok()) - .unwrap_or(0); - - // The body may already be (partly) buffered with the head. - let mut body = raw[head_end..].to_vec(); - while body.len() < content_length { - let n = socket.read(&mut buf).await.ok()?; - if n == 0 { - break; - } - body.extend_from_slice(&buf[..n]); - } - body.truncate(content_length); - - Some(MockRequest { - target, - authorization, - body: String::from_utf8_lossy(&body).to_string(), - }) -} - -/// Index just past the blank line ending the request head. -fn find_head_end(raw: &[u8]) -> Option { - raw.windows(4).position(|w| w == b"\r\n\r\n").map(|p| p + 4) -} - -fn http_response(status: u16, content_type: &str, body: &str) -> String { - let reason = match status { - 200 => "OK", - 400 => "Bad Request", - 401 => "Unauthorized", - 403 => "Forbidden", - 404 => "Not Found", - 429 => "Too Many Requests", - 500 => "Internal Server Error", - 503 => "Service Unavailable", - _ => "Status", - }; - format!( - "HTTP/1.1 {status} {reason}\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", - body.len() - ) -} diff --git a/rust/tests/src/uc_schema_tests.rs b/rust/tests/src/uc_schema_tests.rs deleted file mode 100644 index 73d97e8f..00000000 --- a/rust/tests/src/uc_schema_tests.rs +++ /dev/null @@ -1,394 +0,0 @@ -//! Tests for fetching a table's schema from Unity Catalog and using it as the -//! dynamic-proto descriptor (`uc_schema`, `ZerobusSdk::fetch_message_descriptor`). - -mod mock_uc; -mod utils; - -use databricks_zerobus_ingest_sdk::uc_schema::{fetch_message_descriptor, fetch_table_schema}; -use databricks_zerobus_ingest_sdk::{ZerobusError, ZerobusSdk}; -use mock_uc::{ - simple_columns_json, start_mock_uc, start_mock_uc_serving_columns, table_response, MockReply, -}; -use utils::setup_tracing; - -const TABLE_NAME: &str = "main.sales.orders"; - -/// A token reply minting `access_token`. -fn token_ok() -> MockReply { - MockReply::Json(r#"{"access_token":"test-token","expires_in":3600}"#.to_string()) -} - -mod fetch_tests { - use super::*; - - #[tokio::test] - async fn fetches_and_resolves_descriptor() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - - let descriptor = - fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - // Message name is _
, sanitized. - assert_eq!(descriptor.name(), "SalesOrders"); - // Proto field number is the UC position + 1. - assert_eq!(descriptor.get_field_by_name("id").unwrap().number(), 1); - assert_eq!( - descriptor.get_field_by_name("customer").unwrap().number(), - 2 - ); - - assert_eq!(mock.token_calls(), 1, "should mint exactly one token"); - assert_eq!(mock.schema_calls(), 1, "should make one schema request"); - } - - #[tokio::test] - async fn sends_basic_auth_token_request_then_bearer_schema_request() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - - fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - // The token request authenticates with client credentials... - let token_auth = mock.requests.token_auth.lock().unwrap().clone(); - assert_eq!(token_auth.len(), 1); - assert!( - token_auth[0].starts_with("Basic "), - "expected Basic auth on the token request, got {:?}", - token_auth[0] - ); - - // ...and requests plain all-apis client credentials, not an ingestion - // token (which UC's REST API would reject). - let token_body = mock.requests.token_bodies.lock().unwrap().clone(); - assert!( - token_body[0].contains("grant_type=client_credentials") - && token_body[0].contains("scope=all-apis"), - "unexpected token request body: {:?}", - token_body[0] - ); - assert!( - !token_body[0].contains("authorization_details") - && !token_body[0].contains("zerobusDirectWriteApi"), - "token request must not request ingestion scopes: {:?}", - token_body[0] - ); - - // The schema request presents the minted token as a bearer token. - let schema_auth = mock.requests.schema_auth.lock().unwrap().clone(); - assert_eq!(schema_auth, vec!["Bearer test-token".to_string()]); - } - - #[tokio::test] - async fn requests_the_expected_table_path() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - - fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - let paths = mock.requests.schema_paths.lock().unwrap().clone(); - assert_eq!( - paths, - vec![format!("/api/2.1/unity-catalog/tables/{TABLE_NAME}")] - ); - } - - #[tokio::test] - async fn fetch_table_schema_returns_raw_columns() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - - let schema = fetch_table_schema(&mock.url, TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - assert_eq!(schema.name, "orders"); - assert_eq!(schema.catalog_name, "main"); - assert_eq!(schema.schema_name, "sales"); - assert_eq!(schema.columns.len(), 2); - assert_eq!(schema.columns[0].name, "id"); - assert_eq!(schema.columns[0].type_name, "BIGINT"); - assert!(!schema.columns[0].nullable); - } - - #[tokio::test] - async fn resolves_complex_columns_from_type_json() { - setup_tracing(); - let columns = r#"[ - {"name":"id","type_name":"BIGINT","type_text":"bigint","type_json":"","nullable":false,"position":0}, - {"name":"address","type_name":"STRUCT","type_text":"struct","type_json":"{\"type\":\"struct\",\"fields\":[{\"name\":\"street\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]}","nullable":true,"position":1} - ]"#; - let mock = start_mock_uc_serving_columns("test-token", columns).await; - - let descriptor = - fetch_message_descriptor(&mock.url, TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - let address = descriptor.get_field_by_name("address").unwrap(); - let nested = match address.kind() { - prost_reflect::Kind::Message(m) => m, - other => panic!("expected a nested message, got {other:?}"), - }; - assert!(nested.get_field_by_name("street").is_some()); - } -} - -mod error_classification_tests { - use super::*; - - /// Assert `result` failed with a `SchemaFetchError` whose retryability - /// matches `retryable` and whose message contains `needle`. - fn assert_schema_fetch_error( - result: Result, - retryable: bool, - needle: &str, - ) { - match result { - Err(err @ ZerobusError::SchemaFetchError { .. }) => { - assert_eq!( - err.is_retryable(), - retryable, - "unexpected retryability for {err}" - ); - assert!( - err.to_string().contains(needle), - "expected {needle:?} in error, got: {err}" - ); - } - other => panic!("expected SchemaFetchError, got {other:?}"), - } - } - - #[tokio::test] - async fn server_error_on_schema_request_is_retryable() { - setup_tracing(); - let mock = start_mock_uc( - token_ok(), - MockReply::Status(503, "upstream unavailable".to_string()), - ) - .await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, true, "503"); - } - - #[tokio::test] - async fn too_many_requests_is_retryable() { - setup_tracing(); - let mock = start_mock_uc(token_ok(), MockReply::Status(429, "slow down".to_string())).await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, true, "429"); - } - - #[tokio::test] - async fn client_errors_are_not_retryable() { - setup_tracing(); - for status in [400u16, 401, 403, 404] { - let mock = - start_mock_uc(token_ok(), MockReply::Status(status, "denied".to_string())).await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, &status.to_string()); - } - } - - #[tokio::test] - async fn failed_token_request_is_reported_as_such() { - setup_tracing(); - let mock = start_mock_uc( - MockReply::Status(401, "bad credentials".to_string()), - MockReply::Json(table_response(simple_columns_json())), - ) - .await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, "token request"); - assert_eq!( - mock.schema_calls(), - 0, - "must not request the schema without a token" - ); - } - - #[tokio::test] - async fn token_response_without_access_token_is_rejected() { - setup_tracing(); - let mock = start_mock_uc( - MockReply::Json(r#"{"expires_in":3600}"#.to_string()), - MockReply::Json(table_response(simple_columns_json())), - ) - .await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, "access_token"); - } - - #[tokio::test] - async fn unparseable_schema_response_is_rejected() { - setup_tracing(); - let mock = start_mock_uc(token_ok(), MockReply::Json("this is not json".to_string())).await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, "could not parse"); - } - - #[tokio::test] - async fn empty_column_list_is_rejected() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", "[]").await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, "no columns"); - } - - #[tokio::test] - async fn oversized_response_is_rejected() { - setup_tracing(); - let mock = start_mock_uc( - token_ok(), - MockReply::OverlongContentLength(64 * 1024 * 1024), - ) - .await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, false, "larger than"); - } - - #[tokio::test] - async fn dropped_connection_is_retryable() { - setup_tracing(); - let mock = start_mock_uc(MockReply::Hangup, MockReply::Hangup).await; - - let result = fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await; - assert_schema_fetch_error(result, true, "token request"); - } - - #[tokio::test] - async fn unsupported_column_type_is_invalid_argument() { - setup_tracing(); - let columns = r#"[ - {"name":"span","type_name":"INTERVAL","type_text":"interval","type_json":"","nullable":true,"position":0} - ]"#; - let mock = start_mock_uc_serving_columns("test-token", columns).await; - - // A schema the server returned but that has no protobuf representation is - // a caller-facing argument problem, not a fetch failure. - match fetch_message_descriptor(&mock.url, TABLE_NAME, "cid", "csec").await { - Err(ZerobusError::InvalidArgument(msg)) => { - assert!(msg.contains("INTERVAL"), "unexpected message: {msg}"); - } - other => panic!("expected InvalidArgument, got {other:?}"), - } - } - - #[tokio::test] - async fn invalid_table_name_fails_before_any_request() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - - for bad in ["orders", "sales.orders", "main.sales.orders.extra", ""] { - match fetch_message_descriptor(&mock.url, bad, "cid", "csec").await { - Err(ZerobusError::InvalidTableName(msg)) => { - assert!(msg.contains("catalog.schema.table"), "got: {msg}"); - } - other => panic!("expected InvalidTableName for {bad:?}, got {other:?}"), - } - } - assert_eq!(mock.token_calls(), 0, "must not hit the network"); - } - - #[tokio::test] - async fn invalid_endpoint_fails_before_any_request() { - setup_tracing(); - for bad in ["", " ", "ftp://example.com"] { - match fetch_message_descriptor(bad, TABLE_NAME, "cid", "csec").await { - Err(ZerobusError::InvalidUCEndpointError(_)) => {} - other => panic!("expected InvalidUCEndpointError for {bad:?}, got {other:?}"), - } - } - } -} - -mod sdk_convenience_tests { - use super::*; - - /// An SDK whose `unity_catalog_url` points at `uc_url`. The zerobus endpoint - /// is never dialed here — these tests only exercise the schema fetch and the - /// builder wiring, both of which precede the gRPC connection. - fn sdk_with_uc(uc_url: &str) -> ZerobusSdk { - ZerobusSdk::builder() - .endpoint("http://127.0.0.1:1") - .unity_catalog_url(uc_url) - .no_tls() - .build() - .expect("sdk should build") - } - - #[tokio::test] - async fn sdk_fetch_message_descriptor_uses_configured_uc_url() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - let sdk = sdk_with_uc(&mock.url); - - let descriptor = sdk - .fetch_message_descriptor(TABLE_NAME, "client-id", "client-secret") - .await - .expect("fetch should succeed"); - - assert_eq!(descriptor.name(), "SalesOrders"); - assert_eq!(mock.schema_calls(), 1); - } - - #[tokio::test] - async fn fetched_descriptor_plugs_into_dynamic_proto_builder() { - setup_tracing(); - let mock = start_mock_uc_serving_columns("test-token", simple_columns_json()).await; - let sdk = sdk_with_uc(&mock.url); - - // The whole point of the utility: fetch once, then hand the descriptor to - // the existing `.dynamic_proto()` selector. `validate()` confirms the - // builder accepts it without any further schema work. - let descriptor = sdk - .fetch_message_descriptor(TABLE_NAME, "cid", "csec") - .await - .expect("fetch should succeed"); - - let builder = sdk - .stream_builder() - .table(TABLE_NAME) - .oauth("cid", "csec") - .dynamic_proto(descriptor); - builder.validate().expect("validation should succeed"); - } - - #[tokio::test] - async fn fetch_failure_surfaces_from_sdk_helper() { - setup_tracing(); - let mock = start_mock_uc( - token_ok(), - MockReply::Status(404, "table not found".to_string()), - ) - .await; - let sdk = sdk_with_uc(&mock.url); - - match sdk - .fetch_message_descriptor(TABLE_NAME, "cid", "csec") - .await - { - Err(err @ ZerobusError::SchemaFetchError { .. }) => { - assert!(!err.is_retryable()); - assert!(err.to_string().contains("404"), "got: {err}"); - } - other => panic!("expected SchemaFetchError, got {other:?}"), - } - assert_eq!(mock.schema_calls(), 1); - } -} From 8d4f9de2d6c1688efa4e52701b3f44c214130985 Mon Sep 17 00:00:00 2001 From: andrijast-db Date: Mon, 10 Aug 2026 16:21:57 +0000 Subject: [PATCH 3/3] add test --- rust/tests/Cargo.toml | 4 + rust/tests/src/uc_schema_tests.rs | 139 ++++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 rust/tests/src/uc_schema_tests.rs diff --git a/rust/tests/Cargo.toml b/rust/tests/Cargo.toml index f45c75b0..1ca3f9df 100644 --- a/rust/tests/Cargo.toml +++ b/rust/tests/Cargo.toml @@ -20,6 +20,10 @@ path = "src/multiplexed_stream_tests.rs" name = "arrow_tests" path = "src/arrow_tests.rs" +[[test]] +name = "uc_schema_tests" +path = "src/uc_schema_tests.rs" + [dependencies] async-trait.workspace = true prost.workspace = true diff --git a/rust/tests/src/uc_schema_tests.rs b/rust/tests/src/uc_schema_tests.rs new file mode 100644 index 00000000..966d124b --- /dev/null +++ b/rust/tests/src/uc_schema_tests.rs @@ -0,0 +1,139 @@ +//! Tests for fetching a table's schema from Unity Catalog +//! (`uc_schema::fetch_message_descriptor`), against a tiny in-process HTTP mock. + +use std::sync::{Arc, Mutex}; + +use databricks_zerobus_ingest_sdk::uc_schema::fetch_message_descriptor; +use databricks_zerobus_ingest_sdk::ZerobusError; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +const TABLE: &str = "main.sales.orders"; + +/// One recorded request: its target path and `authorization` header. +type Recorded = Vec<(String, String)>; + +/// A running mock. Serves `POST /oidc/v1/token` with a fixed token, then the +/// table route with `schema_status`/`schema_body`. Dropping it stops the loop. +struct MockUc { + url: String, + requests: Arc>, + _shutdown: tokio::sync::oneshot::Sender<()>, +} + +/// Start a mock replying to the schema route with `schema_status` and `schema_body`. +async fn start_mock(schema_status: u16, schema_body: &'static str) -> MockUc { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Recorded::new())); + let (tx, mut rx) = tokio::sync::oneshot::channel(); + + let recorded = Arc::clone(&requests); + tokio::spawn(async move { + loop { + let sock = tokio::select! { + a = listener.accept() => a, + _ = &mut rx => break, + }; + let Ok((mut sock, _)) = sock else { break }; + let recorded = Arc::clone(&recorded); + tokio::spawn(async move { + // The head is enough; the token request's small body follows the + // client's Content-Length but we don't assert on it. + let mut buf = vec![0u8; 4096]; + let n = sock.read(&mut buf).await.unwrap_or(0); + let head = String::from_utf8_lossy(&buf[..n]); + let target = head + .lines() + .next() + .and_then(|l| l.split_whitespace().nth(1)) + .unwrap_or_default() + .to_string(); + let auth = header(&head, "authorization"); + recorded.lock().unwrap().push((target.clone(), auth)); + + let (status, body) = if target.contains("/oidc/") { + (200, r#"{"access_token":"tok-123","expires_in":3600}"#) + } else { + (schema_status, schema_body) + }; + let resp = format!( + "HTTP/1.1 {status} X\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = sock.write_all(resp.as_bytes()).await; + }); + } + }); + + MockUc { + url, + requests, + _shutdown: tx, + } +} + +fn header(head: &str, name: &str) -> String { + head.lines() + .find(|l| l.to_ascii_lowercase().starts_with(&format!("{name}:"))) + .and_then(|l| l.split_once(':')) + .map(|(_, v)| v.trim().to_string()) + .unwrap_or_default() +} + +fn table_json() -> &'static str { + r#"{"name":"orders","catalog_name":"main","schema_name":"sales","columns":[ + {"name":"id","type_name":"BIGINT","type_text":"bigint","type_json":"","nullable":false,"position":0}, + {"name":"customer","type_name":"STRING","type_text":"string","type_json":"","nullable":true,"position":1} + ]}"# +} + +#[tokio::test] +async fn fetches_descriptor_and_sends_expected_requests() { + let mock = start_mock(200, table_json()).await; + + let descriptor = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .expect("fetch should succeed"); + + // Descriptor: message name is _
, field numbers are position + 1. + assert_eq!(descriptor.name(), "SalesOrders"); + assert_eq!(descriptor.get_field_by_name("id").unwrap().number(), 1); + assert_eq!( + descriptor.get_field_by_name("customer").unwrap().number(), + 2 + ); + + // Request shapes: Basic-auth token mint, then a bearer schema request at the + // expected path. + let reqs = mock.requests.lock().unwrap().clone(); + assert_eq!(reqs.len(), 2, "expected a token then a schema request"); + assert!(reqs[0].0.starts_with("/oidc/v1/token"), "got {}", reqs[0].0); + assert!(reqs[0].1.starts_with("Basic "), "got {}", reqs[0].1); + assert_eq!(reqs[1].0, format!("/api/2.1/unity-catalog/tables/{TABLE}")); + assert_eq!(reqs[1].1, "Bearer tok-123"); +} + +#[tokio::test] +async fn schema_request_failure_is_schema_fetch_error() { + let mock = start_mock(404, "table not found").await; + + match fetch_message_descriptor(&mock.url, TABLE, "cid", "csec").await { + Err(ZerobusError::SchemaFetchError(msg)) => assert!(msg.contains("404"), "got: {msg}"), + other => panic!("expected SchemaFetchError, got {other:?}"), + } +} + +#[tokio::test] +async fn invalid_table_name_fails_before_any_request() { + let mock = start_mock(200, table_json()).await; + + match fetch_message_descriptor(&mock.url, "not.qualified", "cid", "csec").await { + Err(ZerobusError::InvalidTableName(_)) => {} + other => panic!("expected InvalidTableName, got {other:?}"), + } + assert!( + mock.requests.lock().unwrap().is_empty(), + "must not hit the network" + ); +}