Databricks OIDC token federation for AWS IAM

dev/breeze/tests/test_selective_checks.py

@@ -1083,7 +1083,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
             ("providers/amazon/src/airflow/providers/amazon/provider.yaml",),
             {
                 "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
-                "common.compat common.messaging common.sql exasol ftp google http imap microsoft.azure "
+                "common.compat common.messaging common.sql databricks exasol ftp google http imap microsoft.azure "
                 "mongo mysql openlineage postgres salesforce ssh teradata",
                 "all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
                 "all-python-versions-list-as-string": DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
@@ -1109,7 +1109,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
                         {
                             "description": "amazon...google",
                             "test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
-                            "common.compat,common.messaging,common.sql,exasol,ftp,http,imap,"
+                            "common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
                             "microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
                             "Providers[google]",
                         }
@@ -1155,7 +1155,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
             ("providers/amazon/src/airflow/providers/amazon/file.py",),
             {
                 "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
-                "common.compat common.messaging common.sql exasol ftp google http imap microsoft.azure "
+                "common.compat common.messaging common.sql databricks exasol ftp google http imap microsoft.azure "
                 "mongo mysql openlineage postgres salesforce ssh teradata",
                 "all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
                 "all-python-versions-list-as-string": DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
@@ -1178,7 +1178,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
                         {
                             "description": "amazon...google",
                             "test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
-                            "common.compat,common.messaging,common.sql,exasol,ftp,http,imap,"
+                            "common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
                             "microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
                             "Providers[google]",
                         }

providers/databricks/docs/connections/databricks.rst

@@ -52,6 +52,8 @@ There are several ways to connect to Databricks using Airflow.
    i.e. a caller-provided callable returns a short-lived OIDC JWT that is exchanged for a Databricks OAuth token. Unlike the Kubernetes method,
    the subject token is obtained in-process (never read from disk) and can come from any federation-trusted OIDC issuer, so it is not tied to
    Kubernetes and supports both account-wide and service-principal federation policies. Like the Kubernetes method, no long-lived secret is stored in the connection.
+8. Using AWS IAM `OIDC token federation <https://docs.databricks.com/aws/en/dev-tools/auth/provider-aws-iam>`__ (applicable when Airflow runs with AWS credentials)
+   i.e. mint an AWS-signed OIDC JWT via AWS STS ``GetWebIdentityToken`` and exchange it for a Databricks OAuth token. Stores no secrets; requires the provider's ``amazon`` extra.
 
 Default Connection IDs
 ----------------------
@@ -172,6 +174,17 @@ Extra (optional)
               resp.raise_for_status()
               return resp.json()["token"]
 
+    The following parameters enable *AWS IAM OIDC token federation* — mint an AWS-signed OIDC JWT via AWS STS
+    ``GetWebIdentityToken`` and exchange it for a Databricks token. Requires the ``amazon`` extra
+    (``pip install apache-airflow-providers-databricks[amazon]``). See the Databricks `AWS IAM workload identity
+    federation guide <https://docs.databricks.com/aws/en/dev-tools/auth/provider-aws-iam>`_.
+
+    * ``federated_aws``: set ``login`` to ``"federated_aws"`` or add it as a boolean flag (``{"federated_aws": true}``).
+    * ``aws_conn_id``: (optional) AWS connection used for the STS call (default: ``aws_default``).
+    * ``aws_jwt_audience``: (optional) audience of the minted JWT; must match the federation policy (default: ``databricks``).
+    * ``aws_web_identity_token_duration``: (optional) requested JWT lifetime in seconds (default: ``300``, AWS max ``3600``).
+    * ``client_id``: (optional) service principal client UUID for a service principal policy; omit for an account-wide policy.
+
     The following parameters are necessary if using authentication with Kubernetes OIDC token federation:
 
     * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as a boolean flag in extra parameters (``{"federated_k8s": true}``). When enabled, the hook will fetch a JWT token from Kubernetes and exchange it for a Databricks OAuth token using the `OIDC token exchange API <https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-exchange.html>`_. This authentication method only works when Airflow is running inside a Kubernetes cluster (e.g., AWS EKS, Azure AKS, Google GKE).

providers/databricks/docs/index.rst

@@ -126,12 +126,13 @@ You can install such cross-provider dependencies when installing from PyPI. For
 
 .. code-block:: bash
 
-    pip install apache-airflow-providers-databricks[google]
+    pip install apache-airflow-providers-databricks[amazon]
 
 
 ==============================================================================================================  ===============
 Dependent package                                                                                               Extra
 ==============================================================================================================  ===============
+`apache-airflow-providers-amazon <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_            ``amazon``
 `apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_            ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_  ``openlineage``
 ==============================================================================================================  ===============
@@ -151,6 +152,7 @@ Install them when installing from PyPI. For example:
 Extra               Dependencies
 ==================  ================================================================================================================================================================
 ``avro``            ``fastavro>=1.9.0; python_version<"3.14"``, ``fastavro>=1.10.0; python_version>="3.12" and python_version<"3.14"``, ``fastavro>=1.12.1; python_version>="3.14"``
+``amazon``          ``apache-airflow-providers-amazon>=9.22.0``
 ``azure-identity``  ``azure-identity>=1.3.1``
 ``fab``             ``apache-airflow-providers-fab>=2.2.0``
 ``google``          ``apache-airflow-providers-google>=10.24.0``

providers/databricks/pyproject.toml

@@ -83,6 +83,9 @@ dependencies = [
     'fastavro>=1.10.0; python_version>="3.12" and python_version<"3.14"',
     'fastavro>=1.12.1; python_version>="3.14"',
 ]
+"amazon" = [
+    "apache-airflow-providers-amazon>=9.22.0"
+]
 "azure-identity" = [
     "azure-identity>=1.3.1",
 ]
@@ -110,6 +113,7 @@ dev = [
     "apache-airflow",
     "apache-airflow-task-sdk",
     "apache-airflow-devel-common",
+    "apache-airflow-providers-amazon",
     "apache-airflow-providers-common-compat",
     "apache-airflow-providers-common-sql",
     "apache-airflow-providers-google",

providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py

@@ -84,6 +84,11 @@
 DEFAULT_K8S_NAMESPACE_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
 K8S_CA_CERT_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
 
+# AWS IAM OIDC token federation (STS Outbound Identity Federation)
+# https://docs.databricks.com/aws/en/dev-tools/auth/provider-aws-iam
+DEFAULT_AWS_JWT_AUDIENCE = "databricks"
+DEFAULT_AWS_WEB_IDENTITY_TOKEN_DURATION = 300
+
 # RFC 8693 token exchange data template
 TOKEN_EXCHANGE_DATA = {
     "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
@@ -137,6 +142,10 @@ class BaseDatabricksHook(BaseHook):
         "service_principal_oauth",
         "federated_k8s",
         "federated_token_provider",
+        "federated_aws",
+        "aws_conn_id",
+        "aws_jwt_audience",
+        "aws_web_identity_token_duration",
         "k8s_token_path",
         "k8s_namespace_path",
         "k8s_projected_volume_token_path",
@@ -917,17 +926,24 @@ def _get_federation_subject_token(self) -> tuple[str, str | None]:
         """
         Resolve the OIDC JWT to exchange for a Databricks token (RFC 8693 ``subject_token``).
 
-        Two subject-token sources are supported:
+        Three subject-token sources are supported:
 
         * ``federated_token_provider`` -- a dotted path to a ``Callable[[], str]`` that returns
           the JWT. The token is obtained in-process and never written to disk, so a control
           plane can vend a short-lived, per-workload identity token for the exchange. ``client_id``
           is optional here: supply it in the extra for a service principal federation policy, or
           omit it for an account-wide federation policy.
+        * AWS IAM (``federated_aws``) -- mint an AWS-signed OIDC JWT via AWS STS
+          ``GetWebIdentityToken``, configured entirely from the connection extra
+          (``aws_conn_id``/``aws_jwt_audience``/``aws_web_identity_token_duration``). ``client_id`` is
+          optional, as with a supplied provider.
         * Kubernetes service account (default) -- read from the pod. ``client_id`` is required
           because Kubernetes service account tokens cannot carry custom claims, so only
           service-principal-level federation is possible; it is validated before the token is read.
 
+        If more than one is configured, precedence is: supplied provider, then AWS, then Kubernetes
+        (matching the dispatch order in :meth:`_get_token`).
+
         :return: a ``(subject_token, client_id)`` tuple; ``client_id`` is ``None`` when the exchange
             should omit it (account-wide federation policy).
         """
@@ -936,6 +952,8 @@ def _get_federation_subject_token(self) -> tuple[str, str | None]:
             return self._resolve_supplied_subject_token(provider), self.databricks_conn.extra_dejson.get(
                 "client_id"
             )
+        if self._is_aws_federation():
+            return self._get_aws_subject_token(), self.databricks_conn.extra_dejson.get("client_id")
         client_id = self._get_required_client_id()
         return self._get_k8s_jwt_token(), client_id
 
@@ -948,9 +966,57 @@ async def _a_get_federation_subject_token(self) -> tuple[str, str | None]:
             loop = asyncio.get_running_loop()
             subject_token = await loop.run_in_executor(None, self._resolve_supplied_subject_token, provider)
             return subject_token, self.databricks_conn.extra_dejson.get("client_id")
+        if self._is_aws_federation():
+            loop = asyncio.get_running_loop()
+            subject_token = await loop.run_in_executor(None, self._get_aws_subject_token)
+            return subject_token, self.databricks_conn.extra_dejson.get("client_id")
         client_id = self._get_required_client_id()
         return await self._a_get_k8s_jwt_token(), client_id
 
+    def _is_aws_federation(self) -> bool:
+        """Return whether the connection is configured for AWS IAM OIDC token federation."""
+        return self.databricks_conn.login == "federated_aws" or self.databricks_conn.extra_dejson.get(
+            "federated_aws", False
+        )
+
+    def _get_aws_subject_token(self) -> str:
+        """
+        Mint an AWS-signed OIDC JWT for the exchange via AWS STS ``GetWebIdentityToken``.
+
+        Reads ``aws_conn_id``/``aws_jwt_audience``/``aws_web_identity_token_duration`` from the connection
+        extra. The minted JWT's ``sub`` claim is the caller's IAM role ARN, matching the Databricks
+        federation policy. See https://docs.databricks.com/aws/en/dev-tools/auth/provider-aws-iam.
+        """
+        try:
+            from airflow.providers.amazon.aws.hooks.sts import StsHook
+        except ImportError as e:
+            raise AirflowOptionalProviderFeatureException(
+                "The 'apache-airflow-providers-amazon' package (>=9.22.0) is required for AWS OIDC "
+                "token federation. Install it with: pip install 'apache-airflow-providers-amazon>=9.22.0'"
+            ) from e
+
+        extra = self.databricks_conn.extra_dejson
+        aws_conn_id = extra.get("aws_conn_id", "aws_default")
+        audience = extra.get("aws_jwt_audience", DEFAULT_AWS_JWT_AUDIENCE)
+        duration_seconds = int(
+            extra.get("aws_web_identity_token_duration", DEFAULT_AWS_WEB_IDENTITY_TOKEN_DURATION)
+        )
+
+        sts_client = StsHook(aws_conn_id=aws_conn_id).get_conn()
+        if not hasattr(sts_client, "get_web_identity_token"):
+            raise AirflowOptionalProviderFeatureException(
+                "The installed AWS SDK does not support 'sts:GetWebIdentityToken'. AWS IAM outbound "
+                "identity federation requires boto3>=1.41.0 / botocore>=1.41.0 "
+                "(apache-airflow-providers-amazon>=9.22.0)."
+            )
+
+        response = sts_client.get_web_identity_token(
+            Audience=[audience],
+            SigningAlgorithm="RS256",
+            DurationSeconds=duration_seconds,
+        )
+        return response["WebIdentityToken"]
+
     def _resolve_supplied_subject_token(self, provider: str) -> str:
         """
         Import and invoke the ``federated_token_provider`` callable, returning its OIDC JWT.
@@ -983,9 +1049,9 @@ def _get_federated_databricks_token(self, resource: str) -> str:
         """
         Get a Databricks OAuth token by exchanging a federated OIDC JWT.
 
-        Uses RFC 8693 token exchange to convert an OIDC subject token -- supplied either by a
-        ``federated_token_provider`` callable or by the pod's Kubernetes service account (see
-        :meth:`_get_federation_subject_token`) -- into a Databricks OAuth token.
+        Uses RFC 8693 token exchange to convert an OIDC subject token -- supplied by a
+        ``federated_token_provider`` callable, AWS STS (``federated_aws``), or the pod's Kubernetes
+        service account (see :meth:`_get_federation_subject_token`) -- into a Databricks OAuth token.
 
         :param resource: Databricks OIDC token exchange URL
         :return: Databricks OAuth access token
@@ -1167,6 +1233,9 @@ def _get_token(self, raise_error: bool = False) -> str | None:
         if self.databricks_conn.extra_dejson.get("federated_token_provider"):
             self.log.debug("Using OIDC token federation with a supplied token provider.")
             return self._get_federated_databricks_token(self._get_oidc_token_service_url())
+        if self._is_aws_federation():
+            self.log.debug("Using AWS IAM OIDC token federation.")
+            return self._get_federated_databricks_token(self._get_oidc_token_service_url())
         if self.databricks_conn.login == "federated_k8s" or self.databricks_conn.extra_dejson.get(
             "federated_k8s", False
         ):
@@ -1207,6 +1276,9 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None:
         if self.databricks_conn.extra_dejson.get("federated_token_provider"):
             self.log.debug("Using OIDC token federation with a supplied token provider.")
             return await self._a_get_federated_databricks_token(self._get_oidc_token_service_url())
+        if self._is_aws_federation():
+            self.log.debug("Using AWS IAM OIDC token federation.")
+            return await self._a_get_federated_databricks_token(self._get_oidc_token_service_url())
         if self.databricks_conn.login == "federated_k8s" or self.databricks_conn.extra_dejson.get(
             "federated_k8s", False
         ):