[Rust][Python] Add first-class federated-token auth (external IdP / Entra ID) - #741
Open
anilmenon14 wants to merge 2 commits into
Open
[Rust][Python] Add first-class federated-token auth (external IdP / Entra ID)#741anilmenon14 wants to merge 2 commits into
anilmenon14 wants to merge 2 commits into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
HeadersProviderhook; this makes it first-class instead of aworkaround. The Zerobus service is unchanged.
What this changes
FederatedTokenProvider(implements the existingHeadersProvidertrait, includinginvalidate()) and anIdpTokenSuppliercallback 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(...)andfederated_with_client_id(...).auth=FederatedToken(idp_token_supplier=..., databricks_client_id=...)on
create_stream(sync and async), with the Python callback bridged acrossFFI (sync and async callbacks both supported). The
HeadersProvider.invalidate()hook is now forwarded through the Python bridge.
Resolves #740
The two supported modes
databricks_client_idomitted): no Databricksservice principal. The exchanged token's subject resolves to an identity
synced into Databricks via Automatic Identity Management (SCIM). The exchange
request omits
client_id.databricks_client_idset): a Databricksservice 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) andheaders_providerpaths areunchanged. The new behavior is reached only when the caller passes the new,
opt-in
auth=FederatedToken(...)argument (or thefederated*builder methodsin Rust).
Testing evidence
client_id, plusend-to-end provider tests against a mock token endpoint (caching,
invalidate()re-mint, mode independence, supplier-error propagation). All lib tests pass;
clippyandrustfmtclean.black,isort,pycodestyleclean.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
would cover non-SDK and REST callers uniformly); this PR is the client-side
first step.
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
NEXT_CHANGELOG.mdupdated (Rust core and Python).examples/sample added for the new API.