diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 8baa9599..818ec614 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 now rolls back logical offsets and record ranges when an enqueue @@ -36,3 +59,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 a9573e40..76f3a50d 100644 --- a/rust/README.md +++ b/rust/README.md @@ -608,6 +608,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. Compatible schema evolution may be accepted; incompatible changes fail stream creation with `ZerobusError::CreateStreamError`, so re-fetch the descriptor before rebuilding the stream. 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. ### 5. Ingest Data @@ -1069,6 +1106,7 @@ The `examples/` directory contains working examples covering different serializa | `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. @@ -1342,6 +1380,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 993d01c3..8ddcc085 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 5d75ca81..d72f681a 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 Example Files](#update-example-files) @@ -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 @@ -271,6 +273,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..53ae6f16 --- /dev/null +++ b/rust/examples/proto/dynamic/from_uc.rs @@ -0,0 +1,68 @@ +//! 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}; + +// Change constants to match your data. +const TABLE_NAME: &str = ""; +const DATABRICKS_CLIENT_ID: &str = ""; +const DATABRICKS_CLIENT_SECRET: &str = ""; + +// 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"; + +#[tokio::main] +async fn main() -> Result<(), Box> { + // `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()?; + + // Descriptor from live table metadata — no columns or `.proto` needed up front. + 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() + ); + + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .dynamic_proto(descriptor) + .build() + .await?; + + // (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()?; + record + .set("id", i as i64)? + .set("customer_name", *customer_name)? + .set("quantity", *quantity)? + .set("price", *price)?; + // Queue without waiting for the ack. + stream + .ingest_record_offset(ProtoBytes(record.encode()?)) + .await?; + } + + stream.flush().await?; // wait once for all pending acks + stream.close().await?; + println!("Done"); + + 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..3169797a 100644 --- a/rust/sdk/src/errors.rs +++ b/rust/sdk/src/errors.rs @@ -131,6 +131,11 @@ 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`]). + #[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 +215,7 @@ impl ZerobusError { ZerobusError::InvalidStateError(_) => false, ZerobusError::ConnectionTimeout(_) => true, ZerobusError::TokenFetchError(_) => true, + ZerobusError::SchemaFetchError { retryable, .. } => *retryable, } } @@ -261,6 +267,21 @@ pub(crate) fn should_retry_initial_connection( mod tests { use super::*; + #[test] + fn schema_fetch_error_retryable_classification() { + let retryable_err = ZerobusError::SchemaFetchError { + message: "503 service unavailable".to_string(), + retryable: true, + }; + assert!(retryable_err.is_retryable()); + + let non_retryable_err = ZerobusError::SchemaFetchError { + message: "404 not found".to_string(), + retryable: false, + }; + assert!(!non_retryable_err.is_retryable()); + } + #[test] fn initial_connection_auth_retry_is_one_shot() { let unauthenticated = diff --git a/rust/sdk/src/lib.rs b/rust/sdk/src/lib.rs index 177c6e3c..3cae2c2d 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 f05334ef..0bc8469b 100644 --- a/rust/sdk/src/sdk.rs +++ b/rust/sdk/src/sdk.rs @@ -115,6 +115,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..893d0760 --- /dev/null +++ b/rust/sdk/src/uc_schema.rs @@ -0,0 +1,403 @@ +//! 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. Compatible +//! schema evolution may be accepted; incompatible changes fail stream creation +//! with [`ZerobusError::CreateStreamError`], so re-fetch the descriptor before +//! rebuilding the stream. + +use std::time::Duration; + +use prost_reflect::MessageDescriptor; +use tracing::debug; + +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); + +/// Maximum allowed response size in bytes for Unity Catalog HTTP responses (8 MiB). +const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; + +/// Maximum response body bytes to capture in error messages on non-success HTTP status. +const MAX_ERROR_SNIPPET_BYTES: usize = 4096; + +/// 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. +/// - [`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?; + 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 +/// 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| ZerobusError::SchemaFetchError { + message: format!("failed to build HTTP client: {e}"), + retryable: false, + })?; + + debug!(table = %table_name, "fetching UC table schema"); + let token = mint_metadata_token(&client, &base, client_id, client_secret).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 req = client + .get(url) + .bearer_auth(&token) + .header(reqwest::header::ACCEPT, "application/json"); + let body = read_bounded_response(req, "schema").await?; + + let schema: UcTableSchema = + serde_json::from_slice(&body).map_err(|e| ZerobusError::SchemaFetchError { + message: format!("could not parse Unity Catalog response: {e}"), + retryable: false, + })?; + if schema.columns.is_empty() { + return Err(ZerobusError::SchemaFetchError { + message: format!("Unity Catalog returned no columns for table '{table_name}'"), + retryable: false, + }); + } + Ok(schema) +} + +/// 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 req = client + .post(url) + .basic_auth(client_id, Some(client_secret)) + .form(¶ms); + let body = read_bounded_response(req, "token").await?; + + let body: serde_json::Value = + serde_json::from_slice(&body).map_err(|e| ZerobusError::SchemaFetchError { + message: format!("could not parse token response: {e}"), + retryable: false, + })?; + let token = body["access_token"] + .as_str() + .ok_or_else(|| ZerobusError::SchemaFetchError { + message: "token response has no access_token".to_string(), + retryable: 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(ZerobusError::SchemaFetchError { + message: "token response contains an unusable access_token".to_string(), + retryable: false, + }); + } + Ok(token.to_string()) +} + +/// Send a request and stream the response body with size limits and error classification. +async fn read_bounded_response( + request: reqwest::RequestBuilder, + operation: &str, +) -> ZerobusResult> { + let response = request.send().await.map_err(|e| { + let retryable = is_reqwest_error_retryable(&e); + ZerobusError::SchemaFetchError { + message: format!("{operation} request failed: {e}"), + retryable, + } + })?; + + let status = response.status(); + if !status.is_success() { + let retryable = + status.is_server_error() || status == reqwest::StatusCode::TOO_MANY_REQUESTS; + let err_body = read_error_snippet(response, MAX_ERROR_SNIPPET_BYTES).await; + let body_str = String::from_utf8_lossy(&err_body); + let trimmed = body_str.trim(); + let message = if trimmed.is_empty() { + format!("{operation} request failed with status {status}") + } else { + format!("{operation} request failed with status {status}: {trimmed}") + }; + return Err(ZerobusError::SchemaFetchError { message, retryable }); + } + + if let Some(content_length) = response.content_length() { + if content_length > MAX_RESPONSE_BYTES as u64 { + return Err(ZerobusError::SchemaFetchError { + message: format!( + "{operation} response exceeded the size limit of {MAX_RESPONSE_BYTES} bytes (Content-Length: {content_length})" + ), + retryable: false, + }); + } + } + + read_body_chunks(response, MAX_RESPONSE_BYTES, operation).await +} + +async fn read_body_chunks( + mut response: reqwest::Response, + limit: usize, + operation: &str, +) -> ZerobusResult> { + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|e| { + let retryable = is_reqwest_error_retryable(&e); + ZerobusError::SchemaFetchError { + message: format!("reading {operation} response failed: {e}"), + retryable, + } + })? { + if body.len().saturating_add(chunk.len()) > limit { + return Err(ZerobusError::SchemaFetchError { + message: format!("{operation} response exceeded the size limit of {limit} bytes"), + retryable: false, + }); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + +async fn read_error_snippet(mut response: reqwest::Response, limit: usize) -> Vec { + let mut body = Vec::new(); + while let Ok(Some(chunk)) = response.chunk().await { + let remaining = limit.saturating_sub(body.len()); + if remaining == 0 { + break; + } + let to_take = chunk.len().min(remaining); + body.extend_from_slice(&chunk[..to_take]); + if body.len() >= limit { + break; + } + } + body +} + +fn is_reqwest_error_retryable(error: &reqwest::Error) -> bool { + if error.is_timeout() || error.is_connect() { + return true; + } + if let Some(status) = error.status() { + return status.is_server_error() || status == reqwest::StatusCode::TOO_MANY_REQUESTS; + } + !error.is_builder() && !error.is_redirect() +} + +/// 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; 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!("invalid Unity Catalog URL: {e}")) + })?; + if !matches!(url.scheme(), "http" | "https") || !url.has_host() { + return Err(ZerobusError::InvalidUCEndpointError( + "invalid Unity Catalog URL: expected an http or https URL with a host".to_string(), + )); + } + // 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::*; + + #[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 normalize_endpoint_does_not_leak_credentials_on_parse_error() { + let bad = "https://user:secret-pass@/no-host"; + let err = normalize_endpoint(bad).unwrap_err(); + let msg = err.to_string(); + assert!(!msg.contains("user"), "must not leak username in: {msg}"); + assert!( + !msg.contains("secret-pass"), + "must not leak password in: {msg}" + ); + } + + #[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" + ); + } +} diff --git a/rust/tests/Cargo.toml b/rust/tests/Cargo.toml index f76a2854..cf939f17 100644 --- a/rust/tests/Cargo.toml +++ b/rust/tests/Cargo.toml @@ -24,6 +24,10 @@ path = "src/arrow_tests.rs" name = "arrow_c_data_ffi_tests" path = "src/arrow_c_data_ffi_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..8265102a --- /dev/null +++ b/rust/tests/src/uc_schema_tests.rs @@ -0,0 +1,247 @@ +//! 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: impl Into) -> MockUc { + let schema_body: Arc = Arc::from(schema_body.into()); + 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); + let schema_body = Arc::clone(&schema_body); + tokio::spawn(async move { + 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}"#.to_string(), + ) + } else { + (schema_status, schema_body.to_string()) + }; + 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} + ]}"# +} + +fn table_json_with_padding(padding_len: usize) -> String { + let padding = "x".repeat(padding_len); + format!( + r#"{{"name":"orders","catalog_name":"main","schema_name":"sales","padding":"{padding}","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_404_is_non_retryable_schema_fetch_error() { + let mock = start_mock(404, "table not found").await; + + let err = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .unwrap_err(); + match &err { + ZerobusError::SchemaFetchError { + message, retryable, .. + } => { + assert!(message.contains("404"), "got: {message}"); + assert!(message.contains("table not found"), "got: {message}"); + assert!(!retryable); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert!(!err.is_retryable()); +} + +#[tokio::test] +async fn schema_request_503_is_retryable_schema_fetch_error() { + let mock = start_mock(503, "service unavailable").await; + + let err = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .unwrap_err(); + match &err { + ZerobusError::SchemaFetchError { + message, retryable, .. + } => { + assert!(message.contains("503"), "got: {message}"); + assert!(message.contains("service unavailable"), "got: {message}"); + assert!(*retryable); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert!(err.is_retryable()); +} + +#[tokio::test] +async fn schema_request_429_is_retryable_schema_fetch_error() { + let mock = start_mock(429, "rate limit exceeded").await; + + let err = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .unwrap_err(); + match &err { + ZerobusError::SchemaFetchError { + message, retryable, .. + } => { + assert!(message.contains("429"), "got: {message}"); + assert!(message.contains("rate limit exceeded"), "got: {message}"); + assert!(*retryable); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert!(err.is_retryable()); +} + +#[tokio::test] +async fn oversized_schema_response_is_rejected() { + // 8 MiB + 1 byte body + let oversized_body = table_json_with_padding(8 * 1024 * 1024 + 1); + let mock = start_mock(200, oversized_body).await; + + let err = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .unwrap_err(); + match &err { + ZerobusError::SchemaFetchError { + message, retryable, .. + } => { + assert!( + message.contains("size limit") || message.contains("exceeded"), + "got: {message}" + ); + assert!(!retryable); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert!(!err.is_retryable()); +} + +#[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" + ); +} + +#[tokio::test] +async fn empty_columns_returns_non_retryable_schema_fetch_error() { + let empty_schema = + r#"{"name":"orders","catalog_name":"main","schema_name":"sales","columns":[]}"#; + let mock = start_mock(200, empty_schema).await; + + let err = fetch_message_descriptor(&mock.url, TABLE, "cid", "csec") + .await + .unwrap_err(); + match &err { + ZerobusError::SchemaFetchError { + message, retryable, .. + } => { + assert!(message.contains("no columns"), "got: {message}"); + assert!(!retryable); + } + other => panic!("expected SchemaFetchError, got {other:?}"), + } + assert!(!err.is_retryable()); +}