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) 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};