Skip to content

[Rust] Add support to fetch schema from UC for dynamic proto - #704

Open
andrijast-db wants to merge 3 commits into
mainfrom
andrijast-db/effort/rust-dynamic-proto-fetch-from-uc
Open

[Rust] Add support to fetch schema from UC for dynamic proto#704
andrijast-db wants to merge 3 commits into
mainfrom
andrijast-db/effort/rust-dynamic-proto-fetch-from-uc

Conversation

@andrijast-db

@andrijast-db andrijast-db commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

What changes are proposed in this pull request?

Follow-up to #465, closing #521. Adds the ability to fetch a table's schema from Unity Catalog and use it as a dynamic-proto descriptor, so the descriptor no longer has to be built column-by-column in code.

New public API (all additive):

  • uc_schema module — fetch_message_descriptor(...) mints an OAuth token, reads GET /api/2.1/unity-catalog/tables/{name}, and resolves it to a MessageDescriptor; fetch_table_schema(...) returns the raw schema.
  • ZerobusSdk::fetch_message_descriptor(table, client_id, client_secret) — wrapper using the SDK's unity_catalog_url.
  • ZerobusError::SchemaFetchError { message, retryable }.

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?;

The fetch is a separate step feeding the existing .dynamic_proto(...) — rather than a new builder method that fetches inside build() — so .dynamic_proto() stays the single format selector, network I/O stays visible, and one descriptor can be reused across streams. It mints a plain all-apis token (the ingestion token DefaultTokenFactory produces is rejected by the UC REST API), bounds response body size, classifies transport/5xx/429 as retryable, and rejects credential-bearing URLs. No FFI/wrapper changes.

How is this tested?

Full cargo test --workspace passing; clippy (--all --all-features -D warnings) and fmt clean. Unit tests cover validation, endpoint/URL handling, and descriptor conversion; 20 integration tests run against a loopback HTTP mock covering the happy path, request shapes, and error classification (4xx/5xx/429, unparseable/oversized body, dropped connection). Not tested against a live workspace.

@andrijast-db
andrijast-db force-pushed the andrijast-db/effort/rust-dynamic-proto-fetch-from-uc branch from ba4fc6b to 7593a4b Compare August 10, 2026 15:02
Comment thread rust/sdk/src/uc_schema.rs
.await
.and_then(reqwest::Response::error_for_status)
.map_err(|e| fetch_error(format!("schema request failed: {e}")))?
.bytes()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both UC responses are buffered without a size limit. A large schema response here, or a large token response at rust/sdk/src/uc_schema.rs:147, can exhaust the caller's memory before parsing. The request timeout limits elapsed time, but it does not limit how many bytes a fast or compressed response can allocate.

Could we route both responses through one bounded reader that rejects an oversized Content-Length, enforces the same limit while reading chunks, and preserves a short error body for non-success statuses? For example:

const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;

let mut body = Vec::new();
while let Some(chunk) = response
    .chunk()
    .await
    .map_err(|e| fetch_error(format!("reading UC response failed: {e}")))?
{
    if body.len().saturating_add(chunk.len()) > MAX_RESPONSE_BYTES {
        return Err(fetch_error(
            "UC response exceeded the size limit".to_string(),
        ));
    }
    body.extend_from_slice(&chunk);
}

The loopback test can return otherwise valid schema JSON with an ignored padding property that puts the body one byte over the limit. It should succeed before the cap is added and return SchemaFetchError afterwards.

Comment thread rust/sdk/src/errors.rs
ZerobusError::ConnectionTimeout(_) => true,
ZerobusError::TokenFetchError(_) => true,
// A schema fetch is a one-shot setup step, not part of the recovery loop.
ZerobusError::SchemaFetchError(_) => false,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SchemaFetchError is always non-retryable, so a connection reset, timeout, HTTP 429, or HTTP 503 tells callers to give up just like a bad request or missing table. The fetch site still has the transport error or status needed to make this distinction, but fetch_error flattens it into a string. This also leaves the new tuple variant unable to gain retry metadata later without a breaking API change.

Could we keep the classification in a non-exhaustive struct variant and set it where the underlying failure is available?

#[error("Failed to fetch table schema from Unity Catalog: {message}.")]
#[non_exhaustive]
SchemaFetchError {
    message: String,
    retryable: bool,
}

ZerobusError::SchemaFetchError { retryable, .. } => *retryable,

Transport failures, timeouts, HTTP 429, and HTTP 5xx can then be retryable while rejected requests and malformed bodies remain terminal. The current mock already accepts arbitrary schema statuses, so a 503 test can assert is_retryable() is true and the existing 404 case can assert it remains false.

Comment thread rust/sdk/src/uc_schema.rs
//!
//! 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The dynamic protobuf path does not surface stream-creation schema failures as ZerobusError::InvalidSchema. Its gRPC setup returns ZerobusError::CreateStreamError, while InvalidSchema classification is only used by Arrow Flight setup. The blanket statement about every table change is also too strong: compatible evolution such as adding a nullable column can be accepted, while incompatible changes are rejected. The same claim is repeated in rust/README.md:611.

Could we document the behavior callers can actually match, or add equivalent gRPC error classification before promising InvalidSchema? If the implementation stays as-is, wording along these lines would avoid steering recovery code to the wrong variant:

//! 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.

Comment thread rust/sdk/src/uc_schema.rs
};

let url = reqwest::Url::parse(&candidate)
.map_err(|e| ZerobusError::InvalidUCEndpointError(format!("{unity_catalog_url}: {e}")))?;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The credential check below this line protects successfully parsed URLs, but a parse failure includes the original unity_catalog_url verbatim. A malformed credential-bearing value can therefore put its username and password into InvalidUCEndpointError, which is likely to be logged.

Could we omit the raw input from parse errors and only include the parser's sanitized reason?

let url = reqwest::Url::parse(&candidate).map_err(|e| {
    ZerobusError::InvalidUCEndpointError(format!("invalid Unity Catalog URL: {e}"))
})?;

The existing normalize_endpoint unit tests can add a malformed credential-bearing URL and assert that neither credential appears in the rendered error.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants