[Rust][Python] Add first-class federated-token auth (external IdP / Entra ID) by anilmenon14 · Pull Request #760 · databricks/zerobus-sdk

anilmenon14 · 2026-08-19T16:51:05Z

Motivation

Enterprise customers who cannot use Databricks-managed OAuth secrets currently
cannot use Zerobus. This PR adds an opt-in authentication mode that federates an
external identity provider (for example Entra ID) token into a Zerobus-scoped
Databricks token, client-side, so those customers can stream without a
Databricks secret. The platform token exchange already works on gRPC via the
undocumented HeadersProvider hook; this makes it first-class instead of a
workaround. The Zerobus service is unchanged.

What this changes

Rust core: new FederatedTokenProvider (implements the existing
HeadersProvider trait, including invalidate()) and an IdpTokenSupplier
callback type. The client-credentials and token-exchange grants now share one
request-shaping path in default_token_factory.rs, keeping them at parity.
Opt-in StreamBuilder::federated(...) and federated_with_client_id(...).
Python binding: auth=FederatedToken(idp_token_supplier=..., databricks_client_id=...)
on create_stream (sync and async), with the Python callback bridged across
FFI (sync and async callbacks both supported). The HeadersProvider.invalidate()
hook is now forwarded through the Python bridge.
No existing signatures change behavior.

Resolves #740

The two supported modes

Account-level federation: (databricks_client_id omitted): no Databricks
service principal. The exchanged token's subject resolves to an identity
synced into Databricks via Automatic Identity Management (SCIM). The exchange
request omits client_id.
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 request names the service principal via client_id.

Backward compatibility

The client_id/client_secret (OAuth) and headers_provider paths are
unchanged. The new behavior is reached only when the caller passes the new,
opt-in auth=FederatedToken(...) argument (or the federated* builder methods
in Rust).

Testing evidence

Rust: unit tests for request-shaping parity with/without client_id, plus
end-to-end provider tests against a mock token endpoint (caching, invalidate()
re-mint, mode independence, supplier-error propagation). All lib tests pass;
clippy and rustfmt clean.
Python: dispatch/routing and argument-validation tests. black, isort,
pycodestyle clean.
Live (both modes, against a real workspace, streaming to a UC table): a
successful stream on gRPC for account-level and workload identity federation;
caching confirmed (the IdP callback fires once across multiple streams from one
SDK); a >1 hour soak showing the exchanged token auto-refreshes near the
~55-minute mark (token lifetime ~60 min minus the 300s cache buffer) with no
interruption; and an async callback validated via the async SDK.

REST insert path (resolved, not a limitation)

The same federated exchange token works on the REST insert endpoint
(/zerobus/v1/tables/<table>/insert), verified live with an HTTP 200 insert.

Known limitations / follow-ups

Other language bindings (TypeScript, Java, Go, C++) to follow.
A server-side / control-plane exchange remains a reasonable later phase (it
would cover non-SDK and REST callers uniformly); this PR is the client-side
first step.
Suggested maintainer test follow-up: the sync callback is unit-tested and
live-verified; the async callback bridge and the Python FFI error mapping are
verified live but not yet covered by automated unit tests, because both need a
live token+gRPC endpoint or a fuller mock harness. Adding CI coverage for these
two using the repo's test fixtures would be worthwhile.

Housekeeping

DCO sign-off on all commits; GPG-signed.
NEXT_CHANGELOG.md updated (Rust core and Python).
README and an examples/ sample added for the new API.

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 <anil.menon@databricks.com>

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 <anil.menon@databricks.com>

elenagaljak-db · 2026-08-24T15:45:21Z

+    /// 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<String>,
+    ) -> Self {
+        self.auth = Some(AuthConfig::Federated {
+            idp_token_supplier,
+            client_id: Some(client_id.into()),
+        });
+        self
+    }
+


Could these 2 methods be one? They call the same underlying method and they differ by one parameter. We can make the client id an optional parameter:

pub fn federated( mut self, idp_token_supplier: IdpTokenSupplier, client_id: Option<impl Into<String>>, ) -> Self { self.auth = Some(AuthConfig::Federated { idp_token_supplier, client_id: client_id.map(Into::into), }); self }

