From cfab767693cafdeb20eadbeedc933b7d8f60fbf0 Mon Sep 17 00:00:00 2001 From: Anil Menon Date: Sun, 16 Aug 2026 17:41:39 +0200 Subject: [PATCH 1/2] [Rust] Add federated-token (RFC 8693) auth Add first-class external-IdP (e.g. Entra ID) token federation to the Rust core as an opt-in auth mode, alongside the existing OAuth client-credentials path. No existing signatures change. - default_token_factory: factor the Zerobus-scoped request shaping (scope, resource, table-scoped authorization_details) into shared helpers so the client-credentials grant and the new token-exchange grant build an identical request, keeping the two at parity. - headers_provider: add FederatedTokenProvider (implements the existing HeadersProvider trait, including invalidate()) plus the IdpTokenSupplier callback type. It exchanges the current external IdP token via RFC 8693, caches the exchanged token, and supports both account-level federation (no client_id, SCIM) and workload identity federation (client_id, no secret) through one client_id toggle. - token_cache: reused unchanged; federated tokens key by (client_id-or-none, table) so the two modes cache independently. - stream_builder: add opt-in .federated() and .federated_with_client_id() builder methods. Default auth paths are unchanged. Tests: request-shaping parity with/without client_id, and end-to-end provider tests (caching, invalidate re-mint, mode independence) against a mock token endpoint. All lib tests pass; clippy and fmt clean. Signed-off-by: Anil Menon --- rust/NEXT_CHANGELOG.md | 12 + rust/README.md | 37 +++ rust/sdk/src/builder/stream_builder.rs | 61 +++- rust/sdk/src/default_token_factory.rs | 308 +++++++++++++++-- rust/sdk/src/headers_provider.rs | 437 +++++++++++++++++++++++++ rust/sdk/src/lib.rs | 4 +- 6 files changed, 833 insertions(+), 26 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 8cdb50e8..ef09bc9c 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -6,6 +6,18 @@ ### New Features and Improvements +- Added first-class external-IdP token federation (`FederatedTokenProvider`, + `IdpTokenSupplier`) alongside the existing OAuth client-credentials path. It + exchanges an external IdP token (for example an Entra ID token) for a + Zerobus-scoped Databricks token via the RFC 8693 token-exchange grant, caches + and refreshes it through the existing `TokenCache`, and supports both + account-level federation (no `client_id`, identity synced via Automatic + Identity Management) and workload identity federation (a service principal + `client_id` with no secret). Opt in via `StreamBuilder::federated(...)` or + `StreamBuilder::federated_with_client_id(...)`. The client-credentials and + token-exchange grants now share one request-shaping path, keeping them at + parity. Existing `oauth(...)` and `headers_provider(...)` paths are unchanged. + ### Bug Fixes ### Documentation diff --git a/rust/README.md b/rust/README.md index bf291984..cf74a2e8 100644 --- a/rust/README.md +++ b/rust/README.md @@ -487,6 +487,43 @@ let client_secret = "your-client-secret".to_string(); See [`examples/README.md`](https://github.com/databricks/zerobus-sdk/blob/main/rust/examples/README.md) for more information on how to get these credentials. +#### External-IdP federation (e.g. Entra ID) + +To authenticate with an external identity provider instead of a Databricks +OAuth secret, use the federated builder methods. You provide an +[`IdpTokenSupplier`] callback that returns the current external IdP token; the +SDK exchanges it for a Zerobus-scoped Databricks token (RFC 8693 token +exchange) and caches and refreshes it, exactly like the OAuth path. + +```rust,ignore +use std::sync::Arc; + +// Account-level federation: no Databricks service principal. The identity is +// synced into Databricks via Automatic Identity Management (SCIM). +let stream = sdk + .stream_builder() + .table("catalog.schema.table") + .federated(Arc::new(|| Box::pin(async { get_idp_token().await }))) + .json() + .build() + .await?; + +// Workload identity federation: a Databricks service principal with a client_id +// and no secret, with a federation policy attached. +let stream = sdk + .stream_builder() + .table("catalog.schema.table") + .federated_with_client_id( + Arc::new(|| Box::pin(async { get_idp_token().await })), + "", + ) + .json() + .build() + .await?; +``` + +The existing `.oauth(...)` and `.headers_provider(...)` paths are unchanged. + ### 4. Create a Stream Use the `stream_builder()` API to create a stream: diff --git a/rust/sdk/src/builder/stream_builder.rs b/rust/sdk/src/builder/stream_builder.rs index 5514667a..d34c3bf9 100644 --- a/rust/sdk/src/builder/stream_builder.rs +++ b/rust/sdk/src/builder/stream_builder.rs @@ -25,7 +25,9 @@ use crate::callbacks::AckCallback; use crate::databricks::zerobus::RecordType; #[cfg(feature = "testing")] use crate::headers_provider::NoAuthHeadersProvider; -use crate::headers_provider::{HeadersProvider, OAuthHeadersProvider}; +use crate::headers_provider::{ + FederatedTokenProvider, HeadersProvider, IdpTokenSupplier, OAuthHeadersProvider, +}; use crate::stream_configuration::StreamConfigurationOptions; use crate::{ MessageDescriptor, TableProperties, ZerobusError, ZerobusResult, ZerobusSdk, ZerobusStream, @@ -42,6 +44,13 @@ enum AuthConfig { client_id: String, client_secret: String, }, + /// External-IdP federation via RFC 8693 token exchange. `client_id` is + /// `Some` for workload identity federation and `None` for account-level + /// federation. + Federated { + idp_token_supplier: IdpTokenSupplier, + client_id: Option, + }, HeadersProvider(Arc), #[cfg(feature = "testing")] NoAuth, @@ -109,6 +118,7 @@ impl fmt::Debug for StreamBuilder<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let auth_kind = match &self.auth { Some(AuthConfig::OAuth { .. }) => "OAuth", + Some(AuthConfig::Federated { .. }) => "Federated", Some(AuthConfig::HeadersProvider(_)) => "HeadersProvider", #[cfg(feature = "testing")] Some(AuthConfig::NoAuth) => "NoAuth", @@ -170,6 +180,44 @@ impl<'a> StreamBuilder<'a> { self } + /// Authenticate with account-level external-IdP federation (RFC 8693 token + /// exchange), with no Databricks-managed service principal. + /// + /// The `idp_token_supplier` is an async callback that returns the current + /// external IdP token (e.g. an Entra ID JWT). The SDK exchanges it for a + /// Zerobus-scoped Databricks token; the token's subject is resolved to an + /// identity synced into Databricks via Automatic Identity Management (SCIM). + /// Use [`federated_with_client_id`](Self::federated_with_client_id) for + /// workload identity federation (a service principal with a client_id and + /// no secret). + pub fn federated(mut self, idp_token_supplier: IdpTokenSupplier) -> Self { + self.auth = Some(AuthConfig::Federated { + idp_token_supplier, + client_id: None, + }); + self + } + + /// Authenticate with workload identity federation (RFC 8693 token exchange) + /// for a Databricks service principal that has a `client_id` and no secret, + /// with a federation policy attached. + /// + /// The `idp_token_supplier` returns the current external IdP token; the + /// exchange request names the service principal via `client_id`. Use + /// [`federated`](Self::federated) for account-level federation (no service + /// principal). + pub fn federated_with_client_id( + mut self, + idp_token_supplier: IdpTokenSupplier, + client_id: impl Into, + ) -> Self { + self.auth = Some(AuthConfig::Federated { + idp_token_supplier, + client_id: Some(client_id.into()), + }); + self + } + /// Authenticate with a custom headers provider. pub fn headers_provider(mut self, provider: Arc) -> Self { self.auth = Some(AuthConfig::HeadersProvider(provider)); @@ -401,6 +449,17 @@ impl<'a> StreamBuilder<'a> { self.sdk.unity_catalog_url.clone(), Arc::clone(&self.sdk.token_cache), ))), + Some(AuthConfig::Federated { + idp_token_supplier, + client_id, + }) => Ok(Arc::new(FederatedTokenProvider::with_cache( + client_id.clone(), + Arc::clone(idp_token_supplier), + self.table_name.clone(), + self.sdk.workspace_id.clone(), + self.sdk.unity_catalog_url.clone(), + Arc::clone(&self.sdk.token_cache), + ))), Some(AuthConfig::HeadersProvider(p)) => Ok(Arc::clone(p)), #[cfg(feature = "testing")] Some(AuthConfig::NoAuth) => Ok(Arc::new(NoAuthHeadersProvider)), diff --git a/rust/sdk/src/default_token_factory.rs b/rust/sdk/src/default_token_factory.rs index 0b178cd7..de704886 100644 --- a/rust/sdk/src/default_token_factory.rs +++ b/rust/sdk/src/default_token_factory.rs @@ -88,7 +88,8 @@ impl DefaultTokenFactory { /// /// This is the caching-aware variant of [`get_token`](Self::get_token): in /// addition to the token it returns the `expires_in` value from the OAuth - /// response so callers can cache the token until it nears expiry. + /// response so callers can cache the token until it nears expiry. Uses the + /// client-credentials grant (client_id + secret). pub(crate) async fn fetch_token( uc_endpoint: &str, table_name: &str, @@ -97,7 +98,7 @@ impl DefaultTokenFactory { workspace_id: &str, reason: MintReason, ) -> ZerobusResult { - debug!(table = %table_name, "requesting UC OAuth token"); + debug!(table = %table_name, "requesting Zerobus token (client_credentials)"); let started = Instant::now(); let result = Self::fetch_token_inner( uc_endpoint, @@ -107,37 +108,99 @@ impl DefaultTokenFactory { workspace_id, ) .await; - let elapsed_ms = started.elapsed().as_millis() as u64; - match &result { + Self::log_mint_outcome( + table_name, + reason, + started.elapsed().as_millis() as u64, + &result, + "client_credentials", + ); + result + } + + /// Obtains a Databricks access token by exchanging an external IdP token + /// (e.g. an Entra ID JWT) for a Zerobus-scoped Databricks token via the + /// RFC 8693 token-exchange grant. + /// + /// Used by [`FederatedTokenProvider`](crate::FederatedTokenProvider). The + /// request is shaped identically to the client-credentials grant (same + /// Zerobus resource, scope, and table-scoped authorization details); only + /// the grant-specific parameters differ: `grant_type=token-exchange`, the + /// `subject_token` carrying the IdP JWT, and — for workload identity + /// federation — the Databricks service principal `client_id`. `client_id` + /// is `None` for account-level federation (identity resolved via SCIM). + pub(crate) async fn fetch_exchanged_token( + uc_endpoint: &str, + table_name: &str, + client_id: Option<&str>, + subject_token: &str, + workspace_id: &str, + reason: MintReason, + ) -> ZerobusResult { + debug!(table = %table_name, "requesting Zerobus token (token-exchange)"); + let started = Instant::now(); + let result = Self::fetch_exchanged_token_inner( + uc_endpoint, + table_name, + client_id, + subject_token, + workspace_id, + ) + .await; + Self::log_mint_outcome( + table_name, + reason, + started.elapsed().as_millis() as u64, + &result, + "token_exchange", + ); + result + } + + /// Emits the structured mint log shared by every grant type. `grant` + /// distinguishes the client-credentials path from the token-exchange path. + fn log_mint_outcome( + table_name: &str, + reason: MintReason, + elapsed_ms: u64, + result: &ZerobusResult, + grant: &'static str, + ) { + match result { Ok(FetchedToken { expires_in: Some(ttl), .. }) => info!( table = %table_name, reason = reason.as_str(), + grant, expires_in_secs = ttl.as_secs(), elapsed_ms, - "minted UC OAuth token" + "minted Zerobus token" ), Ok(FetchedToken { expires_in: None, .. }) => warn!( table = %table_name, reason = reason.as_str(), + grant, elapsed_ms, - "minted UC OAuth token but UC returned no expires_in; token will not be cached" + "minted Zerobus token but UC returned no expires_in; token will not be cached" ), Err(err) => warn!( table = %table_name, reason = reason.as_str(), + grant, retryable = err.is_retryable(), elapsed_ms, - "failed to mint UC OAuth token: {err}" + "failed to mint Zerobus token: {err}" ), } - result } + /// Client-credentials grant: builds the shared Zerobus-scoped request and + /// adds `grant_type=client_credentials`, authenticating with HTTP Basic + /// (client_id + secret). async fn fetch_token_inner( uc_endpoint: &str, table_name: &str, @@ -145,12 +208,82 @@ impl DefaultTokenFactory { client_secret: &str, workspace_id: &str, ) -> ZerobusResult { - let (catalog, schema, table) = Self::parse_table_name(table_name)?; + let params = Self::client_credentials_form_params(table_name, workspace_id)?; + Self::post_token_request(uc_endpoint, ¶ms, Some((client_id, client_secret))).await + } + + /// Builds the full client-credentials form parameters: the shared + /// Zerobus-scoped parameters plus `grant_type=client_credentials`. + #[allow(clippy::result_large_err)] + fn client_credentials_form_params( + table_name: &str, + workspace_id: &str, + ) -> ZerobusResult> { + let mut params = Self::zerobus_scoped_form_params(table_name, workspace_id)?; + params.push(("grant_type", "client_credentials".to_string())); + Ok(params) + } + + /// Token-exchange grant (RFC 8693): builds the same shared Zerobus-scoped + /// request and adds the exchange-specific parameters — `grant_type`, + /// `subject_token` (the external IdP JWT), `subject_token_type`, and, for + /// workload identity federation, the Databricks SP `client_id`. No HTTP + /// Basic auth: the subject token is the credential. + async fn fetch_exchanged_token_inner( + uc_endpoint: &str, + table_name: &str, + client_id: Option<&str>, + subject_token: &str, + workspace_id: &str, + ) -> ZerobusResult { + let params = + Self::exchange_form_params(table_name, client_id, subject_token, workspace_id)?; + Self::post_token_request(uc_endpoint, ¶ms, None).await + } - let uc_endpoint = uc_endpoint.to_string(); - let databricks_client_id = client_id.to_string(); - let databricks_client_secret = client_secret.to_string(); - let workspace_id = workspace_id.to_string(); + /// Builds the full RFC 8693 token-exchange form parameters: the shared + /// Zerobus-scoped parameters plus the exchange-specific parameters. The + /// Databricks SP `client_id` is included only for workload identity + /// federation (Story 2) and omitted for account-level federation (Story 1). + #[allow(clippy::result_large_err)] + fn exchange_form_params( + table_name: &str, + client_id: Option<&str>, + subject_token: &str, + workspace_id: &str, + ) -> ZerobusResult> { + let mut params = Self::zerobus_scoped_form_params(table_name, workspace_id)?; + params.push(( + "grant_type", + "urn:ietf:params:oauth:grant-type:token-exchange".to_string(), + )); + params.push(("subject_token", subject_token.to_string())); + params.push(( + "subject_token_type", + "urn:ietf:params:oauth:token-type:jwt".to_string(), + )); + // Present for workload identity federation (Story 2), naming the + // Databricks service principal; omitted for account-level federation + // (Story 1), where the subject is resolved to a SCIM-synced identity. + if let Some(client_id) = client_id { + params.push(("client_id", client_id.to_string())); + } + Ok(params) + } + + /// Builds the Zerobus-scoped OAuth form parameters shared by every grant + /// type: `scope=all-apis`, the `zerobusDirectWriteApi` resource for this + /// workspace, and the table-scoped Unity Catalog `authorization_details`. + /// Both the client-credentials grant and the RFC 8693 token-exchange grant + /// send an identical Zerobus-scoped request; only the grant-specific + /// parameters (grant_type, credentials/subject_token) differ. Keeping this + /// in one place keeps the two grants at parity. + #[allow(clippy::result_large_err)] + fn zerobus_scoped_form_params( + table_name: &str, + workspace_id: &str, + ) -> ZerobusResult> { + let (catalog, schema, table) = Self::parse_table_name(table_name)?; let authorization_details = serde_json::json!([ { @@ -174,27 +307,37 @@ impl DefaultTokenFactory { } ]); - let client = reqwest::Client::new(); - - let params = [ - ("grant_type", "client_credentials".to_string()), + Ok(vec![ ("scope", "all-apis".to_string()), ( "resource", format!( "api://databricks/workspaces/{}/zerobusDirectWriteApi", workspace_id - ) - .to_string(), + ), ), ("authorization_details", authorization_details.to_string()), - ]; + ]) + } + /// Posts a token request to the UC OIDC endpoint and parses the response. + /// Shared by every grant: `basic_auth` carries the client-credentials + /// Basic header when present, and is `None` for the token-exchange grant. + async fn post_token_request( + uc_endpoint: &str, + params: &[(&str, String)], + basic_auth: Option<(&str, &str)>, + ) -> ZerobusResult { + let client = reqwest::Client::new(); let token_endpoint = format!("{}/oidc/v1/token", uc_endpoint); - let resp = client - .post(&token_endpoint) - .basic_auth(databricks_client_id, Some(databricks_client_secret)) - .form(¶ms) + + let mut request = client.post(&token_endpoint); + if let Some((client_id, client_secret)) = basic_auth { + request = request.basic_auth(client_id, Some(client_secret)); + } + + let resp = request + .form(params) .send() .await .map_err(Self::handle_http_error)?; @@ -383,6 +526,123 @@ mod tests { assert!(!DefaultTokenFactory::is_usable_as_header("bad\0token")); } + /// Looks up the single value for `key` in a form-param list, asserting it is + /// present exactly once. + fn param<'a>(params: &'a [(&str, String)], key: &str) -> &'a str { + let matches: Vec<&str> = params + .iter() + .filter(|(k, _)| *k == key) + .map(|(_, v)| v.as_str()) + .collect(); + assert_eq!( + matches.len(), + 1, + "expected exactly one '{}' param, found {}", + key, + matches.len() + ); + matches[0] + } + + fn has_param(params: &[(&str, String)], key: &str) -> bool { + params.iter().any(|(k, _)| *k == key) + } + + /// The Zerobus-scoped parameters (scope, resource, authorization_details) + /// must be byte-identical between the client-credentials grant and the + /// token-exchange grant. This is the parity guarantee the refactor exists to + /// enforce. + #[test] + fn client_credentials_and_exchange_share_identical_zerobus_scoping() { + let table = "cat.sch.tbl"; + let workspace = "1234567890"; + + let cc = DefaultTokenFactory::client_credentials_form_params(table, workspace).unwrap(); + let ex = + DefaultTokenFactory::exchange_form_params(table, Some("sp-id"), "idp-jwt", workspace) + .unwrap(); + + for key in ["scope", "resource", "authorization_details"] { + assert_eq!( + param(&cc, key), + param(&ex, key), + "'{}' must match across grants", + key + ); + } + } + + #[test] + fn client_credentials_form_params_shape() { + let params = + DefaultTokenFactory::client_credentials_form_params("cat.sch.tbl", "42").unwrap(); + + assert_eq!(param(¶ms, "grant_type"), "client_credentials"); + assert_eq!(param(¶ms, "scope"), "all-apis"); + assert_eq!( + param(¶ms, "resource"), + "api://databricks/workspaces/42/zerobusDirectWriteApi" + ); + // The token-exchange-only params must never appear on this grant. + assert!(!has_param(¶ms, "subject_token")); + assert!(!has_param(¶ms, "subject_token_type")); + + // authorization_details is downscoped to the specific table. + let details: serde_json::Value = + serde_json::from_str(param(¶ms, "authorization_details")).unwrap(); + assert_eq!(details[2]["object_full_path"], "cat.sch.tbl"); + assert_eq!(details[2]["operations"][0], "zerobuswrite"); + } + + #[test] + fn exchange_form_params_without_client_id_is_account_level() { + // Account-level federation (Story 1): no client_id in the request. + let params = + DefaultTokenFactory::exchange_form_params("cat.sch.tbl", None, "idp-jwt-token", "99") + .unwrap(); + + assert_eq!( + param(¶ms, "grant_type"), + "urn:ietf:params:oauth:grant-type:token-exchange" + ); + assert_eq!(param(¶ms, "subject_token"), "idp-jwt-token"); + assert_eq!( + param(¶ms, "subject_token_type"), + "urn:ietf:params:oauth:token-type:jwt" + ); + assert!( + !has_param(¶ms, "client_id"), + "account-level federation must omit client_id" + ); + // Client-credentials-only auth is never sent on the exchange grant. + assert!(!has_param(¶ms, "client_secret")); + } + + #[test] + fn exchange_form_params_with_client_id_is_workload_identity() { + // Workload identity federation (Story 2): client_id names the SP. + let params = DefaultTokenFactory::exchange_form_params( + "cat.sch.tbl", + Some("sp-client-id-uuid"), + "idp-jwt-token", + "99", + ) + .unwrap(); + + assert_eq!( + param(¶ms, "grant_type"), + "urn:ietf:params:oauth:grant-type:token-exchange" + ); + assert_eq!(param(¶ms, "client_id"), "sp-client-id-uuid"); + assert_eq!(param(¶ms, "subject_token"), "idp-jwt-token"); + } + + #[test] + fn exchange_form_params_rejects_bad_table_name() { + let err = DefaultTokenFactory::exchange_form_params("not_three_parts", None, "jwt", "1"); + assert!(matches!(err, Err(ZerobusError::InvalidTableName(_)))); + } + #[test] fn test_parse_table_name_invalid() { let invalid_cases = vec![ diff --git a/rust/sdk/src/headers_provider.rs b/rust/sdk/src/headers_provider.rs index 9c361f04..903f2678 100644 --- a/rust/sdk/src/headers_provider.rs +++ b/rust/sdk/src/headers_provider.rs @@ -3,6 +3,8 @@ use crate::token_cache::{TokenCache, DEFAULT_REFRESH_BUFFER}; use crate::ZerobusResult; use async_trait::async_trait; use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; /// A trait for providing custom headers for gRPC requests. @@ -152,6 +154,144 @@ impl HeadersProvider for OAuthHeadersProvider { } } +/// An async callback that yields the current external IdP token (for example an +/// Entra ID / OIDC JWT). +/// +/// The federated auth mode takes a supplier callback rather than a bare token +/// on purpose: an external IdP token is short-lived (typically ~1 hour), so a +/// bare token would strand the stream at the token's expiry with no way to +/// refresh. The supplier is invoked only when a fresh Databricks token must be +/// minted (a cache miss or a proactive refresh), never on every request, so a +/// cache hit incurs neither the callback nor the exchange. A caller that truly +/// holds a static token can wrap it in a trivial closure. +pub type IdpTokenSupplier = + Arc Pin> + Send>> + Send + Sync>; + +/// A headers provider that federates an external IdP token into a Zerobus-scoped +/// Databricks token via the RFC 8693 token-exchange grant. +/// +/// This is the first-class implementation of external-IdP (e.g. Entra ID) +/// federation. It supports the two supported federation modes through a single +/// `client_id` toggle: +/// +/// * **Account-level federation** (`client_id = None`): no Databricks-managed +/// service principal. The exchanged token's subject is resolved to an +/// identity synced into Databricks via Automatic Identity Management (SCIM). +/// * **Workload identity federation** (`client_id = Some(sp_id)`): a Databricks +/// service principal with a client_id and no secret, with a federation policy +/// attached. The exchange request names the service principal via `client_id`. +/// +/// It obtains the current IdP token from an [`IdpTokenSupplier`], performs the +/// exchange with the same request shaping as the client-credentials path, and +/// caches the exchanged Databricks token in the shared [`TokenCache`] keyed by +/// `(client_id-or-none, table)` so account-level and workload modes (and +/// distinct service principals) cache independently. +pub struct FederatedTokenProvider { + /// The Databricks service principal client_id for workload identity + /// federation, or `None` for account-level federation. + client_id: Option, + idp_token_supplier: IdpTokenSupplier, + table_name: String, + workspace_id: String, + unity_catalog_url: String, + token_cache: Arc, +} + +impl FederatedTokenProvider { + /// Creates a new `FederatedTokenProvider`. + /// + /// This standalone constructor caches tokens for the lifetime of the + /// returned provider only. When streams are created via + /// [`ZerobusSdk::stream_builder`](crate::ZerobusSdk::stream_builder) the SDK + /// supplies a shared cache so tokens are reused across streams; see + /// [`with_cache`](Self::with_cache). + pub fn new( + client_id: Option, + idp_token_supplier: IdpTokenSupplier, + table_name: String, + workspace_id: String, + unity_catalog_url: String, + ) -> Self { + Self::with_cache( + client_id, + idp_token_supplier, + table_name, + workspace_id, + unity_catalog_url, + Arc::new(TokenCache::new(true, DEFAULT_REFRESH_BUFFER)), + ) + } + + /// Creates a new `FederatedTokenProvider` backed by a shared token cache. + /// + /// Used internally so all streams created from one `ZerobusSdk` reuse cached + /// exchanged tokens rather than re-exchanging per stream. + pub(crate) fn with_cache( + client_id: Option, + idp_token_supplier: IdpTokenSupplier, + table_name: String, + workspace_id: String, + unity_catalog_url: String, + token_cache: Arc, + ) -> Self { + Self { + client_id, + idp_token_supplier, + table_name, + workspace_id, + unity_catalog_url, + token_cache, + } + } + + /// The cache key's client-id component: the service principal id for + /// workload identity federation, or the empty string for account-level + /// federation. There is no secret in either mode, so the secret component + /// of the key is always empty; the two modes therefore cache independently + /// by client id. + fn cache_client_id(&self) -> &str { + self.client_id.as_deref().unwrap_or("") + } +} + +#[async_trait] +impl HeadersProvider for FederatedTokenProvider { + async fn get_headers(&self) -> ZerobusResult> { + let token = self + .token_cache + .get_or_fetch( + self.cache_client_id(), + "", + &self.table_name, + |reason| async move { + // Only reached on a cache miss/refresh: fetch the current IdP + // token, then exchange it for a Zerobus-scoped Databricks token. + let idp_token = (self.idp_token_supplier)().await?; + DefaultTokenFactory::fetch_exchanged_token( + &self.unity_catalog_url, + &self.table_name, + self.client_id.as_deref(), + &idp_token, + &self.workspace_id, + reason, + ) + .await + }, + ) + .await?; + let mut headers = HashMap::new(); + headers.insert("authorization", format!("Bearer {}", token)); + headers.insert("x-databricks-zerobus-table-name", self.table_name.clone()); + Ok(headers) + } + + async fn invalidate(&self) { + self.token_cache + .invalidate(self.cache_client_id(), "", &self.table_name) + .await; + } +} + /// A headers provider that returns no headers. /// /// Intended only for local testing against a Zerobus endpoint not enforcing authentication. @@ -177,3 +317,300 @@ impl HeadersProvider for NoAuthHeadersProvider { Ok(HashMap::new()) } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// A minimal blocking HTTP mock of the UC `/oidc/v1/token` endpoint. It runs + /// on its own OS thread (blocking std IO) so the async test can drive the + /// reqwest-based exchange against it. Each request is answered with a fresh + /// `dbx-token-` so tests can tell a real mint from a cache hit, and every + /// request body is captured for assertions on the request shape. + struct MockTokenEndpoint { + base_url: String, + request_bodies: Arc>>, + mint_count: Arc, + } + + impl MockTokenEndpoint { + fn start() -> Self { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let request_bodies = Arc::new(std::sync::Mutex::new(Vec::new())); + let mint_count = Arc::new(AtomicUsize::new(0)); + + let bodies = Arc::clone(&request_bodies); + let count = Arc::clone(&mint_count); + // Detached daemon thread: it serves connections for the lifetime of + // the test process. Bounded accept is deliberately avoided so that a + // caching regression (too many mints) fails an assertion rather than + // deadlocking on a missing connection. + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { continue }; + let body = read_http_body(&mut stream); + let n = count.fetch_add(1, Ordering::SeqCst); + bodies.lock().unwrap().push(body); + write_json_response(&mut stream, &format!("dbx-token-{n}")); + } + }); + + Self { + base_url, + request_bodies, + mint_count, + } + } + + fn mint_count(&self) -> usize { + self.mint_count.load(Ordering::SeqCst) + } + + fn last_request_body(&self) -> String { + self.request_bodies.lock().unwrap().last().cloned().unwrap() + } + } + + fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + haystack.windows(needle.len()).position(|w| w == needle) + } + + fn read_http_body(stream: &mut std::net::TcpStream) -> String { + use std::io::Read; + let mut buf = Vec::new(); + let mut tmp = [0u8; 2048]; + loop { + let n = stream.read(&mut tmp).unwrap_or(0); + if n == 0 { + break; + } + buf.extend_from_slice(&tmp[..n]); + if let Some(pos) = find_subslice(&buf, b"\r\n\r\n") { + let header_str = String::from_utf8_lossy(&buf[..pos]).to_string(); + let content_length = header_str + .lines() + .find_map(|line| { + let lower = line.to_ascii_lowercase(); + lower + .strip_prefix("content-length:") + .map(|v| v.trim().parse::().unwrap_or(0)) + }) + .unwrap_or(0); + let body_start = pos + 4; + while buf.len() < body_start + content_length { + let n = stream.read(&mut tmp).unwrap_or(0); + if n == 0 { + break; + } + buf.extend_from_slice(&tmp[..n]); + } + let end = (body_start + content_length).min(buf.len()); + return String::from_utf8_lossy(&buf[body_start..end]).to_string(); + } + } + String::new() + } + + fn write_json_response(stream: &mut std::net::TcpStream, access_token: &str) { + use std::io::Write; + let body = format!(r#"{{"access_token":"{access_token}","expires_in":3600}}"#); + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + let _ = stream.write_all(response.as_bytes()); + let _ = stream.flush(); + } + + /// Builds an [`IdpTokenSupplier`] that returns `token` and counts its calls, + /// so tests can assert the supplier is invoked only on a real mint. + fn counting_supplier(token: &'static str, calls: Arc) -> IdpTokenSupplier { + Arc::new(move || { + let calls = Arc::clone(&calls); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(token.to_string()) + }) + }) + } + + const TOKEN_EXCHANGE_GRANT_ENCODED: &str = + "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange"; + + #[tokio::test] + async fn account_level_exchanges_and_returns_headers() { + let mock = MockTokenEndpoint::start(); + let idp_calls = Arc::new(AtomicUsize::new(0)); + let provider = FederatedTokenProvider::new( + None, + counting_supplier("entra-jwt", Arc::clone(&idp_calls)), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + ); + + let headers = provider.get_headers().await.unwrap(); + + assert_eq!(headers.get("authorization").unwrap(), "Bearer dbx-token-0"); + assert_eq!( + headers.get("x-databricks-zerobus-table-name").unwrap(), + "cat.sch.tbl" + ); + assert_eq!(mock.mint_count(), 1); + assert_eq!(idp_calls.load(Ordering::SeqCst), 1); + + // The exchange request carried the RFC 8693 grant and the IdP token, and + // omitted client_id (account-level federation). + let body = mock.last_request_body(); + assert!(body.contains(TOKEN_EXCHANGE_GRANT_ENCODED), "body: {body}"); + assert!(body.contains("subject_token=entra-jwt"), "body: {body}"); + assert!(!body.contains("client_id="), "body: {body}"); + } + + #[tokio::test] + async fn workload_identity_sends_client_id() { + let mock = MockTokenEndpoint::start(); + let idp_calls = Arc::new(AtomicUsize::new(0)); + let provider = FederatedTokenProvider::new( + Some("sp-uuid".to_string()), + counting_supplier("entra-jwt", Arc::clone(&idp_calls)), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + ); + + provider.get_headers().await.unwrap(); + + let body = mock.last_request_body(); + assert!(body.contains(TOKEN_EXCHANGE_GRANT_ENCODED), "body: {body}"); + assert!(body.contains("client_id=sp-uuid"), "body: {body}"); + } + + #[tokio::test] + async fn caches_exchanged_token_across_calls() { + let mock = MockTokenEndpoint::start(); + let idp_calls = Arc::new(AtomicUsize::new(0)); + let provider = FederatedTokenProvider::new( + None, + counting_supplier("entra-jwt", Arc::clone(&idp_calls)), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + ); + + let first = provider.get_headers().await.unwrap(); + let second = provider.get_headers().await.unwrap(); + + assert_eq!(first.get("authorization"), second.get("authorization")); + assert_eq!(mock.mint_count(), 1, "second call must reuse cached token"); + assert_eq!( + idp_calls.load(Ordering::SeqCst), + 1, + "IdP supplier must not be called on a cache hit" + ); + } + + #[tokio::test] + async fn invalidate_forces_remint() { + let mock = MockTokenEndpoint::start(); + let idp_calls = Arc::new(AtomicUsize::new(0)); + let provider = FederatedTokenProvider::new( + None, + counting_supplier("entra-jwt", Arc::clone(&idp_calls)), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + ); + + let first = provider.get_headers().await.unwrap(); + provider.invalidate().await; + let second = provider.get_headers().await.unwrap(); + + assert_eq!(first.get("authorization").unwrap(), "Bearer dbx-token-0"); + assert_eq!(second.get("authorization").unwrap(), "Bearer dbx-token-1"); + assert_eq!(mock.mint_count(), 2, "invalidate must force a re-mint"); + assert_eq!(idp_calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn account_level_and_workload_cache_independently() { + let mock = MockTokenEndpoint::start(); + // A shared cache, as the SDK supplies to every stream from one instance. + let cache = Arc::new(TokenCache::new(true, DEFAULT_REFRESH_BUFFER)); + + let account_level = FederatedTokenProvider::with_cache( + None, + counting_supplier("entra-jwt", Arc::new(AtomicUsize::new(0))), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + Arc::clone(&cache), + ); + let workload = FederatedTokenProvider::with_cache( + Some("sp-uuid".to_string()), + counting_supplier("entra-jwt", Arc::new(AtomicUsize::new(0))), + "cat.sch.tbl".to_string(), + "12345".to_string(), + mock.base_url.clone(), + Arc::clone(&cache), + ); + + // Same table, same shared cache, but different client_id => two mints. + account_level.get_headers().await.unwrap(); + workload.get_headers().await.unwrap(); + assert_eq!( + mock.mint_count(), + 2, + "account-level and workload modes must key independently" + ); + + // Each then serves its own cached token. + account_level.get_headers().await.unwrap(); + workload.get_headers().await.unwrap(); + assert_eq!(mock.mint_count(), 2, "both must now be cache hits"); + } + + #[tokio::test] + async fn supplier_error_propagates_and_is_not_cached() { + // A supplier whose token fetch fails (e.g. the external IdP rejected the + // credentials). The error must surface from get_headers, and because + // nothing was cached, a subsequent call must invoke the supplier again + // rather than serving a stale/absent token. No network is used: the + // supplier fails before the exchange is ever attempted. + let calls = Arc::new(AtomicUsize::new(0)); + let calls_in_cb = Arc::clone(&calls); + let supplier: IdpTokenSupplier = Arc::new(move || { + let calls = Arc::clone(&calls_in_cb); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(crate::ZerobusError::InvalidUCTokenError( + "external IdP token fetch failed".to_string(), + )) + }) + }); + let provider = FederatedTokenProvider::new( + None, + supplier, + "cat.sch.tbl".to_string(), + "12345".to_string(), + "http://127.0.0.1:1".to_string(), + ); + + assert!( + provider.get_headers().await.is_err(), + "first call must error" + ); + assert!( + provider.get_headers().await.is_err(), + "second call must error too" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a failed mint must not be cached; the supplier is retried each call" + ); + } +} diff --git a/rust/sdk/src/lib.rs b/rust/sdk/src/lib.rs index 177c6e3c..cde78736 100644 --- a/rust/sdk/src/lib.rs +++ b/rust/sdk/src/lib.rs @@ -68,7 +68,9 @@ pub use dynamic_proto::{ pub use errors::{SchemaValidationCause, ZerobusError}; #[cfg(feature = "testing")] pub use headers_provider::NoAuthHeadersProvider; -pub use headers_provider::{HeadersProvider, OAuthHeadersProvider}; +pub use headers_provider::{ + FederatedTokenProvider, HeadersProvider, IdpTokenSupplier, OAuthHeadersProvider, +}; #[cfg(feature = "testing")] pub use multiplexed_stream::{MessageId, MultiplexedStream}; pub use offset_generator::{OffsetId, OffsetIdGenerator}; From 98f5215f0379a75a0138ede10e32c2678b3e102c Mon Sep 17 00:00:00 2001 From: Anil Menon Date: Sun, 16 Aug 2026 17:41:41 +0200 Subject: [PATCH 2/2] [Python] Add federated-token auth to the Python SDK Expose the Rust core's external-IdP (e.g. Entra ID) federation through the Python binding as an opt-in `auth=FederatedToken(...)` argument to create_stream, in both the sync and async SDKs. No existing signatures change in behavior. - auth.rs: add make_idp_token_supplier(), bridging a Python IdP-token callback to the Rust IdpTokenSupplier. Supports sync callbacks (return a str) and async callbacks (return an awaitable, driven via pyo3_async_runtimes::into_future). Also forward invalidate() through HeadersProviderWrapper to the Python provider's optional invalidate() hook, closing a prior gap. - sync_wrapper/async_wrapper: add create_stream_federated(), dispatching to the builder's .federated() / .federated_with_client_id(). - FederatedToken: a pure-Python dataclass (idp_token_supplier + optional databricks_client_id), exported from `zerobus`. - create_stream: accept auth=FederatedToken(...); client_id/client_secret become optional when auth or headers_provider is given (validated). Precedence: auth > headers_provider > OAuth. Existing paths unchanged. - Update type stubs for create_stream_federated (sync + async). Tests: new test_federated_auth.py covers export, dispatch routing (account -level vs workload), precedence, and arg validation. All Python tests pass; clippy, rustfmt, black, isort, and pycodestyle are clean. Signed-off-by: Anil Menon --- python/NEXT_CHANGELOG.md | 12 ++ python/README.md | 53 ++++++++ python/examples/sync_example_federated.py | 119 ++++++++++++++++++ python/rust/src/async_wrapper.rs | 38 +++++- python/rust/src/auth.rs | 101 ++++++++++++++- python/rust/src/sync_wrapper.rs | 43 ++++++- python/tests/test_federated_auth.py | 145 ++++++++++++++++++++++ python/zerobus/__init__.py | 2 + python/zerobus/_zerobus_core.pyi | 47 ++++++- python/zerobus/sdk/aio/zerobus_sdk.py | 37 ++++-- python/zerobus/sdk/shared/__init__.py | 3 + python/zerobus/sdk/shared/auth.py | 57 +++++++++ python/zerobus/sdk/sync/zerobus_sdk.py | 36 ++++-- 13 files changed, 668 insertions(+), 25 deletions(-) create mode 100644 python/examples/sync_example_federated.py create mode 100644 python/tests/test_federated_auth.py create mode 100644 python/zerobus/sdk/shared/auth.py diff --git a/python/NEXT_CHANGELOG.md b/python/NEXT_CHANGELOG.md index f1cb6f57..aa0411d3 100644 --- a/python/NEXT_CHANGELOG.md +++ b/python/NEXT_CHANGELOG.md @@ -6,6 +6,18 @@ ### New Features and Improvements +- Added `FederatedToken` for external-IdP (for example Entra ID) authentication. + Pass `auth=FederatedToken(idp_token_supplier=..., databricks_client_id=...)` + to `create_stream` and the SDK exchanges the external IdP token for a + Zerobus-scoped Databricks token (RFC 8693 token exchange), caching and + refreshing it. Supports account-level federation (omit `databricks_client_id`, + identity synced via Automatic Identity Management) and workload identity + federation (set `databricks_client_id` to the service principal, no secret). + The `idp_token_supplier` callback may be synchronous or asynchronous. Existing + `client_id`/`client_secret` and `headers_provider` calls are unchanged. +- Forwarded the `HeadersProvider.invalidate()` hook through the Python bridge, so + a custom provider can drop cached auth state when the server rejects a token. + ### Bug Fixes ### Documentation diff --git a/python/README.md b/python/README.md index 970217e7..c31c4b9e 100644 --- a/python/README.md +++ b/python/README.md @@ -287,6 +287,59 @@ asyncio.run(main()) See the [`examples/`](examples/) directory for complete runnable examples. +## Authentication + +`create_stream()` supports three authentication methods, in this precedence: +`auth` (federation) > `headers_provider` > `client_id`/`client_secret`. + +### OAuth client credentials (default) + +```python +stream = sdk.create_stream( + client_id="your-client-id", + client_secret="your-client-secret", + table_properties=table_properties, +) +``` + +### External-IdP federation (e.g. Entra ID) + +Use `auth=FederatedToken(...)` to authenticate with an external identity +provider instead of a Databricks OAuth secret. You provide a callback that +returns the current external IdP token; the SDK exchanges it for a +Zerobus-scoped Databricks token (RFC 8693 token exchange) and caches and +refreshes that token for you. The callback may be synchronous (sync SDK) or +asynchronous (async SDK). + +```python +from zerobus import FederatedToken + +def get_idp_token(): + # Return the current external IdP (e.g. Entra ID) access token. + ... + +# Account-level federation: no Databricks service principal. The identity is +# synced into Databricks via Automatic Identity Management (SCIM). +stream = sdk.create_stream( + table_properties=table_properties, + auth=FederatedToken(idp_token_supplier=get_idp_token), +) + +# Workload identity federation: a Databricks service principal with a client_id +# and no secret, with a federation policy attached. +stream = sdk.create_stream( + table_properties=table_properties, + auth=FederatedToken(idp_token_supplier=get_idp_token, databricks_client_id=""), +) +``` + +See [`examples/sync_example_federated.py`](examples/sync_example_federated.py) for a complete example. + +### Custom headers provider + +For advanced cases you can supply your own `HeadersProvider` via +`headers_provider=`; see the [`HeadersProvider`](#headersprovider) reference. + ## Configuration Configure stream behavior by passing a `StreamConfigurationOptions` object to `create_stream()`: diff --git a/python/examples/sync_example_federated.py b/python/examples/sync_example_federated.py new file mode 100644 index 00000000..76d4a333 --- /dev/null +++ b/python/examples/sync_example_federated.py @@ -0,0 +1,119 @@ +""" +Synchronous Ingestion Example - Federated Authentication (external IdP / Entra ID) + +This example demonstrates streaming to Zerobus using an external identity +provider (for example Microsoft Entra ID) instead of a Databricks OAuth +client_id/client_secret. You supply a callback that returns the current external +IdP token; the SDK exchanges it for a Zerobus-scoped Databricks token (RFC 8693 +token exchange), then caches and refreshes that token for you. + +Two federation modes, selected by `databricks_client_id`: + - Account-level federation (databricks_client_id omitted): no Databricks + service principal. The identity is synced into Databricks via Automatic + Identity Management (SCIM). + - Workload identity federation (databricks_client_id set): a Databricks service + principal with a client_id and no secret, with a federation policy attached. + +Record Type Mode: JSON (omitting a descriptor from TableProperties selects JSON). + +Note: the existing client_id/client_secret and headers_provider authentication +paths are unchanged; `auth=FederatedToken(...)` is a new, opt-in argument. +""" + +import json +import logging +import os + +import requests + +from zerobus import FederatedToken +from zerobus.sdk.shared import StreamConfigurationOptions, TableProperties +from zerobus.sdk.sync import ZerobusSdk + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +# Configuration - update these with your values. +SERVER_ENDPOINT = os.getenv( + "ZEROBUS_SERVER_ENDPOINT", + "https://your-shard-id.zerobus.region.cloud.databricks.com", +) +UNITY_CATALOG_ENDPOINT = os.getenv("DATABRICKS_WORKSPACE_URL", "https://your-workspace.cloud.databricks.com") +TABLE_NAME = os.getenv("ZEROBUS_TABLE_NAME", "catalog.schema.table") + +# External IdP (Entra ID) app registration used to mint the IdP token. +ENTRA_TENANT_ID = os.getenv("ENTRA_TENANT_ID", "your-entra-tenant-id") +ENTRA_CLIENT_ID = os.getenv("ENTRA_CLIENT_ID", "your-entra-client-id") +ENTRA_CLIENT_SECRET = os.getenv("ENTRA_CLIENT_SECRET", "your-entra-client-secret") +# Audience/scope your Databricks federation policy expects. +ENTRA_SCOPE = os.getenv("ENTRA_SCOPE", "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default") + +# Set to the Databricks service principal for workload identity federation; +# leave unset (None) for account-level federation. +DATABRICKS_SP_CLIENT_ID = os.getenv("DATABRICKS_SP_CLIENT_ID") or None + +NUM_RECORDS = 100 + + +def get_entra_token(): + """Return a fresh external IdP (Entra ID) access token. + + This is the `idp_token_supplier` callback. The SDK calls it only when it + needs to mint or refresh the exchanged Databricks token (a cache miss or a + proactive refresh), not on every record, so doing a network fetch here is + fine. A synchronous callback like this works with the sync SDK; the async + SDK also accepts an `async def` callback. + """ + resp = requests.post( + f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}/oauth2/v2.0/token", + data={ + "grant_type": "client_credentials", + "client_id": ENTRA_CLIENT_ID, + "client_secret": ENTRA_CLIENT_SECRET, + "scope": ENTRA_SCOPE, + }, + timeout=15, + ) + resp.raise_for_status() + return resp.json()["access_token"] + + +def create_sample_json_record(index): + """Create a sample record as a dict. JSON mode accepts a dict or a JSON string.""" + return { + "device_name": f"sensor-{index % 10}", + "temp": 20 + (index % 15), + "humidity": 50 + (index % 40), + } + + +def main(): + sdk = ZerobusSdk(SERVER_ENDPOINT, UNITY_CATALOG_ENDPOINT) + + # Opt into federated auth. Omit databricks_client_id for account-level + # federation; set it for workload identity federation. + auth = FederatedToken( + idp_token_supplier=get_entra_token, + databricks_client_id=DATABRICKS_SP_CLIENT_ID, + ) + mode = "workload identity" if DATABRICKS_SP_CLIENT_ID else "account-level" + logger.info("Creating stream with %s federation to %s", mode, TABLE_NAME) + + # No descriptor => JSON record format. + table_properties = TableProperties(TABLE_NAME) + options = StreamConfigurationOptions(max_inflight_records=100_000, recovery=True) + + stream = sdk.create_stream(table_properties=table_properties, options=options, auth=auth) + try: + # Queue records in a loop, then flush once. Never wait per record. + for i in range(NUM_RECORDS): + stream.ingest_record_offset(json.dumps(create_sample_json_record(i))) + stream.flush() + logger.info("Ingested and acknowledged %d records", NUM_RECORDS) + finally: + stream.close() + + +if __name__ == "__main__": + main() diff --git a/python/rust/src/async_wrapper.rs b/python/rust/src/async_wrapper.rs index 424569fd..af3fc64a 100644 --- a/python/rust/src/async_wrapper.rs +++ b/python/rust/src/async_wrapper.rs @@ -12,7 +12,7 @@ use databricks_zerobus_ingest_sdk::{ use crate::arrow; use crate::arrow::{ArrowStreamConfigurationOptions, AsyncZerobusArrowStream}; -use crate::auth::HeadersProviderWrapper; +use crate::auth::{make_idp_token_supplier, HeadersProviderWrapper}; use crate::common::{ apply_grpc_options, encoded_record_to_pybytes, extract_record_payload, extract_record_payloads, map_error, StreamConfigurationOptions, TableProperties, SDK_IDENTIFIER_PREFIX, @@ -362,6 +362,42 @@ impl ZerobusSdk { }) } + /// Create a stream with external-IdP federation (RFC 8693 token exchange) + /// (async). `idp_token_supplier` is a callback (sync or async) returning the + /// current external IdP token. `databricks_client_id` is `Some` for workload + /// identity federation and `None` for account-level federation. + #[pyo3(signature = (table_properties, idp_token_supplier, databricks_client_id = None, options = None))] + fn create_stream_federated<'py>( + &self, + py: Python<'py>, + table_properties: &TableProperties, + idp_token_supplier: Py, + databricks_client_id: Option, + options: Option, + ) -> PyResult> { + let sdk = self.inner.clone(); + let table_properties = table_properties.clone(); + let opts = options.unwrap_or_default(); + opts.validate()?; + let supplier = make_idp_token_supplier(idp_token_supplier); + + future_into_py(py, async move { + let sdk_guard = sdk.read().await; + let builder = match databricks_client_id { + Some(client_id) => sdk_guard + .stream_builder() + .federated_with_client_id(supplier, client_id), + None => sdk_guard.stream_builder().federated(supplier), + }; + let builder = apply_table_and_format(builder, &table_properties); + let builder = apply_grpc_options(builder, &opts)?; + let stream = builder.build().await.map_err(map_error)?; + Ok(ZerobusStream { + inner: Arc::new(RwLock::new(stream)), + }) + }) + } + /// Create a new Arrow Flight stream with OAuth authentication (async). #[pyo3(signature = (table_name, schema_ipc_bytes, client_id, client_secret, options = None))] fn create_arrow_stream<'py>( diff --git a/python/rust/src/auth.rs b/python/rust/src/auth.rs index e6bed4a4..c60df0c9 100644 --- a/python/rust/src/auth.rs +++ b/python/rust/src/auth.rs @@ -1,14 +1,26 @@ use async_trait::async_trait; -use pyo3::exceptions::PyNotImplementedError; +use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError}; use pyo3::prelude::*; +use pyo3_async_runtimes::TaskLocals; use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; use databricks_zerobus_ingest_sdk::{ - HeadersProvider as RustHeadersProvider, ZerobusError as RustError, ZerobusResult as RustResult, + HeadersProvider as RustHeadersProvider, IdpTokenSupplier, ZerobusError as RustError, + ZerobusResult as RustResult, }; use crate::common::intern_header_name; +/// Maps a Python error raised while bridging into the Rust SDK's error type, +/// prefixed with a short context so failures are attributable. +fn py_err_to_rust(context: &str, err: PyErr) -> RustError { + let msg = format!("{}: {}", context, err); + RustError::CreateStreamError(tonic::Status::new(tonic::Code::InvalidArgument, msg)) +} + /// Base class for custom authentication headers (subclassable from Python) /// /// The Rust SDK handles OAuth authentication internally by default. @@ -91,4 +103,89 @@ impl RustHeadersProvider for HeadersProviderWrapper { } Ok(map) } + + async fn invalidate(&self) { + // Forward to the Python provider's optional `invalidate()` hook so a + // custom provider can drop cached auth state when the server rejects a + // token. Providers that do not define one (the common case) are a + // no-op. Best-effort: any error is swallowed because `invalidate` is a + // cache-drop hint with no return channel. + let _ = Python::attach(|py| -> PyResult<()> { + let obj = self.py_obj.bind(py); + if let Ok(method) = obj.getattr("invalidate") { + method.call0()?; + } + Ok(()) + }); + } +} + +// ============================================================================= +// FEDERATED IDP TOKEN SUPPLIER BRIDGE +// ============================================================================= + +/// The outcome of invoking the Python IdP-token callback: either a token was +/// returned directly (sync callback) or an awaitable was returned that must be +/// driven to completion (async callback). +enum TokenOutcome { + Ready(String), + Awaitable(Pin>> + Send>>), +} + +/// Bridges a Python IdP-token callback to the Rust SDK's [`IdpTokenSupplier`]. +/// +/// The callback is invoked only when a fresh Databricks token must be minted +/// (a cache miss or refresh), and must return the current external IdP token as +/// a string. Both sync callbacks (return the string directly) and async +/// callbacks (return an awaitable) are supported. Async callbacks are driven +/// via the running event loop, so they require the async SDK. +pub fn make_idp_token_supplier(py_callable: Py) -> IdpTokenSupplier { + // Capture the running asyncio event loop's task-locals up front, on the + // Python thread constructing this supplier. The async SDK calls + // create_stream from inside a running loop, so this succeeds there; the + // sync SDK has no running loop, so it stays None and only sync callbacks + // are supported. These locals are required to drive an async (awaitable) + // callback later, because the supplier runs on a Rust worker thread where + // no event loop is running (calling into_future there fails with + // "no running event loop"). + let task_locals: Option = + Python::attach(|py| pyo3_async_runtimes::tokio::get_current_locals(py).ok()); + + Arc::new(move || { + // Invoke the callback under the GIL. If it returned an awaitable, + // convert it to a Rust future here (GIL held) using the captured + // event-loop locals, then await it below without holding the GIL. + let outcome = Python::attach(|py| -> PyResult { + let result = py_callable.bind(py).call0()?; + if result.hasattr("__await__")? { + match task_locals.as_ref() { + Some(locals) => { + let fut = pyo3_async_runtimes::into_future_with_locals(locals, result)?; + Ok(TokenOutcome::Awaitable(Box::pin(fut))) + } + None => Err(PyRuntimeError::new_err( + "async idp_token_supplier requires the async SDK (a running event loop); \ + use a synchronous callback with the sync SDK", + )), + } + } else { + Ok(TokenOutcome::Ready(result.extract::()?)) + } + }); + + Box::pin(async move { + match outcome { + Ok(TokenOutcome::Ready(token)) => Ok(token), + Ok(TokenOutcome::Awaitable(fut)) => { + let awaited = fut + .await + .map_err(|e| py_err_to_rust("federated IdP token callback failed", e))?; + Python::attach(|py| awaited.bind(py).extract::()).map_err(|e| { + py_err_to_rust("federated IdP token callback returned a non-string", e) + }) + } + Err(e) => Err(py_err_to_rust("federated IdP token callback failed", e)), + } + }) + }) } diff --git a/python/rust/src/sync_wrapper.rs b/python/rust/src/sync_wrapper.rs index e5b74a04..3d9923bb 100644 --- a/python/rust/src/sync_wrapper.rs +++ b/python/rust/src/sync_wrapper.rs @@ -11,7 +11,7 @@ use databricks_zerobus_ingest_sdk::{ use crate::arrow; use crate::arrow::{ArrowStreamConfigurationOptions, ZerobusArrowStream}; -use crate::auth::HeadersProviderWrapper; +use crate::auth::{make_idp_token_supplier, HeadersProviderWrapper}; use crate::common::{ apply_grpc_options, encoded_record_to_pybytes, extract_record_payload, extract_record_payloads, map_error, StreamConfigurationOptions, TableProperties, SDK_IDENTIFIER_PREFIX, @@ -403,6 +403,47 @@ impl ZerobusSdk { }) } + /// Create a new stream with external-IdP federation (RFC 8693 token + /// exchange). `idp_token_supplier` is a callback returning the current + /// external IdP token. `databricks_client_id` is `Some` for workload + /// identity federation and `None` for account-level federation. + #[pyo3(signature = (table_properties, idp_token_supplier, databricks_client_id = None, options = None))] + fn create_stream_federated( + &self, + py: Python, + table_properties: TableProperties, + idp_token_supplier: Py, + databricks_client_id: Option, + options: Option, + ) -> PyResult { + let opts = options.unwrap_or_default(); + opts.validate()?; + let supplier = make_idp_token_supplier(idp_token_supplier); + let sdk = self.inner.clone(); + let runtime = self.runtime.clone(); + let runtime_for_stream = self.runtime.clone(); + + let stream = py.detach(|| { + runtime.block_on(async move { + let sdk_guard = sdk.read().await; + let builder = match databricks_client_id { + Some(client_id) => sdk_guard + .stream_builder() + .federated_with_client_id(supplier, client_id), + None => sdk_guard.stream_builder().federated(supplier), + }; + let builder = apply_table_and_format(builder, &table_properties); + let builder = apply_grpc_options(builder, &opts)?; + builder.build().await.map_err(map_error) + }) + })?; + + Ok(ZerobusStream { + inner: Arc::new(RwLock::new(stream)), + runtime: runtime_for_stream, + }) + } + /// Create a new Arrow Flight stream with OAuth authentication. /// /// **Beta**: Arrow Flight ingestion is in Beta. The API is stabilising diff --git a/python/tests/test_federated_auth.py b/python/tests/test_federated_auth.py new file mode 100644 index 00000000..98512b3f --- /dev/null +++ b/python/tests/test_federated_auth.py @@ -0,0 +1,145 @@ +"""Tests for the external-IdP federation auth surface (FederatedToken). + +These cover the pure-Python dispatch layer of ``create_stream`` with a fake +native ``_inner`` so no network or gRPC server is needed. The Rust core's +exchange/caching behavior is covered by the Rust unit tests; the Python->Rust +callback bridge is exercised end-to-end separately against a mock token +endpoint. +""" + +import pytest + +import zerobus +from zerobus import FederatedToken, TableProperties, ZerobusSdk +from zerobus.sdk.aio import ZerobusSdk as AsyncZerobusSdk + + +class _FakeInner: + """Records which native create_stream_* method the wrapper dispatched to.""" + + def __init__(self): + self.calls = [] + + def create_stream_federated(self, table_properties, idp_token_supplier, databricks_client_id, options): + self.calls.append(("federated", table_properties, idp_token_supplier, databricks_client_id, options)) + return object() + + def create_stream(self, client_id, client_secret, table_properties, options): + self.calls.append(("oauth", client_id, client_secret, table_properties, options)) + return object() + + def create_stream_with_headers_provider(self, table_properties, headers_provider, options): + self.calls.append(("headers", table_properties, headers_provider, options)) + return object() + + +class _AsyncFakeInner: + def __init__(self): + self.calls = [] + + async def create_stream_federated(self, table_properties, idp_token_supplier, databricks_client_id, options): + self.calls.append(("federated", table_properties, idp_token_supplier, databricks_client_id, options)) + return object() + + async def create_stream(self, client_id, client_secret, table_properties, options): + self.calls.append(("oauth", client_id, client_secret, table_properties, options)) + return object() + + async def create_stream_with_headers_provider(self, table_properties, headers_provider, options): + self.calls.append(("headers", table_properties, headers_provider, options)) + return object() + + +def _sync_sdk_with_fake(): + sdk = ZerobusSdk(host="https://example", unity_catalog_url="https://example") + fake = _FakeInner() + sdk._inner = fake + return sdk, fake + + +def _props(): + return TableProperties("cat.sch.tbl") + + +def test_federated_token_exported_and_constructs(): + assert "FederatedToken" in zerobus.__all__ + account = FederatedToken(idp_token_supplier=lambda: "tok") + assert account.databricks_client_id is None + workload = FederatedToken(idp_token_supplier=lambda: "tok", databricks_client_id="sp-uuid") + assert workload.databricks_client_id == "sp-uuid" + + +def test_native_methods_present(): + import zerobus._zerobus_core as _core + + assert hasattr(_core.sync.ZerobusSdk, "create_stream_federated") + assert hasattr(_core.aio.ZerobusSdk, "create_stream_federated") + + +def test_create_stream_routes_auth_to_federated_account_level(): + sdk, fake = _sync_sdk_with_fake() + + def supplier(): + return "tok" + + sdk.create_stream(table_properties=_props(), auth=FederatedToken(idp_token_supplier=supplier)) + + assert len(fake.calls) == 1 + kind, _tp, passed_supplier, client_id, _opts = fake.calls[0] + assert kind == "federated" + assert passed_supplier is supplier + assert client_id is None + + +def test_create_stream_routes_auth_to_federated_workload(): + sdk, fake = _sync_sdk_with_fake() + sdk.create_stream( + table_properties=_props(), + auth=FederatedToken(idp_token_supplier=lambda: "tok", databricks_client_id="sp-uuid"), + ) + kind, _tp, _sup, client_id, _opts = fake.calls[0] + assert kind == "federated" + assert client_id == "sp-uuid" + + +def test_create_stream_oauth_path_unchanged(): + sdk, fake = _sync_sdk_with_fake() + sdk.create_stream("cid", "secret", _props()) + assert fake.calls[0][0] == "oauth" + assert fake.calls[0][1] == "cid" + + +def test_auth_takes_precedence_over_headers_provider(): + sdk, fake = _sync_sdk_with_fake() + sdk.create_stream( + table_properties=_props(), + auth=FederatedToken(idp_token_supplier=lambda: "tok"), + headers_provider=object(), + ) + assert fake.calls[0][0] == "federated" + + +def test_create_stream_requires_auth_or_credentials(): + sdk, _fake = _sync_sdk_with_fake() + with pytest.raises(ValueError): + sdk.create_stream(table_properties=_props()) + + +def test_create_stream_requires_table_properties(): + sdk, _fake = _sync_sdk_with_fake() + with pytest.raises(ValueError): + sdk.create_stream(auth=FederatedToken(idp_token_supplier=lambda: "tok")) + + +@pytest.mark.asyncio +async def test_async_create_stream_routes_to_federated(): + sdk = AsyncZerobusSdk(host="https://example", unity_catalog_url="https://example") + fake = _AsyncFakeInner() + sdk._inner = fake + + await sdk.create_stream( + table_properties=_props(), + auth=FederatedToken(idp_token_supplier=lambda: "tok", databricks_client_id="sp"), + ) + assert fake.calls[0][0] == "federated" + assert fake.calls[0][3] == "sp" diff --git a/python/zerobus/__init__.py b/python/zerobus/__init__.py index 24eff233..85e7f9c6 100644 --- a/python/zerobus/__init__.py +++ b/python/zerobus/__init__.py @@ -56,6 +56,7 @@ # Import from Rust core import zerobus._zerobus_core as _core from zerobus.sdk.shared.arrow import ArrowStreamConfigurationOptions, IPCCompression +from zerobus.sdk.shared.auth import FederatedToken from zerobus.sdk.sync import ZerobusArrowStream, ZerobusSdk, ZerobusStream __version__ = "1.6.1" @@ -86,6 +87,7 @@ "AckCallback", # Authentication "HeadersProvider", + "FederatedToken", # Exceptions "ZerobusException", "NonRetriableException", diff --git a/python/zerobus/_zerobus_core.pyi b/python/zerobus/_zerobus_core.pyi index 7eed9e65..0f90d468 100644 --- a/python/zerobus/_zerobus_core.pyi +++ b/python/zerobus/_zerobus_core.pyi @@ -1,6 +1,6 @@ """Type stubs for _zerobus_core Rust module.""" -from typing import Any, List, Optional, Tuple, Union +from typing import Any, Awaitable, Callable, List, Optional, Tuple, Union from typing_extensions import Self @@ -375,6 +375,28 @@ class sync: """ ... + def create_stream_federated( + self, + table_properties: TableProperties, + idp_token_supplier: Callable[[], Union[str, Awaitable[str]]], + databricks_client_id: Optional[str] = None, + options: Optional[StreamConfigurationOptions] = None, + ) -> "ZerobusStream": + """ + Create a new stream with external-IdP federation (RFC 8693 token exchange). + + Args: + table_properties: Table properties + idp_token_supplier: Callback returning the current external IdP token + databricks_client_id: Service principal client_id for workload + identity federation, or None for account-level federation + options: Optional configuration options + + Returns: + A new ZerobusStream + """ + ... + def recreate_stream(self, old_stream: "ZerobusStream") -> "ZerobusStream": """ Recreate a closed stream with the same configuration. @@ -528,6 +550,29 @@ class aio: """ ... + async def create_stream_federated( + self, + table_properties: TableProperties, + idp_token_supplier: Callable[[], Union[str, Awaitable[str]]], + databricks_client_id: Optional[str] = None, + options: Optional[StreamConfigurationOptions] = None, + ) -> "ZerobusStream": + """ + Create a new stream with external-IdP federation (RFC 8693 token exchange). + + Args: + table_properties: Table properties + idp_token_supplier: Callback (sync or async) returning the current + external IdP token + databricks_client_id: Service principal client_id for workload + identity federation, or None for account-level federation + options: Optional configuration options + + Returns: + A new ZerobusStream + """ + ... + async def recreate_stream(self, old_stream: "ZerobusStream") -> "ZerobusStream": """ Recreate a closed stream with the same configuration. diff --git a/python/zerobus/sdk/aio/zerobus_sdk.py b/python/zerobus/sdk/aio/zerobus_sdk.py index 5ee3ac9b..6c07da5e 100644 --- a/python/zerobus/sdk/aio/zerobus_sdk.py +++ b/python/zerobus/sdk/aio/zerobus_sdk.py @@ -322,29 +322,46 @@ async def recreate_arrow_stream(self, old_stream: ZerobusArrowStream) -> Zerobus async def create_stream( self, - client_id: str, - client_secret: str, - table_properties, + client_id: str = None, + client_secret: str = None, + table_properties=None, options=None, headers_provider=None, + auth=None, ): """ - Create a stream with OAuth authentication or custom headers provider. + Create a stream with OAuth, external-IdP federation, or a custom headers provider. + + Exactly one authentication method is used, in this precedence: + ``auth`` (federation) > ``headers_provider`` > ``client_id``/``client_secret``. Args: - client_id: OAuth client ID - client_secret: OAuth client secret - table_properties: Table configuration - options: Optional stream configuration - headers_provider: Optional custom headers provider (if set, overrides OAuth) + client_id: OAuth client ID (client-credentials auth). + client_secret: OAuth client secret (client-credentials auth). + table_properties: Table configuration (required). + options: Optional stream configuration. + headers_provider: Optional custom headers provider (if set, overrides OAuth). + auth: Optional :class:`~zerobus.sdk.shared.auth.FederatedToken` for + external-IdP (e.g. Entra ID) federation. When set, ``client_id`` + and ``client_secret`` are not required. Its ``idp_token_supplier`` + may be a sync or async callable. """ - if headers_provider is not None: + if table_properties is None: + raise ValueError("table_properties is required") + if auth is not None: + # External-IdP federation (RFC 8693 token exchange). + rust_stream = await self._inner.create_stream_federated( + table_properties, auth.idp_token_supplier, auth.databricks_client_id, options + ) + elif headers_provider is not None: # Use custom headers provider (ignores client_id/client_secret) rust_stream = await self._inner.create_stream_with_headers_provider( table_properties, headers_provider, options ) else: # Use OAuth authentication + if client_id is None or client_secret is None: + raise ValueError("client_id and client_secret are required unless auth= or headers_provider= is given") rust_stream = await self._inner.create_stream(client_id, client_secret, table_properties, options) return ZerobusStream(rust_stream) diff --git a/python/zerobus/sdk/shared/__init__.py b/python/zerobus/sdk/shared/__init__.py index 7e21e709..20b846c3 100644 --- a/python/zerobus/sdk/shared/__init__.py +++ b/python/zerobus/sdk/shared/__init__.py @@ -15,11 +15,14 @@ ) # Import Python wrappers with documentation +from zerobus.sdk.shared.auth import FederatedToken, IdpTokenSupplier from zerobus.sdk.shared.config import AckCallback, StreamConfigurationOptions __all__ = [ "AckCallback", + "FederatedToken", "HeadersProvider", + "IdpTokenSupplier", "NonRetriableException", "RecordType", "StreamConfigurationOptions", diff --git a/python/zerobus/sdk/shared/auth.py b/python/zerobus/sdk/shared/auth.py new file mode 100644 index 00000000..aecce152 --- /dev/null +++ b/python/zerobus/sdk/shared/auth.py @@ -0,0 +1,57 @@ +""" +Federated external-IdP authentication for Zerobus streams. + +This module defines :class:`FederatedToken`, the opt-in configuration for +authenticating a stream with an external identity provider (for example +Entra ID) instead of a Databricks OAuth client_id/client_secret. +""" + +from dataclasses import dataclass +from typing import Awaitable, Callable, Optional, Union + +# A zero-arg callback returning the current external IdP token (e.g. an Entra ID +# JWT), either synchronously (``str``) or asynchronously (an awaitable of +# ``str``). +IdpTokenSupplier = Callable[[], Union[str, Awaitable[str]]] + + +@dataclass +class FederatedToken: + """Authenticate a Zerobus stream by federating an external IdP token. + + The SDK exchanges the external IdP token returned by ``idp_token_supplier`` + for a Zerobus-scoped Databricks token via RFC 8693 token exchange. The + exchange happens client-side, in the SDK; the Zerobus service is unchanged. + + Two federation modes are selected by ``databricks_client_id``: + + * **Account-level federation** (``databricks_client_id=None``): no + Databricks-managed service principal. The token subject is resolved to an + identity synced into Databricks via Automatic Identity Management (SCIM). + * **Workload identity federation** (``databricks_client_id`` set): a + Databricks service principal with a client_id and no secret, with a + federation policy attached. The exchange names the service principal via + its client_id. + + Pass an instance as the ``auth`` argument to ``create_stream``. + + Args: + idp_token_supplier: A zero-arg callable returning the current external + IdP token as a string. May be synchronous (returns ``str``) or + asynchronous (returns an awaitable of ``str``); async suppliers + require the async SDK. It is called only when a fresh Databricks + token must be minted (a cache miss or refresh), never on every + request, so a callable that fetches a token is fine here. + databricks_client_id: The Databricks service principal client_id for + workload identity federation, or ``None`` for account-level + federation. + """ + + idp_token_supplier: IdpTokenSupplier + databricks_client_id: Optional[str] = None + + +__all__ = [ + "FederatedToken", + "IdpTokenSupplier", +] diff --git a/python/zerobus/sdk/sync/zerobus_sdk.py b/python/zerobus/sdk/sync/zerobus_sdk.py index 1343a1b5..0550991e 100644 --- a/python/zerobus/sdk/sync/zerobus_sdk.py +++ b/python/zerobus/sdk/sync/zerobus_sdk.py @@ -304,27 +304,43 @@ def recreate_arrow_stream(self, old_stream: ZerobusArrowStream) -> ZerobusArrowS def create_stream( self, - client_id: str, - client_secret: str, - table_properties, + client_id: str = None, + client_secret: str = None, + table_properties=None, options=None, headers_provider=None, + auth=None, ): """ - Create a stream with OAuth authentication or custom headers provider. + Create a stream with OAuth, external-IdP federation, or a custom headers provider. + + Exactly one authentication method is used, in this precedence: + ``auth`` (federation) > ``headers_provider`` > ``client_id``/``client_secret``. Args: - client_id: OAuth client ID - client_secret: OAuth client secret - table_properties: Table configuration - options: Optional stream configuration - headers_provider: Optional custom headers provider (if set, overrides OAuth) + client_id: OAuth client ID (client-credentials auth). + client_secret: OAuth client secret (client-credentials auth). + table_properties: Table configuration (required). + options: Optional stream configuration. + headers_provider: Optional custom headers provider (if set, overrides OAuth). + auth: Optional :class:`~zerobus.sdk.shared.auth.FederatedToken` for + external-IdP (e.g. Entra ID) federation. When set, ``client_id`` + and ``client_secret`` are not required. """ - if headers_provider is not None: + if table_properties is None: + raise ValueError("table_properties is required") + if auth is not None: + # External-IdP federation (RFC 8693 token exchange). + rust_stream = self._inner.create_stream_federated( + table_properties, auth.idp_token_supplier, auth.databricks_client_id, options + ) + elif headers_provider is not None: # Use custom headers provider (ignores client_id/client_secret) rust_stream = self._inner.create_stream_with_headers_provider(table_properties, headers_provider, options) else: # Use OAuth authentication + if client_id is None or client_secret is None: + raise ValueError("client_id and client_secret are required unless auth= or headers_provider= is given") rust_stream = self._inner.create_stream(client_id, client_secret, table_properties, options) return ZerobusStream(rust_stream)