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

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

Agree that would make the API surface for the method simpler by having one instead of variations for each. Merged into a single federated(supplier, client_id: Option<impl Into<String>>) and trimmed the doc to the two-line form you suggested. Updated the two Python native wrappers (the Some/None match collapses to one call), plus the README examples and the Rust changelog.
A quick note on ergonomics that with Option<impl Into<String>>, a bare None can't infer its type, so account-level reads .federated(supplier, None::<String>); workload stays clean as .federated(supplier, Some("sp-id")).

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

Agree. Added .federated() to the missing-auth error so it now lists alongside .oauth(), .headers_provider(), and .no_auth(). I covered both the default and the testing builds.

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

I trimmed it down to three lines to match fetch_token_inner, since the parameter detail already lives on exchange_form_params right below it.

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

Agree, that was a miss since these shouldn't have been there. I shortened the exchange_form_params doc, and pulled every Story 1 / Story 2 reference out of the code: the doc, the inline comment, and the two form-params tests. They just say "account-level" and "workload identity federation" now. I kept the Story framing to the internal design docs and PR only.

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

This one needs to be solved indeed since we could have multiple streams potentially instantiated , using different IdP callback functions (i.e., wrapping different identities from the IdP) , from the same ZerobusSdk instance. I will have a fix proposed , test it and add as new commit. Will respond back here in this thread after that is done.

Have implemented this by adding a cache_key that partitions the 'account-level' cache. In the Python SDK each FederatedToken auto-generates a stable per-instance key, so two different FederatedToken objects isolate (no more cross-identity token reuse), while reusing the same FederatedToken keeps the cache shared as before. The story around 'Workload identity' is unchanged, since it already keys by the SP client_id.

IMO stamping a per-instance key is safer than trying to derive the key from the identity itself. It keys by "same FederatedToken instance", which is a proxy for identity rather than the verified identity. I looked at keying by the token's actual sub, but that needs either calling the supplier on every request (defeats the cache) or a mint per provider (defeats cross-stream sharing), so it isn't workable without giving up the caching. The proxy converts the dangerous failure (silently writing as the wrong identity) into a harmless one (an occasional extra mint if someone builds a fresh FederatedToken per stream for the same identity). The Rust builder exposes the same thing as an optional cache_key on .federated() for direct Rust callers.

Added unit tests on both sides (Rust: distinct cache_keys isolate, same shares; Python: distinct FederatedTokens carry distinct keys, reuse carries the same), and validated live that reusing one FederatedToken still hits the cache.

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

This was a gap that had to be sorted. Great catch!
I implemented the cap on TTL now. fetch_exchanged_token now does a best-effort decode of the subject JWT's exp and caps the cached TTL at min(expires_in, exp - now), so we never serve a token past the point its subject expired. The decode is unverified on purpose, since it only bounds our own cache and never authorizes anything (the server already validated the token during the exchange). If the subject isn't a decodable JWT with a future exp, or UC returns no expires_in, it falls back to the old behavior, so it can never be worse than before.

One thing I liked about doing it this way is that the cap just shrinks the value before it reaches the cache, so all of #701 features I rebased on keep working unchanged. I added 5 unit tests for the cap, plus a base64 dep (already in the tree via reqwest). The client-credentials path is untouched since it has no subject token.

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?

True, that was unnecessary. I reverted it and the existing Args are back to their original wording, and I dropped the precedence paragraph and the :class: role. The only changes left now are the summary line (it mentions federation, since there's a third auth method) and the appended auth parameter. Same treatment on the async version.

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

Agree that this would lead to confusing behavior/ error.

auth is now keyword-only (a * in the signature). That's safe since nothing released passes it positionally.
More importantly, I added a guard for the exact case you hit. If auth= is combined with a positional that lands in client_id/client_secret, it now raises a clear error ("client_id/client_secret cannot be combined with auth=; pass table_properties as a keyword, e.g. create_stream(table_properties=..., auth=...)") instead of the misleading "table_properties is required".
I deliberately kept client_id first and didn't reorder the parameters. The SDK is already released, so reordering positional params would break existing OAuth callers. That means I can't make create_stream(table_properties, auth=...) work positionally, but the guard at least makes the right usage obvious. I also added regression tests proving the existing positional OAuth calls (including the 4-positional form) still bind correctly.

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,


Same fix on the async SDK as previous comment.

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

Good call, and I used the mechanism from your previous PR #759 rather than hand-rolling anything. Instead of classifying in the bridge, I make the callback failure produce the right ZerobusError variant so map_error / is_retryable() handles it:

A transient failure (the callback raised, or its awaitable rejected) is now TokenFetchError, which is retryable, exactly like an OAuth mint failure.
Caller misuse (a non-string return, or an async callback on the sync SDK) is InvalidArgument, which is non-retryable.
Each failure is classified at its source now, rather than everything getting lumped into one mapping. I added two unit tests that exercise this without a live server (they fail during minting, before the connect), a raising callback surfaces the base ZerobusException, and a non-string return surfaces NonRetriableException. I left the custom HeadersProvider error mapping alone, since that's a separate path from the IdP callback.

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

Makes sense. I trimmed it to a one-line summary plus Args, matching the neighboring types. The two federation modes are already covered in the README and the example, so I dropped the inline bullet list from the docstring.

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. Cap an exchanged token's cached TTL at the subject JWT's exp so it is never served past the point its subject expired. - 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 the workload client_id, or by an optional account-level cache_key, plus the table, so distinct identities and the two modes cache independently. - stream_builder: add an opt-in .federated(supplier, client_id, cache_key) builder method. Default auth paths are unchanged. Tests: request-shaping parity with/without client_id, subject-JWT TTL capping, account-level cache-key isolation, 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). A transient callback failure (it raised) is classified retryable (TokenFetchError), while caller misuse (a non-string return, or an async callback on the sync SDK) is not. 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(supplier, client_id, cache_key) method. - FederatedToken: a pure-Python dataclass (idp_token_supplier + optional databricks_client_id), exported from `zerobus`. Each instance carries an auto-generated cache key that partitions the shared token cache for account-level federation, so two different identities used from one ZerobusSdk do not collide; reusing one instance keeps it shared. - create_stream: accept auth=FederatedToken(...); client_id/client_secret become optional when auth or headers_provider is given (validated). auth is keyword-only, and combining it with client_id/client_secret raises a clear error. Precedence: auth > headers_provider > OAuth. Existing paths unchanged. - Update type stubs for create_stream_federated (sync + async). Tests: test_federated_auth.py covers export, dispatch routing (account -level vs workload), precedence, keyword-only auth, arg validation, and per-instance cache-key isolation; test_exceptions.py covers the retryable-vs-non-retryable callback mapping. All Python tests pass; clippy, rustfmt, black, isort, and pycodestyle are clean. Signed-off-by: Anil Menon <anil.menon@databricks.com>

anilmenon14 · 2026-08-27T12:09:11Z

@elenagaljak-db , thanks for the detailed review and surfacing some issues. I addressed the issues from the past review and now ready for another round of review with the refinement.

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

anilmenon14 force-pushed the feature/federated-token-auth branch from 98f5215 to add5f35 Compare August 27, 2026 04:44

anilmenon14 force-pushed the feature/federated-token-auth branch from add5f35 to c6273ca Compare August 27, 2026 10:17

anilmenon14 force-pushed the feature/federated-token-auth branch from c6273ca to 7fdc53a Compare August 27, 2026 10:39

anilmenon14 requested a review from elenagaljak-db August 27, 2026 12:07

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!

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!

anilmenon14 Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

anilmenon14 commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anilmenon14 Aug 27, 2026 •

edited

Loading