Skip to content

Add supplied-token OIDC federation to the Databricks provider - #69272

Merged
kaxil merged 2 commits into
apache:mainfrom
astronomer:databricks-supplied-token-federation
Jul 8, 2026
Merged

Add supplied-token OIDC federation to the Databricks provider#69272
kaxil merged 2 commits into
apache:mainfrom
astronomer:databricks-supplied-token-federation

Conversation

@kaxil

@kaxil kaxil commented Jul 2, 2026

Copy link
Copy Markdown
Member

The Databricks provider can already federate an OIDC token into a Databricks OAuth token, but only from the pod's Kubernetes service account (federated_k8s), read from disk. This adds a second subject-token source: a federated_token_provider connection extra, a dotted path to a Callable[[], str] that returns the OIDC JWT in-process, which the hook exchanges for a Databricks OAuth token through the same RFC 8693 token-exchange endpoint the federated_k8s path uses. No long-lived secret is stored, the subject token is never written to disk, and the source works in any environment (not only Kubernetes) with any issuer a Databricks federation policy trusts.

Why a token provider callable

The subject token is a short-lived, per-workload OIDC JWT that has to be minted fresh at connect time, so a static value in the connection extra would not work (it would be a stored long-lived secret that never refreshes). A callable lets a control plane vend a fresh token per task, in-process. This mirrors the OpenAI provider's token_provider custom source (#69069) and resolves the dotted path with import_string, the same pattern other providers use for callables in connection config.

Because the callable is imported and executed in the process running the hook, point it only at trusted code. The connection extra is an operator/admin surface.

Design notes

  • client_id is optional for this path: include it for a service principal federation policy, or omit it for an account-wide policy (an account-wide policy is what lets the issuer stamp a per-workload subject claim). The Kubernetes path is unchanged: client_id stays required (service account tokens cannot carry custom claims) and is still validated before the token is read.
  • The async path offloads the synchronous, usually network-bound provider call to a worker thread with run_in_executor, so it cannot stall the triggerer event loop used by deferrable operators.
  • No connector version bump: the exchange reuses the existing raw requests call to /oidc/v1/token, identical to the shipped Kubernetes path.
  • Like the other extra-based auth methods, this path is only used when no higher-precedence credential (a PAT, token, Azure credentials, or service_principal_oauth) is set on the connection.

Usage

{
  "conn_type": "databricks",
  "host": "my-workspace.cloud.databricks.com",
  "extra": {
    "federated_token_provider": "my_package.identity.get_oidc_token"
  }
}

Add client_id for a service principal federation policy; omit it for an account-wide policy.

Add a `federated_token_provider` connection extra: a dotted-path callable that
supplies a short-lived OIDC JWT in-process, which the hook exchanges for a
Databricks OAuth token via the same RFC 8693 endpoint the `federated_k8s` path
uses. Works with any federation-trusted issuer and in any environment, not only
Kubernetes. `client_id` is optional (account-wide policy) or supplied (service
principal policy). The Kubernetes federation path is behaviourally unchanged.
Comment thread providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py Outdated
Address review: validate the federated_token_provider return on its stripped value and
return the stripped token, so a whitespace-only token is rejected and a valid token with
surrounding whitespace (e.g. a trailing newline) is normalised before the exchange -- matching
the Kubernetes path, which already strips. Add whitespace-only and newline-padded test cases.
@kaxil
kaxil requested review from phanikumv and vincbeck July 7, 2026 17:58

@vincbeck vincbeck left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not too familiar wth Databricks ... But code looks okay to me

@eladkal

eladkal commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

We have @moomindani from Databricks to verify if needed :)

@moomindani

Copy link
Copy Markdown
Contributor

Verified this PR against a real Databricks workspace, including the full happy path:

  • Created a service principal federation policy (inline JWKS of a locally generated RSA key) on a test account, pointed federated_token_provider at a callable that signs a matching RS256 JWT in-process, and both _get_token and _a_get_token exchanged it at the real /oidc/v1/token and returned working OAuth tokens (verified with an authenticated API call as that service principal). A second call reused the cached token without re-invoking the provider.
  • Differential checks against the real endpoint: omitting client_id with an SP-scoped policy fails with invalid_grant (consistent with the docs — account-wide policies are the no-client_id case); an untrusted JWT fails with invalid_grant: TOKEN_INVALID; a whitespace-only provider return fails locally with ValueError before any network call.

LGTM.


Drafted-by: Claude Code (Fable 5); reviewed by @moomindani before posting

@kaxil
kaxil merged commit 7dc3c2b into apache:main Jul 8, 2026
82 checks passed
@kaxil
kaxil deleted the databricks-supplied-token-federation branch July 8, 2026 07:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants