Skip to content

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

Open
anilmenon14 wants to merge 2 commits into
databricks:mainfrom
anilmenon14:feature/federated-token-auth
Open

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

Conversation

@anilmenon14

Copy link
Copy Markdown

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Rust][Python] Add first-class external-IdP (Entra ID) token federation to the SDK

1 participant