Add supplied-token OIDC federation to the Databricks provider

providers/databricks/docs/connections/databricks.rst

@@ -48,6 +48,10 @@ There are several ways to connect to Databricks using Airflow.
    i.e. automatically fetch JWT tokens from Kubernetes Service Account via projected volume path or TokenRequest API and exchange them for Databricks OAuth tokens.
    This is the recommended method when Airflow runs in Kubernetes. This method requires no secrets to be stored in the connection and eliminates the need
    for token management (no rotation, expiration handling, or credential storage).
+7. Using `OIDC token federation <https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation>`_ with a supplied token provider,
+   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.
 
 Default Connection IDs
 ----------------------
@@ -131,6 +135,43 @@ Extra (optional)
       a user inside workspace)
     * ``azure_managed_identity_client_id``: optional client ID of the user-assigned managed identity. This parameter is only required if you're using a user-assigned managed identity. If not specified, the hook will attempt to authenticate using a system-assigned managed identity.
 
+    The following parameter enables *OIDC token federation with a supplied token provider* (an alternative to
+    the Kubernetes method below that works in any environment, not only Kubernetes):
+
+    * ``federated_token_provider``: dotted path to a ``Callable[[], str]`` that returns an OIDC JWT (the RFC 8693
+      ``subject_token``). The hook imports and calls it in-process to obtain the token, then exchanges 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>`_;
+      the subject token is never written to disk. It may come from any OIDC issuer trusted by a Databricks
+      `federation policy <https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-policy>`_. ``client_id`` is
+      optional here: supply it for a service principal federation policy, or omit it for an account-wide federation
+      policy. When both ``federated_token_provider`` and ``federated_k8s`` are set, the supplied provider takes precedence.
+      Like the other extra-based methods, it is only used when no higher-precedence credential (a PAT in the ``Password``
+      field, ``token``, Azure credentials, or ``service_principal_oauth``) is set on the connection.
+
+      Because the dotted path is imported and executed in the process running the hook, point it only at trusted code.
+      The connection ``extra`` is an operator/admin surface, consistent with how other providers resolve callables from configuration.
+
+      .. code-block:: json
+
+          {
+            "federated_token_provider": "my_package.identity.get_oidc_token"
+          }
+
+      The callable takes no arguments and returns the OIDC JWT as a string. Obtain the token however your
+      environment provides it (for example, request it from your identity provider or a control-plane token
+      endpoint); the returned value is the ``subject_token`` that Databricks exchanges for an OAuth token:
+
+      .. code-block:: python
+
+          # my_package/identity.py
+          import requests
+
+
+          def get_oidc_token() -> str:
+              resp = requests.get("https://id.example.com/oidc/token", timeout=10)
+              resp.raise_for_status()
+              return resp.json()["token"]
+
     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/src/airflow/providers/databricks/hooks/databricks_base.py

@@ -25,6 +25,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import copy
 import platform
 import ssl
@@ -50,6 +51,7 @@
 )
 
 from airflow import __version__
+from airflow.providers.common.compat.module_loading import import_string
 from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException
 from airflow.providers.databricks.exceptions import DatabricksApiError
 from airflow.providers_manager import ProvidersManager
@@ -60,6 +62,8 @@
     from airflow.hooks.base import BaseHook as BaseHook  # type: ignore
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from airflow.models import Connection
 
 # https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token
@@ -123,6 +127,7 @@ class BaseDatabricksHook(BaseHook):
         "proxies",
         "service_principal_oauth",
         "federated_k8s",
+        "federated_token_provider",
         "k8s_token_path",
         "k8s_namespace_path",
         "k8s_projected_volume_token_path",
@@ -893,12 +898,79 @@ def _get_required_client_id(self) -> str:
             )
         return client_id
 
+    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:
+
+        * ``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.
+        * 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.
+
+        :return: a ``(subject_token, client_id)`` tuple; ``client_id`` is ``None`` when the exchange
+            should omit it (account-wide federation policy).
+        """
+        provider = self.databricks_conn.extra_dejson.get("federated_token_provider")
+        if provider:
+            return self._resolve_supplied_subject_token(provider), self.databricks_conn.extra_dejson.get(
+                "client_id"
+            )
+        client_id = self._get_required_client_id()
+        return self._get_k8s_jwt_token(), client_id
+
+    async def _a_get_federation_subject_token(self) -> tuple[str, str | None]:
+        """Async version of :meth:`_get_federation_subject_token`."""
+        provider = self.databricks_conn.extra_dejson.get("federated_token_provider")
+        if provider:
+            # The provider is a synchronous callable that typically makes a blocking network call to
+            # mint the token. Offload it to a worker thread so it can't stall the triggerer event loop.
+            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")
+        client_id = self._get_required_client_id()
+        return await self._a_get_k8s_jwt_token(), client_id
+
+    def _resolve_supplied_subject_token(self, provider: str) -> str:
+        """
+        Import and invoke the ``federated_token_provider`` callable, returning its OIDC JWT.
+
+        The dotted path is resolved and executed in-process; its return value is the RFC 8693
+        ``subject_token`` and is never written to disk. Surrounding whitespace is stripped (matching
+        the Kubernetes path), and a value that is not a non-empty string raises, so a misconfigured
+        provider fails with a clear error rather than posting a blank or newline-padded
+        ``subject_token`` to the exchange.
+        """
+        token_provider: Callable[[], str] = import_string(provider)
+        token = token_provider()
+        if not isinstance(token, str) or not token.strip():
+            raise ValueError(f"federated_token_provider {provider!r} must return a non-empty string token.")
+        return token.strip()
+
+    def _build_federation_exchange_data(self, subject_token: str, client_id: str | None) -> dict[str, str]:
+        """
+        Build the RFC 8693 token-exchange form data.
+
+        ``client_id`` is included when set -- required for Kubernetes/service principal federation,
+        optional for a supplied provider -- and omitted for an account-wide federation policy.
+        """
+        data = {**TOKEN_EXCHANGE_DATA, "subject_token": subject_token}
+        if client_id:
+            data["client_id"] = client_id
+        return data
+
     def _get_federated_databricks_token(self, resource: str) -> str:
         """
-        Get Databricks OAuth token by exchanging Kubernetes JWT token.
+        Get a Databricks OAuth token by exchanging a federated OIDC JWT.
 
-        Uses RFC 8693 token exchange to convert a Kubernetes service account JWT
-        into a Databricks OAuth token. Requires service principal-level federation.
+        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.
 
         :param resource: Databricks OIDC token exchange URL
         :return: Databricks OAuth access token
@@ -909,15 +981,12 @@ def _get_federated_databricks_token(self, resource: str) -> str:
 
         self.log.info("Existing federated token is expired or missing. Fetching new token...")
 
-        client_id = self._get_required_client_id()
-
-        # Get JWT from Kubernetes
-        jwt_token = self._get_k8s_jwt_token()
-        self.log.debug("JWT Token obtained from Kubernetes: %s", jwt_token)
+        subject_token, client_id = self._get_federation_subject_token()
 
-        # Prepare token exchange request following RFC 8693
+        # Prepare token exchange request following RFC 8693. The subject token is never logged --
+        # it is a short-lived credential.
         token_exchange_url = resource
-        data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token, "client_id": client_id}
+        data = self._build_federation_exchange_data(subject_token, client_id)
 
         try:
             for attempt in self._get_retry_object():
@@ -941,10 +1010,10 @@ def _get_federated_databricks_token(self, resource: str) -> str:
                     break
         except RetryError:
             raise AirflowException(
-                f"Failed to exchange Kubernetes JWT for Databricks token after {self.retry_limit} retries. Giving up."
+                f"Failed to exchange the federated OIDC token for a Databricks token after {self.retry_limit} retries. Giving up."
             )
         except requests_exceptions.HTTPError as e:
-            msg = f"Failed to exchange Kubernetes JWT for Databricks token. Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
+            msg = f"Failed to exchange the federated OIDC token for a Databricks token. Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
             raise AirflowException(msg)
 
         return jsn["access_token"]
@@ -957,15 +1026,12 @@ async def _a_get_federated_databricks_token(self, resource: str) -> str:
 
         self.log.info("Existing federated token is expired or missing. Fetching new token...")
 
-        client_id = self._get_required_client_id()
-
-        # Get JWT from Kubernetes
-        jwt_token = await self._a_get_k8s_jwt_token()
-        self.log.debug("JWT Token obtained from Kubernetes: %s", jwt_token)
+        subject_token, client_id = await self._a_get_federation_subject_token()
 
-        # Prepare token exchange request following RFC 8693
+        # Prepare token exchange request following RFC 8693. The subject token is never logged --
+        # it is a short-lived credential.
         token_exchange_url = resource
-        data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token, "client_id": client_id}
+        data = self._build_federation_exchange_data(subject_token, client_id)
 
         try:
             async for attempt in self._a_get_retry_object():
@@ -989,11 +1055,11 @@ async def _a_get_federated_databricks_token(self, resource: str) -> str:
                     break
         except RetryError:
             raise AirflowException(
-                f"Failed to exchange Kubernetes JWT for Databricks token after {self.retry_limit} retries. Giving up."
+                f"Failed to exchange the federated OIDC token for a Databricks token after {self.retry_limit} retries. Giving up."
             )
         except aiohttp.ClientResponseError as err:
             raise AirflowException(
-                f"Failed to exchange Kubernetes JWT for Databricks token. Response: {err.message}, Status Code: {err.status}"
+                f"Failed to exchange the federated OIDC token for a Databricks token. Response: {err.message}, Status Code: {err.status}"
             )
 
         return jsn["access_token"]
@@ -1083,6 +1149,9 @@ def _get_token(self, raise_error: bool = False) -> str | None:
                 raise AirflowException("Service Principal credentials aren't provided")
             self.log.debug("Using Service Principal Token.")
             return self._get_sp_token(self._get_oidc_token_service_url())
+        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.databricks_conn.login == "federated_k8s" or self.databricks_conn.extra_dejson.get(
             "federated_k8s", False
         ):
@@ -1120,6 +1189,9 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None:
                 raise AirflowException("Service Principal credentials aren't provided")
             self.log.debug("Using Service Principal Token.")
             return await self._a_get_sp_token(self._get_oidc_token_service_url())
+        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.databricks_conn.login == "federated_k8s" or self.databricks_conn.extra_dejson.get(
             "federated_k8s", False
         ):