From 2f5cc84a6f9bc2c318712ed25f972a005dfff2c6 Mon Sep 17 00:00:00 2001 From: Allen Chen <113149736+allenc97@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:51:57 -0700 Subject: [PATCH 1/3] Add AWS IAM OIDC token federation to Databricks connection --- dev/breeze/tests/test_selective_checks.py | 8 +- .../docs/connections/databricks.rst | 13 ++ providers/databricks/docs/index.rst | 4 +- providers/databricks/pyproject.toml | 4 + .../databricks/hooks/databricks_base.py | 71 +++++++- .../databricks/hooks/test_databricks_base.py | 159 +++++++++++++++++- 6 files changed, 252 insertions(+), 7 deletions(-) diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index 530e6ad74aa46..8f173a91200c8 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 fb7eb80157c74..a7e6b801e8033 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 `__ (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 `_. + + * ``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 `_. This authentication method only works when Airflow is running inside a Kubernetes cluster (e.g., AWS EKS, Azure AKS, Google GKE). diff --git a/providers/databricks/docs/index.rst b/providers/databricks/docs/index.rst index fb5f5193acd46..c95e238d777ae 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 `_ ``amazon`` `apache-airflow-providers-google `_ ``google`` `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 6b6e846bcc6c4..80da493a144f9 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 07f5500f7b371..1bc9c2de850fd 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_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,13 +926,17 @@ 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. @@ -936,6 +949,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 +963,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. @@ -1172,6 +1235,9 @@ def _get_token(self, raise_error: bool = False) -> str | None: ): self.log.debug("Using Kubernetes OIDC token federation.") 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 raise_error: raise AirflowException("Token authentication isn't configured") @@ -1212,6 +1278,9 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None: ): self.log.debug("Using Kubernetes OIDC token federation.") 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 raise_error: raise AirflowException("Token authentication isn't configured") 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 ed7ed054df258..dd44ba9dd1aee 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 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, @@ -1693,6 +1693,163 @@ def test_get_token_with_federated_k8s_extra(self, mock_file): 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) From 0f5c29f3555b3cd93f8eed3064d4ae8921abc201 Mon Sep 17 00:00:00 2001 From: Allen Chen <113149736+allenc97@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:27:20 -0700 Subject: [PATCH 2/3] Update Databricks OIDC to match described precendence --- .../databricks/hooks/databricks_base.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) 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 1bc9c2de850fd..ad48f09a9360b 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -941,6 +941,9 @@ def _get_federation_subject_token(self) -> tuple[str, str | None]: 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). """ @@ -1046,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 @@ -1230,14 +1233,14 @@ 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 ): self.log.debug("Using Kubernetes OIDC token federation.") 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 raise_error: raise AirflowException("Token authentication isn't configured") @@ -1273,14 +1276,14 @@ 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 ): self.log.debug("Using Kubernetes OIDC token federation.") 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 raise_error: raise AirflowException("Token authentication isn't configured") From 2aa6aed3a4307840e3e1748c8ba3618bb2202e47 Mon Sep 17 00:00:00 2001 From: Allen Chen <113149736+allenc97@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:10:13 -0700 Subject: [PATCH 3/3] Add regression precedence test for Databricks OIDC --- .../databricks/hooks/test_databricks_base.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) 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 dd44ba9dd1aee..9b0e340ef12e7 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -1358,6 +1358,45 @@ def test_supplied_provider_takes_precedence_over_federated_k8s(self, mock_post): 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):