This is an automated email from the ASF dual-hosted git repository.
eladkal 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 47e5d03826f Databricks OIDC token federation for AWS IAM (#69509)
47e5d03826f is described below
commit 47e5d03826fc4262e93f90ad0613579c1492d330
Author: Allen <[email protected]>
AuthorDate: Thu Jul 16 10:17:17 2026 -0700
Databricks OIDC token federation for AWS IAM (#69509)
* Add AWS IAM OIDC token federation to Databricks connection
* Update Databricks OIDC to match described precendence
* Add regression precedence test for Databricks OIDC
---
dev/breeze/tests/test_selective_checks.py | 8 +-
.../databricks/docs/connections/databricks.rst | 13 ++
providers/databricks/docs/index.rst | 4 +-
providers/databricks/pyproject.toml | 4 +
.../providers/databricks/hooks/databricks_base.py | 80 ++++++++-
.../unit/databricks/hooks/test_databricks_base.py | 198 ++++++++++++++++++++-
6 files changed, 297 insertions(+), 10 deletions(-)
diff --git a/dev/breeze/tests/test_selective_checks.py
b/dev/breeze/tests/test_selective_checks.py
index 25da64bd517..60371125ade 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/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]",
}
diff --git a/providers/databricks/docs/connections/databricks.rst
b/providers/databricks/docs/connections/databricks.rst
index fb7eb80157c..a7e6b801e80 100644
--- a/providers/databricks/docs/connections/databricks.rst
+++ b/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, Goog [...]
diff --git a/providers/databricks/docs/index.rst
b/providers/databricks/docs/index.rst
index fb5f5193acd..c95e238d777 100644
--- a/providers/databricks/docs/index.rst
+++ b/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``
diff --git a/providers/databricks/pyproject.toml
b/providers/databricks/pyproject.toml
index 6b6e846bcc6..80da493a144 100644
--- a/providers/databricks/pyproject.toml
+++ b/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",
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 07f5500f7b3..ad48f09a936 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -84,6 +84,11 @@ DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH =
"/var/run/secrets/kubernetes.io/service
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 @@ class BaseDatabricksHook(BaseHook):
"""
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
"""
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
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
):
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 ed7ed054df2..9b0e340ef12 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
@@ -30,7 +30,7 @@ from requests.auth import HTTPBasicAuth
from tenacity import AsyncRetrying, Future, RetryError, retry_if_exception,
stop_after_attempt, wait_fixed
from airflow.models import Connection
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
from airflow.providers.databricks.hooks.databricks_base import (
DEFAULT_AZURE_CREDENTIAL_SETTING_KEY,
DEFAULT_DATABRICKS_SCOPE,
@@ -1358,6 +1358,45 @@ class TestBaseDatabricksHook:
mock_k8s.assert_not_called()
assert mock_post.call_args.kwargs["data"]["subject_token"] ==
_SUPPLIED_JWT
+ @mock.patch("airflow.providers.amazon.aws.hooks.sts.StsHook")
+ @mock.patch("requests.post")
+ def test_federated_aws_takes_precedence_over_federated_k8s(self,
mock_post, mock_sts_hook):
+ """With both configured, the AWS path wins and the Kubernetes disk
path is never used."""
+
mock_sts_hook.return_value.get_conn.return_value.get_web_identity_token.return_value
= {
+ "WebIdentityToken": "aws_signed_jwt"
+ }
+ 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_aws": True,
+ "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"] ==
"aws_signed_jwt"
+
@pytest.mark.asyncio
@mock.patch("aiohttp.ClientSession.post")
async def test_a_get_token_with_supplied_provider(self, mock_post):
@@ -1693,6 +1732,163 @@ class TestBaseDatabricksHook:
assert token == "databricks_token"
+ @mock.patch("airflow.providers.amazon.aws.hooks.sts.StsHook")
+ @mock.patch("requests.post")
+ def test_get_token_with_federated_aws_login(self, mock_post,
mock_sts_hook):
+ """_get_token with login='federated_aws' mints an STS token and
exchanges it (service principal)."""
+
mock_sts_hook.return_value.get_conn.return_value.get_web_identity_token.return_value
= {
+ "WebIdentityToken": "aws_signed_jwt"
+ }
+ 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="federated_aws",
+ password=None,
+ extra=json.dumps({"client_id": "test-client-id"}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ assert hook._get_token() == "databricks_token"
+ mock_sts_hook.assert_called_once_with(aws_conn_id="aws_default")
+ data = mock_post.call_args.kwargs["data"]
+ assert data["subject_token"] == "aws_signed_jwt"
+ assert data["grant_type"] ==
"urn:ietf:params:oauth:grant-type:token-exchange"
+ assert data["client_id"] == "test-client-id"
+
+ @mock.patch("airflow.providers.amazon.aws.hooks.sts.StsHook")
+ @mock.patch("requests.post")
+ def test_get_token_with_federated_aws_extra_account_wide(self, mock_post,
mock_sts_hook):
+ """federated_aws via extra flag: custom settings honored, no client_id
(account-wide federation)."""
+ mock_client = mock_sts_hook.return_value.get_conn.return_value
+ mock_client.get_web_identity_token.return_value = {"WebIdentityToken":
"aws_signed_jwt"}
+ 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_aws": True,
+ "aws_conn_id": "aws_prod",
+ "aws_jwt_audience": "databricks-prod",
+ "aws_web_identity_token_duration": 900,
+ }
+ ),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ assert hook._get_token() == "databricks_token"
+ mock_sts_hook.assert_called_once_with(aws_conn_id="aws_prod")
+ mock_client.get_web_identity_token.assert_called_once_with(
+ Audience=["databricks-prod"], SigningAlgorithm="RS256",
DurationSeconds=900
+ )
+ data = mock_post.call_args.kwargs["data"]
+ assert data["subject_token"] == "aws_signed_jwt"
+ assert "client_id" not in data
+
+ @mock.patch("airflow.providers.amazon.aws.hooks.sts.StsHook")
+ def test_get_aws_subject_token_old_boto3(self, mock_sts_hook):
+ """federated_aws errors clearly when the installed AWS SDK lacks
GetWebIdentityToken."""
+ mock_sts_hook.return_value.get_conn.return_value =
mock.Mock(spec=["get_caller_identity"])
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login="federated_aws",
+ password=None,
+ extra=json.dumps({"client_id": "test-client-id"}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException, match="does not support
'sts:GetWebIdentityToken'"
+ ):
+ hook._get_aws_subject_token()
+
+ def test_get_aws_subject_token_amazon_not_installed(self):
+ """federated_aws errors clearly when the amazon provider is not
installed."""
+ conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host="my-workspace.cloud.databricks.com",
+ login="federated_aws",
+ password=None,
+ extra=json.dumps({"client_id": "test-client-id"}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+
+ with mock.patch.dict("sys.modules",
{"airflow.providers.amazon.aws.hooks.sts": None}):
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException,
match="apache-airflow-providers-amazon"
+ ):
+ hook._get_aws_subject_token()
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.amazon.aws.hooks.sts.StsHook")
+ @mock.patch("aiohttp.ClientSession.post")
+ @time_machine.travel("2025-07-12 12:00:00", tick=False)
+ async def test_a_get_token_with_federated_aws(self, mock_post,
mock_sts_hook):
+ """Async _a_get_token with federated_aws mints and exchanges the STS
token."""
+
mock_sts_hook.return_value.get_conn.return_value.get_web_identity_token.return_value
= {
+ "WebIdentityToken": "aws_signed_jwt"
+ }
+ db_response = mock.AsyncMock()
+ 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="federated_aws",
+ password=None,
+ extra=json.dumps({"client_id": "test-client-id"}),
+ )
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = conn
+ hook.user_agent_header = {"User-Agent": "test-agent"}
+
+ async with aiohttp.ClientSession() as session:
+ hook._session = session
+ token = await hook._a_get_token()
+
+ assert token == "async_databricks_token"
+ assert mock_post.call_args.kwargs["data"]["subject_token"] ==
"aws_signed_jwt"
+ assert mock_post.call_args.kwargs["data"]["client_id"] ==
"test-client-id"
+
@pytest.mark.asyncio
@mock.patch("aiohttp.ClientSession.post")
@time_machine.travel("2025-07-12 12:00:00", tick=False)
