This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7dc3c2bfe0b Add supplied-token OIDC federation to the Databricks
provider (#69272)
7dc3c2bfe0b is described below
commit 7dc3c2bfe0b57c60fb0bc2385bdad6c01431679f
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Jul 8 08:25:41 2026 +0100
Add supplied-token OIDC federation to the Databricks provider (#69272)
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.
---
.../databricks/docs/connections/databricks.rst | 41 ++++
.../providers/databricks/hooks/databricks_base.py | 114 +++++++++--
.../unit/databricks/hooks/test_databricks_base.py | 228 ++++++++++++++++++++-
3 files changed, 359 insertions(+), 24 deletions(-)
diff --git a/providers/databricks/docs/connections/databricks.rst
b/providers/databricks/docs/connections/databricks.rst
index cc830aee08f..fb7eb80157c 100644
--- a/providers/databricks/docs/connections/databricks.rst
+++ b/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, Goog [...]
diff --git
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
index c90d10cf69a..07f5500f7b3 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -25,6 +25,7 @@ operators talk to the ``api/2.0/jobs/runs/submit``
from __future__ import annotations
+import asyncio
import copy
import platform
import ssl
@@ -50,6 +51,7 @@ from tenacity import (
)
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 @@ except ImportError:
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
@@ -132,6 +136,7 @@ class BaseDatabricksHook(BaseHook):
"proxies",
"service_principal_oauth",
"federated_k8s",
+ "federated_token_provider",
"k8s_token_path",
"k8s_namespace_path",
"k8s_projected_volume_token_path",
@@ -908,12 +913,79 @@ class BaseDatabricksHook(BaseHook):
)
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
@@ -924,15 +996,12 @@ class BaseDatabricksHook(BaseHook):
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():
@@ -956,10 +1025,10 @@ class BaseDatabricksHook(BaseHook):
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"]
@@ -972,15 +1041,12 @@ class BaseDatabricksHook(BaseHook):
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():
@@ -1004,11 +1070,11 @@ class BaseDatabricksHook(BaseHook):
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"]
@@ -1098,6 +1164,9 @@ class BaseDatabricksHook(BaseHook):
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
):
@@ -1135,6 +1204,9 @@ class BaseDatabricksHook(BaseHook):
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
):
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
index 2c976312811..ed7ed054df2 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
@@ -25,7 +25,7 @@ import aiohttp
import pytest
import time_machine
from aiohttp.client_exceptions import ClientConnectorError
-from requests import exceptions as requests_exceptions
+from requests import Response, exceptions as requests_exceptions
from requests.auth import HTTPBasicAuth
from tenacity import AsyncRetrying, Future, RetryError, retry_if_exception,
stop_after_attempt, wait_fixed
@@ -43,6 +43,45 @@ from airflow.providers.databricks.hooks.databricks_base
import (
DEFAULT_CONN_ID = "databricks_default"
PROXIES = {"http": "http://proxy.example.com:8080";, "https":
"http://proxy.example.com:8443"}
+_SUPPLIED_JWT = "supplied_oidc_jwt"
+
+
+def _supplied_token_provider() -> str:
+ """Module-level ``federated_token_provider`` callable used in tests
(in-process token source)."""
+ return _SUPPLIED_JWT
+
+
+# Real, importable dotted path to the callable above, resolved the same way
``import_string`` will.
+_SUPPLIED_TOKEN_PROVIDER_PATH = (
+
f"{_supplied_token_provider.__module__}.{_supplied_token_provider.__qualname__}"
+)
+
+
+def _empty_token_provider() -> str:
+ """A misconfigured ``federated_token_provider`` that returns an empty
token (tests validation)."""
+ return ""
+
+
+_EMPTY_TOKEN_PROVIDER_PATH =
f"{_empty_token_provider.__module__}.{_empty_token_provider.__qualname__}"
+
+
+def _whitespace_token_provider() -> str:
+ """A misconfigured provider that returns a whitespace-only token (tests
the strip guard)."""
+ return " \t\n"
+
+
+_WHITESPACE_TOKEN_PROVIDER_PATH = (
+
f"{_whitespace_token_provider.__module__}.{_whitespace_token_provider.__qualname__}"
+)
+
+
+def _padded_token_provider() -> str:
+ """A provider that returns a valid token with surrounding whitespace (e.g.
a trailing newline)."""
+ return f" {_SUPPLIED_JWT}\n"
+
+
+_PADDED_TOKEN_PROVIDER_PATH =
f"{_padded_token_provider.__module__}.{_padded_token_provider.__qualname__}"
+
class TestBaseDatabricksHook:
def test_init_exception(self):
@@ -1173,9 +1212,191 @@ class TestBaseDatabricksHook:
hook.user_agent_header = {"User-Agent": "test-agent"}
resource = f"https://{mock_conn.host}/oidc/v1/token";
- with pytest.raises(AirflowException, match="Failed to exchange
Kubernetes JWT for Databricks token"):
+ with pytest.raises(
+ AirflowException, match="Failed to exchange the federated OIDC
token for a Databricks token"
+ ):
hook._get_federated_databricks_token(resource)
+ @pytest.mark.parametrize(
+ "client_id",
+ [
+ pytest.param(None, id="account-wide-policy"),
+ pytest.param("sp-client-id", id="service-principal-policy"),
+ ],
+ )
+ @mock.patch("requests.post")
+ def test_get_token_with_supplied_provider(self, mock_post, client_id):
+ """Provider supplies the subject token in-process; client_id is sent
only when configured."""
+ db_response = mock.Mock(spec=Response)
+ db_response.json.return_value = {
+ "access_token": "databricks_token",
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ }
+ db_response.raise_for_status.return_value = None
+ mock_post.return_value = db_response
+
+ extra = {"federated_token_provider": _SUPPLIED_TOKEN_PROVIDER_PATH}
+ if client_id:
+ extra["client_id"] = client_id
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login=None,
+ password=None,
+ extra=json.dumps(extra),
+ )
+
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ with mock.patch.object(hook, "_get_k8s_jwt_token") as mock_k8s:
+ token = hook._get_token()
+
+ assert token == "databricks_token"
+ # Routed through _get_token to the exchange, never touching the
Kubernetes (disk) path.
+ mock_k8s.assert_not_called()
+ assert mock_post.call_count == 1
+ assert mock_post.call_args.args[0] ==
"https://my-workspace.cloud.databricks.com/oidc/v1/token";
+ data = mock_post.call_args.kwargs["data"]
+ assert data["subject_token"] == _SUPPLIED_JWT
+ assert data["subject_token_type"] ==
"urn:ietf:params:oauth:token-type:jwt"
+ assert data["grant_type"] ==
"urn:ietf:params:oauth:grant-type:token-exchange"
+ assert data["scope"] == "all-apis"
+ if client_id:
+ assert data["client_id"] == client_id # service principal
federation policy
+ else:
+ assert "client_id" not in data # account-wide federation policy
+
+ @pytest.mark.parametrize(
+ "provider_path",
+ [
+ pytest.param(_EMPTY_TOKEN_PROVIDER_PATH, id="empty"),
+ pytest.param(_WHITESPACE_TOKEN_PROVIDER_PATH,
id="whitespace-only"),
+ ],
+ )
+ def test_get_token_supplied_provider_invalid_return(self, provider_path):
+ """A provider returning an empty or whitespace-only token fails
clearly, before any exchange."""
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login=None,
+ password=None,
+ extra=json.dumps({"federated_token_provider": provider_path}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+
+ with mock.patch("requests.post") as mock_post:
+ with pytest.raises(ValueError, match="must return a non-empty
string token"):
+ hook._get_token()
+ mock_post.assert_not_called()
+
+ @mock.patch("requests.post")
+ def test_get_token_supplied_provider_strips_token(self, mock_post):
+ """A token returned with surrounding whitespace is stripped before it
is exchanged."""
+ db_response = mock.Mock(spec=Response)
+ db_response.json.return_value = {
+ "access_token": "databricks_token",
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ }
+ db_response.raise_for_status.return_value = None
+ mock_post.return_value = db_response
+
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login=None,
+ password=None,
+ extra=json.dumps({"federated_token_provider":
_PADDED_TOKEN_PROVIDER_PATH}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ assert hook._get_token() == "databricks_token"
+ # The newline-padded token is stripped -- the raw value is never
posted to the exchange.
+ assert mock_post.call_args.kwargs["data"]["subject_token"] ==
_SUPPLIED_JWT
+
+ @mock.patch("requests.post")
+ def test_supplied_provider_takes_precedence_over_federated_k8s(self,
mock_post):
+ """With both configured, the supplied provider wins and the Kubernetes
disk path is never used."""
+ db_response = mock.Mock(spec=Response)
+ db_response.json.return_value = {
+ "access_token": "databricks_token",
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ }
+ db_response.raise_for_status.return_value = None
+ mock_post.return_value = db_response
+
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login=None,
+ password=None,
+ extra=json.dumps(
+ {
+ "federated_token_provider": _SUPPLIED_TOKEN_PROVIDER_PATH,
+ "federated_k8s": True,
+ "client_id": "sp-client-id",
+ }
+ ),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ with mock.patch.object(hook, "_get_k8s_jwt_token") as mock_k8s:
+ assert hook._get_token() == "databricks_token"
+ mock_k8s.assert_not_called()
+ assert mock_post.call_args.kwargs["data"]["subject_token"] ==
_SUPPLIED_JWT
+
+ @pytest.mark.asyncio
+ @mock.patch("aiohttp.ClientSession.post")
+ async def test_a_get_token_with_supplied_provider(self, mock_post):
+ """Async _get_token resolves the provider off the event loop and
exchanges the token."""
+ db_response = mock.AsyncMock(spec=aiohttp.ClientResponse)
+ db_response.__aenter__.return_value = db_response
+ db_response.__aexit__.return_value = None
+ db_response.raise_for_status.return_value = None
+ db_response.json.return_value = {
+ "access_token": "async_databricks_token",
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ }
+ mock_post.return_value = db_response
+
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login=None,
+ password=None,
+ extra=json.dumps({"federated_token_provider":
_SUPPLIED_TOKEN_PROVIDER_PATH}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+ hook.token_timeout_seconds = 10
+
+ with mock.patch.object(hook, "_a_get_k8s_jwt_token") as mock_k8s:
+ async with aiohttp.ClientSession() as session:
+ hook._session = session
+ token = await hook._a_get_token()
+
+ mock_k8s.assert_not_called()
+ assert token == "async_databricks_token"
+ data = mock_post.call_args.kwargs["data"]
+ assert data["subject_token"] == _SUPPLIED_JWT
+ assert "client_id" not in data
+
@mock.patch("builtins.open")
@mock.patch("requests.post")
def test_get_k8s_jwt_token_with_custom_extra_settings(self, mock_post,
mock_open):
@@ -1647,7 +1868,8 @@ class TestBaseDatabricksHook:
hook._session = session
resource = f"https://{mock_conn.host}/oidc/v1/token";
with pytest.raises(
- AirflowException, match="Failed to exchange Kubernetes JWT
for Databricks token"
+ AirflowException,
+ match="Failed to exchange the federated OIDC token for a
Databricks token",
):
await hook._a_get_federated_databricks_token(resource)
