[Rust][Python] Add first-class federated-token auth (external IdP / Entra ID)

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

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="<sp-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()`:

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()

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<PyAny>,
+        databricks_client_id: Option<String>,
+        options: Option<StreamConfigurationOptions>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        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>(

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<Box<dyn Future<Output = PyResult<Py<PyAny>>> + 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<PyAny>) -> 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<TaskLocals> =
+        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<TokenOutcome> {
+            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::<String>()?))
+            }
+        });
+
+        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::<String>()).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)),
+            }
+        })
+    })
 }