Also I would shorten the docs comments to match style of the file. Something like

/// Authenticate with external-IdP federation (RFC 8693). /// `client_id`: `None` for account-level federation; Databricks SP id for workload identity.

There is already more info in other files, no need for that in the stream_builder files

elenagaljak-db · 2026-08-24T16:01:40Z

 const fn missing_auth_error() -> &'static str {
    #[cfg(feature = "testing")]
    {
        "authentication is required: call .oauth(), .headers_provider(), or .no_auth()"
    }
    #[cfg(not(feature = "testing"))]
    {
        "authentication is required: call .oauth() or .headers_provider()"
    }
 }


we should add .federated() here too

elenagaljak-db · 2026-08-24T16:04:16Z

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


This can also be shortened

elenagaljak-db · 2026-08-24T16:04:52Z

+    /// 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)]


Ditto. Also, I don't think we need this story1/story2 in production code

elenagaljak-db · 2026-08-24T16:16:47Z

+                self.cache_client_id(),
+                "",
+                &self.table_name,


Is the shared cache supposed to be one federated identity per Zerobus sdk? Per my understanding of the code, account-level keys as ("", "", table), so two .federated() streams on the same SDK and table share one slot even if their IdpTokenSuppliers are different accounts. The second caller would get a cache hit and send the first caller's Databricks token. Is that intended? If multi-user on one SDK is in scope, should the key include an identity (JWT sub, or something the caller passes)?

elenagaljak-db · 2026-08-24T16:19:25Z

            ));
        }

        let expires_in = Self::parse_expires_in(&body);


Does UC expires_in always match the remaining lifetime of the subject JWT?
Federation docs say the exchanged token inherits the JWT exp, and the AWS WIF example uses a 300s token. We cache only from expires_in. If that field can be 3600 while the JWT has minutes left, we'd keep serving a dead token until the refresh buffer. should we cap TTL with min(expires_in, jwt.exp - now)?

elenagaljak-db · 2026-08-24T16:21:07Z

+            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.


Why these changes?

elenagaljak-db · 2026-08-24T16:23:53Z

+        client_id: str = None,
+        client_secret: str = None,
+        table_properties=None,
        options=None,
        headers_provider=None,
+        auth=None,


With client_id still first, create_stream(table_properties, auth=FederatedToken(...)) binds the table to client_id and then raises "table_properties is required". The example uses keywords so it works but it won't work otherwise

elenagaljak-db · 2026-08-24T16:24:06Z

+        client_id: str = None,
+        client_secret: str = None,
+        table_properties=None,
        options=None,
        headers_provider=None,
+        auth=None,


elenagaljak-db · 2026-08-24T16:24:52Z

+fn py_err_to_rust(context: &str, err: PyErr) -> RustError {
+    let msg = format!("{}: {}", context, err);
+    RustError::CreateStreamError(tonic::Status::new(tonic::Code::InvalidArgument, msg))
+}
+


This maps every IdP callback failure to CreateStreamError + InvalidArgument, which is non-retryable. OAuth mint failures go through TokenFetchError (retryable). We should have a way to distinguish

elenagaljak-db · 2026-08-24T16:26:29Z

+    """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.
+    """


Could this be Args + one sentence, like neighboring types? We already have this in README and the example

anilmenon14 added 2 commits August 16, 2026 17:41

anilmenon14 requested review from bogdanb-db, elenagaljak-db and teodordelibasic-db August 19, 2026 16:51

anilmenon14 mentioned this pull request Aug 19, 2026

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

Closed

elenagaljak-db reviewed Aug 24, 2026

View reviewed changes

Navigation Menu

Uh oh!

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

[Rust][Python] Add first-class federated-token auth (external IdP / Entra ID)#760
anilmenon14 wants to merge 2 commits into
mainfrom
feature/federated-token-auth

anilmenon14 commented Aug 19, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

elenagaljak-db Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

anilmenon14 commented Aug 19, 2026

Motivation

What this changes

The two supported modes

Backward compatibility

Testing evidence

REST insert path (resolved, not a limitation)

Known limitations / follow-ups

Housekeeping

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants