Databricks OIDC token federation for AWS IAM

[allenc97]

Extends the Databricks provider hook with OIDC token federation for Airflow running with AWS credentials, parallel to the existing Kubernetes path (#61458). The hook mints an AWS-signed OIDC JWT via AWS STS GetWebIdentityToken and exchanges it for a Databricks service-principal token using the same RFC 8693 exchange as federated_k8s — only the JWT source differs. Like the Kubernetes path it stores no secrets in the connection; identity comes from the caller's assumed IAM role. Set login=federated_aws and provide client_id and aws_conn_id in the connection extra. See the Databricks AWS IAM workload identity federation guide.

• closes: Add support for Databricks oauth federation auth #54291

Integration testing

Verified end-to-end against a real workspace with this connection (no secrets stored) and DAG:

{
  "conn_type": "databricks",
  "host": "<workspace>.cloud.databricks.com",
  "login": "federated_aws",
  "extra": {
    "client_id": "<databricks-sp-client-id>",
    "aws_conn_id": "<aws_conn_id>"
  }
}

from __future__ import annotations

from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.sdk import dag
from airflow.sdk import task


@dag(
    dag_id="test_databricks_federated_aws",
    schedule=None,
    catchup=False,
    description="Integration test: authenticate to Databricks via AWS IAM federation and confirm the SP identity.",
)
def test_databricks_federated_aws():
    @task
    def test_connection():
        hook = DatabricksHook(databricks_conn_id="databricks_federated_aws")
        print(hook.test_connection())
        print(hook._do_api_call(("GET", "/2.0/preview/scim/v2/Me")))

    test_connection()


test_databricks_federated_aws()

Requires outbound web identity federation enabled on the AWS account, the role granted sts:GetWebIdentityToken, and a Databricks service-principal federation policy whose subject is the role ARN and audiences is ["databricks"].

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Generated-by: Claude Code (Opus 4.8) following the guidelines

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

---

Important

🛠️ Maintainer triage note for @allenc97 · by @potiuk · 2026-07-15 16:00 UTC

Some review feedback from @eladkal is waiting on you (see our Pull Request quality criteria):

• ❌ Unresolved review comments: 1 thread(s). See docs.

The ball is in your court — you've been assigned to this PR. Reply or push a fix in each thread, then mark them resolved.

_{Automated triage — may be imperfect; a maintainer takes the next look.}

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[allenc97]

I refactored this PR against #69272 - given that the connection already supports other cloud based identity schemes I extended it to support AWS STS as well. However, if it's preferred based on the above PR's supplied OIDC callable changes that other Databricks supported OIDC sources remain user custom I am also happy to close this PR.

[eladkal]

I refactored this PR against #69272 - given that the connection already supports other cloud based identity schemes I extended it to support AWS STS as well. However, if it's preferred based on the above PR's supplied OIDC callable changes that other Databricks supported OIDC sources remain user custom I am also happy to close this PR.

That is a question for Databricks experts cc @moomindani

[moomindani]

Reviewed this PR, including checking out the branch and verifying it locally. Sharing my take as an individual working with these hooks — this is my personal technical opinion, not an official position.

On the built-in vs. user-custom question: I think built-in is the right call, and this PR is worth moving forward.

Some history that I believe settles most of the question:

• When the Kubernetes federation path (#61458) was being developed, its author explicitly announced the roadmap in #54291: Kubernetes first, then AWS STS outbound federation next. This PR is that announced next step, not scope creep — and merging it would also answer the still-open question in #54291 about what remains after #61458.
• The hook already ships three Azure-specific built-in auth paths (SPN, managed identity, default Azure credential) plus the Kubernetes one. Leaving AWS as the only major platform whose users must hand-roll a token callable would be asymmetric, and Airflow-on-AWS (MWAA, EC2, EKS) is one of the largest deployment bases for Databricks users.
• The optional cross-provider amazon extra follows an established pattern — google, mysql, postgres, and teradata providers already declare the same optional dependency. The import is lazy with a clean AirflowOptionalProviderFeatureException, so environments without the amazon provider are unaffected.
• I don't read #69272 as a mandate that cloud-specific sources stay user-custom. Its callable is the escape hatch for issuers Airflow does not model; a tested, documented first-class path for the most common platform complements it rather than competing with it. Asking every AWS deployment to write and maintain its own boto3 callable would be strictly worse UX for the same result.

A few checks I did beyond what CI covers (the unit tests mock StsHook, so CI never exercises the real API contract):

• The GetWebIdentityToken call matches the current botocore service model exactly (Audience list and SigningAlgorithm required, WebIdentityToken in the output), and the chosen defaults — audience "databricks", RS256, 300 seconds — match the reference implementation in the Databricks AWS IAM federation guide parameter for parameter, including the documented IAM policy conditions (sts:IdentityTokenAudience = databricks, duration <= 300s).
• apache-airflow-providers-amazon>=9.22.0 is precisely the first release that ships boto3>=1.41.0, so both the pin and the fallback error message for older SDKs are accurate.
• On the Databricks side, the token-exchange path (/oidc/v1/token, RFC 8693, client_id switching between service-principal and account-wide federation policies) is the same code already validated end-to-end against a real workspace in #69272 — this PR only swaps in a different subject-token source, and the author's documented integration test covers the full AWS flow.

Two small non-blocking nits as inline comments. LGTM once the docstring is touched up.

---

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

[boring-cyborg]

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.