From b56fe145cbc2cebd76d900fa4cbd18587085e796 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 29 Jan 2026 15:28:37 +0100 Subject: [PATCH 01/21] Added OIDC token federation as authentication option for Databricks provider --- .../docs/connections/databricks.rst | 95 ++++- .../databricks/hooks/databricks_base.py | 254 +++++++++++++ .../databricks/hooks/test_databricks_base.py | 352 ++++++++++++++++++ 3 files changed, 694 insertions(+), 7 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 2d92bb67eba8a..4f021718cee29 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -31,7 +31,7 @@ There are several ways to connect to Databricks using Airflow. 1. Use a `Personal Access Token (PAT) `_ - i.e. add a token to the Airflow connection. This is the recommended method. + i.e. add a token to the Airflow connection. 2. Use Databricks login credentials i.e. add the username and password used to login to the Databricks account to the Airflow connection. Note that username/password authentication is discouraged and not supported for @@ -41,6 +41,12 @@ There are several ways to connect to Databricks using Airflow. `user inside workspace `_, or `outside of workspace having Owner or Contributor permissions `_ 4. Using Azure Active Directory (AAD) token obtained for `Azure managed identity `_, when Airflow runs on the VM with assigned managed identity (system-assigned or user-assigned) +5. Using Databricks Service Principal OAuth + i.e. use OAuth authentication with Service Principal client ID and secret. + See `Authentication using OAuth for service principals `_. +6. Using Kubernetes `OIDC token federation `_ (applicable only when Airflow runs in Kubernetes) + i.e. automatically fetch JWT tokens from Kubernetes Service Account 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. Default Connection IDs ---------------------- @@ -57,13 +63,15 @@ Login (optional) * If authentication with *Databricks login credentials* is used then specify the ``username`` used to login to Databricks. * If authentication with *Azure Service Principal* is used then specify the ID of the Azure Service Principal * If authentication with *PAT* is used then either leave this field empty or use 'token' as login (both work, the only difference is that if login is empty then token will be sent in request header as Bearer token, if login is 'token' then it will be sent using Basic Auth which is allowed by Databricks API, this may be useful if you plan to reuse this connection with e.g. HttpOperator) - * If authentication with *Databricks Service Principal OAuth* is used then specify the ID of the Service Principal (Databricks on AWS) + * If authentication with *Databricks Service Principal OAuth* is used then specify the ID of the Service Principal + * If authentication with *Kubernetes OIDC token federation* is used then specify ``federated_k8s`` as the login Password (optional) * If authentication with *Databricks login credentials* is used then specify the ``password`` used to login to Databricks. * If authentication with *Azure Service Principal* is used then specify the secret of the Azure Service Principal * If authentication with *PAT* is used, then specify PAT (recommended) - * If authentication with *Databricks Service Principal OAuth* is used then specify the secret of the Service Principal (Databricks on AWS) + * If authentication with *Databricks Service Principal OAuth* is used then specify the secret of the Service Principal + * If authentication with *Kubernetes OIDC token federation* is used then leave this field empty (not required) Extra (optional) Specify the extra parameter (as json dictionary) that can be used in the Databricks connection. @@ -72,9 +80,9 @@ Extra (optional) * ``token``: Specify PAT to use. Consider to switch to specification of PAT in the Password field as it's more secure. - Following parameters are necessary if using authentication with OAuth token for AWS Databricks Service Principal: + Following parameters are necessary if using authentication with OAuth token for Databricks Service Principal: - * ``service_principal_oauth``: required boolean flag. If specified as ``true``, use the Client ID and Client Secret as the Username and Password. See `Authentication using OAuth for service principals `_. + * ``service_principal_oauth``: required boolean flag. If specified as ``true``, use the Client ID and Client Secret as the Username and Password. See `Authentication using OAuth for service principals `_. Following parameters are necessary if using authentication with AAD token: @@ -91,6 +99,74 @@ Extra (optional) * ``azure_resource_id``: optional Resource ID of the Azure Databricks workspace (required if managed identity isn't a user inside workspace) + The following parameters are necessary if using authentication with Kubernetes OIDC token federation: + + * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as extra parameter. When enabled, the hook will fetch a JWT token from the Kubernetes Service Account TokenRequest API 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). + * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). The hook will call the Kubernetes TokenRequest API to obtain a token with this audience. Note: The default audience works for most use cases. Only customize if required by your specific Databricks federation policy configuration. + * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). + * ``k8s_token_path``: (optional) path to the Kubernetes service account token file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). Override this for custom Kubernetes configurations. + * ``k8s_namespace_path``: (optional) path to the Kubernetes namespace file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/namespace``). Override this for custom Kubernetes configurations. + + **Example configuration for Kubernetes OIDC token federation:** + + Minimal configuration (uses defaults): + + * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. + * Set ``Host`` to your Databricks workspace URL + + With custom configuration: + + * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. + * Set ``Host`` to your Databricks workspace URL + * Add optional parameters as necessary in ``Extra`` field, e.g. ``{"audience": "custom-audience", "expiration_seconds": 7200, "k8s_token_path": "/custom/path/token", "k8s_namespace_path": "/custom/path/namespace"}`` + + **Important Notes:** + + * This method requires no secrets to be stored in the Airflow connection + * The Kubernetes service account must be configured with appropriate permissions to create TokenRequest resources (typically granted by default in most Kubernetes clusters) + * The Databricks workspace must have federation policies configured in Databricks Account for the Kubernetes identity provider. This can be Account-level or Service principal-level policies. + * Both Kubernetes JWT and Databricks OAuth tokens are short-lived and automatically refreshed. + + **Databricks Federation Policy Configuration:** + + Before using Kubernetes OIDC token federation, you must configure a federation policy in your Databricks account. The policy configuration is the same regardless of whether you create it at the account level or for a specific service principal. The difference is only where you create the policy: + + * **Account-level policy:** All users and service principals in the account can use this authentication method + * **Service principal policy:** Only the specific service principal can use this authentication method (recommended for workloads) + + **Federation Policy Configuration**: + + * **Issuer:** ``https://kubernetes.default.svc`` + * **Audience:** ``https://kubernetes.default.svc`` (default) + * **Subject:** ``system:serviceaccount::`` + * **Token Signature Validation (JWKS JSON):** (Optional) Use your Kubernetes cluster's public keys or configure Databricks to fetch them from the well-known endpoint + + Example matching JWT token that Kubernetes will provide: + + .. code-block:: json + + { + "iss": "https://kubernetes.default.svc", + "aud": ["https://kubernetes.default.svc"], + "sub": "system:serviceaccount:airflow:airflow-worker" + } + + (where ``airflow`` is your Kubernetes namespace and ``airflow-worker`` is your service account name) + + **Troubleshooting Common Issues:** + + * **"Kubernetes service account token not found" error:** This authentication method only works when Airflow is running inside a Kubernetes cluster. Ensure your pods have the service account token mounted (default behavior in Kubernetes). + + * **Permission denied errors:** Verify your Kubernetes service account has permission to create TokenRequest resources. In most clusters, this is granted by default to all service accounts. + + * **Token exchange failures:** Ensure your Databricks federation policy configuration matches the JWT token properties: + + - The Issuer in your policy must match the ``iss`` claim in the JWT (typically ``https://kubernetes.default.svc``) + - The Audience in your policy must match the ``aud`` claim (set via the ``audience`` connection extra, default is ``https://kubernetes.default.svc``) + - The Subject pattern in your policy must match the ``sub`` claim (``system:serviceaccount::``) + + **Note:** For complete step-by-step instructions on creating and configuring federation policies in Databricks, including how to obtain and configure JWKS public keys, see the `Databricks federation policy configuration guide `_. + Following parameters could be set when using :class:`~airflow.providers.databricks.operators.databricks_sql.DatabricksSqlOperator`: @@ -99,8 +175,7 @@ Extra (optional) * named internal arguments to the ``Connection`` object from ``databricks-sql-connector`` package. -When specifying the connection using an environment variable you should specify -it using URI syntax. +When specifying the connection using an environment variable you should specify it using URI syntax. Note that all components of the URI should be URL-encoded. @@ -109,3 +184,9 @@ For example: .. code-block:: bash export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://@host-url?token=yourtoken' + + # For Kubernetes OIDC token federation (minimal): + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com' + + # For Kubernetes OIDC token federation with custom audience and expiration: + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?audience=custom-audience&expiration_seconds=7200' 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 5cf4049a43de1..6af9ecb80e174 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -33,6 +33,7 @@ from typing import TYPE_CHECKING, Any from urllib.parse import urlsplit +import aiofiles import aiohttp import requests from aiohttp.client_exceptions import ClientConnectorError @@ -71,6 +72,12 @@ DEFAULT_AZURE_CREDENTIAL_SETTING_KEY = "use_default_azure_credential" +# Kubernetes OIDC token federation +K8S_TOKEN_SERVICE_URL = "https://kubernetes.default.svc" +DEFAULT_K8S_AUDIENCE = "https://kubernetes.default.svc" +DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token" +DEFAULT_K8S_NAMESPACE_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/namespace" + class BaseDatabricksHook(BaseHook): """ @@ -100,6 +107,9 @@ class BaseDatabricksHook(BaseHook): "azure_resource_id", "azure_tenant_id", "service_principal_oauth", + "federated_k8s", + "k8s_token_path", + "k8s_namespace_path", ] def __init__( @@ -525,6 +535,236 @@ def _is_oauth_token_valid(token: dict, time_key="expires_on") -> bool: return int(token[time_key]) > (int(time.time()) + TOKEN_REFRESH_LEAD_TIME) + def _get_k8s_jwt_token(self) -> str: + """ + Get JWT token from Kubernetes Service Account Token service using TokenRequest API. + + :return: JWT Service Account token string + """ + audience = self.databricks_conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) + expiration_seconds = self.databricks_conn.extra_dejson.get("expiration_seconds", 3600) + token_path = self.databricks_conn.extra_dejson.get( + "k8s_token_path", DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH + ) + namespace_path = self.databricks_conn.extra_dejson.get( + "k8s_namespace_path", DEFAULT_K8S_NAMESPACE_PATH + ) + + try: + with open(token_path) as f: + in_cluster_token = f.read().strip() + + with open(namespace_path) as f: + namespace = f.read().strip() + + # Call Kubernetes TokenRequest API with the in-cluster token + token_request_url = ( + f"{K8S_TOKEN_SERVICE_URL}/api/v1/namespaces/{namespace}/serviceaccounts/default/token" + ) + + for attempt in self._get_retry_object(): + with attempt: + resp = requests.post( + token_request_url, + headers={ + "Authorization": f"Bearer {in_cluster_token}", + "Content-Type": "application/json", + }, + json={ + "apiVersion": "authentication.k8s.io/v1", + "kind": "TokenRequest", + "spec": { + "audiences": [audience], + "expirationSeconds": expiration_seconds, + }, + }, + verify=False, # K8s in-cluster uses self-signed certs + timeout=self.token_timeout_seconds, + ) + resp.raise_for_status() + return resp.json()["status"]["token"] + except FileNotFoundError as e: + raise AirflowException( + "Kubernetes service account token not found. " + "This authentication method only works when running inside a Kubernetes cluster." + ) from e + except RetryError: + raise AirflowException( + f"Failed to get Kubernetes JWT token after {self.retry_limit} retries. Giving up." + ) + except requests_exceptions.HTTPError as e: + msg = f"Failed to get Kubernetes JWT token. Response: {e.response.content.decode()}, Status Code: {e.response.status_code}" + raise AirflowException(msg) + + async def _a_get_k8s_jwt_token(self) -> str: + """Async version of _get_k8s_jwt_token().""" + audience = self.databricks_conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) + expiration_seconds = self.databricks_conn.extra_dejson.get("expiration_seconds", 3600) + token_path = self.databricks_conn.extra_dejson.get( + "k8s_token_path", DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH + ) + namespace_path = self.databricks_conn.extra_dejson.get( + "k8s_namespace_path", DEFAULT_K8S_NAMESPACE_PATH + ) + + try: + async with aiofiles.open(token_path) as f: + in_cluster_token = (await f.read()).strip() + + async with aiofiles.open(namespace_path) as f: + namespace = (await f.read()).strip() + + # Call Kubernetes TokenRequest API with the in-cluster token + token_request_url = ( + f"{K8S_TOKEN_SERVICE_URL}/api/v1/namespaces/{namespace}/serviceaccounts/default/token" + ) + + async for attempt in self._a_get_retry_object(): + with attempt: + async with self._session.post( + token_request_url, + headers={ + "Authorization": f"Bearer {in_cluster_token}", + "Content-Type": "application/json", + }, + json={ + "apiVersion": "authentication.k8s.io/v1", + "kind": "TokenRequest", + "spec": { + "audiences": [audience], + "expirationSeconds": expiration_seconds, + }, + }, + ssl=False, # K8s in-cluster uses self-signed certs + timeout=self.token_timeout_seconds, + ) as resp: + resp.raise_for_status() + jsn = await resp.json() + return jsn["status"]["token"] + except FileNotFoundError as e: + raise AirflowException( + "Kubernetes service account token not found. " + "This authentication method only works when running inside a Kubernetes cluster." + ) from e + except RetryError: + raise AirflowException( + f"Failed to get Kubernetes JWT token after {self.retry_limit} retries. Giving up." + ) + except aiohttp.ClientResponseError as err: + raise AirflowException( + f"Failed to get Kubernetes JWT token. Response: {err.message}, Status Code: {err.status}" + ) + + def _get_federated_databricks_token(self, resource: str) -> str: + """ + Get Databricks OAuth token by exchanging Kubernetes JWT token. + + Uses RFC 8693 token exchange to convert a Kubernetes service account JWT + into a Databricks OAuth token. Supports both account-level and service principal + federation policies. + + :param resource: Databricks OIDC token exchange URL + :return: Databricks OAuth access token + """ + federated_token = self.oauth_tokens.get(resource) + if federated_token and self._is_oauth_token_valid(federated_token): + return federated_token["access_token"] + + self.log.info("Existing federated token is expired or missing. Fetching new token...") + + # Get JWT from Kubernetes + jwt_token = self._get_k8s_jwt_token() + + # Prepare token exchange request following RFC 8693 + token_exchange_url = resource + + data = { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token": jwt_token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "scope": "all-apis", + } + + try: + for attempt in self._get_retry_object(): + with attempt: + resp = requests.post( + token_exchange_url, + data=data, + headers={ + **self.user_agent_header, + "Content-Type": "application/x-www-form-urlencoded", + }, + timeout=self.token_timeout_seconds, + ) + resp.raise_for_status() + jsn = resp.json() + jsn["expires_on"] = int(time.time() + jsn["expires_in"]) + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn + break + except RetryError: + raise AirflowException( + f"Failed to exchange Kubernetes JWT for 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}" + raise AirflowException(msg) + + return jsn["access_token"] + + async def _a_get_federated_databricks_token(self, resource: str) -> str: + """Async version of _get_federated_databricks_token().""" + federated_token = self.oauth_tokens.get(resource) + if federated_token and self._is_oauth_token_valid(federated_token): + return federated_token["access_token"] + + self.log.info("Existing federated token is expired or missing. Fetching new token...") + + # Get JWT from Kubernetes + jwt_token = await self._a_get_k8s_jwt_token() + + # Prepare token exchange request following RFC 8693 + token_exchange_url = resource + + data = { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token": jwt_token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "scope": "all-apis", + } + + try: + async for attempt in self._a_get_retry_object(): + with attempt: + async with self._session.post( + token_exchange_url, + data=data, + headers={ + **self.user_agent_header, + "Content-Type": "application/x-www-form-urlencoded", + }, + timeout=self.token_timeout_seconds, + ) as resp: + resp.raise_for_status() + jsn = await resp.json() + jsn["expires_on"] = int(time.time() + jsn["expires_in"]) + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn + break + except RetryError: + raise AirflowException( + f"Failed to exchange Kubernetes JWT for 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}" + ) + + return jsn["access_token"] + def _check_azure_metadata_service(self) -> None: """ Check for Azure Metadata Service (with caching). @@ -610,6 +850,13 @@ 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(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) + 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( + OIDC_TOKEN_SERVICE_URL.format(f"https://{self.databricks_conn.host}") + ) if raise_error: raise AirflowException("Token authentication isn't configured") @@ -642,6 +889,13 @@ 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(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) + 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( + OIDC_TOKEN_SERVICE_URL.format(f"https://{self.databricks_conn.host}") + ) 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 89beb6a8e9cae..121e63497890a 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -910,3 +910,355 @@ async def test_a_check_azure_metadata_service_retry_error(self, mock_get): ): await hook._a_check_azure_metadata_service() assert mock_get.call_count == 3 + + @mock.patch( + "builtins.open", + side_effect=[ + mock.mock_open(read_data="in_cluster_token").return_value, + mock.mock_open(read_data="default").return_value, + ], + ) + @mock.patch("requests.post") + @time_machine.travel("2025-07-12 12:00:00") # mock current timestamp for token expiry calculation + def test_get_federated_token(self, mock_post, mock_file): + """Test Kubernetes OIDC token federation flow.""" + # Mock K8s TokenRequest API response + k8s_response = mock.Mock() + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + k8s_response.raise_for_status.return_value = None + + # Mock Databricks token exchange response + expiry_date = int((datetime(2025, 7, 12, 12, 0, 0) + timedelta(minutes=60)).timestamp()) + db_response = mock.Mock() + 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.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + resource = f"https://{mock_conn.host}/oidc/v1/token" + token = hook._get_federated_databricks_token(resource) + + assert token == "databricks_token" + assert resource in hook.oauth_tokens + assert hook.oauth_tokens[resource]["access_token"] == "databricks_token" + assert hook.oauth_tokens[resource]["expires_on"] == expiry_date + + # Verify both K8s and Databricks calls were made + assert mock_post.call_count == 2 + + # Verify K8s TokenRequest API call + k8s_call_args = mock_post.call_args_list[0] + assert ( + "https://kubernetes.default.svc/api/v1/namespaces/default/serviceaccounts/default/token" + in k8s_call_args[0] + ) + + # Verify K8s request headers + assert k8s_call_args[1]["headers"]["Authorization"] == "Bearer in_cluster_token" + assert k8s_call_args[1]["headers"]["Content-Type"] == "application/json" + + # Verify K8s TokenRequest JSON payload + k8s_json = k8s_call_args[1]["json"] + assert k8s_json["apiVersion"] == "authentication.k8s.io/v1" + assert k8s_json["kind"] == "TokenRequest" + assert k8s_json["spec"]["audiences"] == ["https://kubernetes.default.svc"] + assert k8s_json["spec"]["expirationSeconds"] == 3600 + + # Verify SSL verification is disabled for in-cluster self-signed certs + assert k8s_call_args[1]["verify"] is False + + # Verify Databricks token exchange call + db_call_args = mock_post.call_args_list[1] + assert "https://my-workspace.cloud.databricks.com/oidc/v1/token" in db_call_args[0] + assert db_call_args[1]["data"]["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert db_call_args[1]["data"]["subject_token"] == "k8s_jwt_token" + assert db_call_args[1]["data"]["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" + assert db_call_args[1]["data"]["scope"] == "all-apis" + assert "client_id" not in db_call_args[1]["data"] + + @time_machine.travel("2025-07-12 12:00:00") + def test_get_federated_token_cached_valid(self): + """Test that cached valid token is returned without fetching new one.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + # Set expiration far in the future + future_expiry = int(datetime(2025, 7, 12, 12, 0, 0).timestamp()) + 10000 + hook.oauth_tokens[resource] = { + "access_token": "cached_token", + "expires_on": future_expiry, + "token_type": "Bearer", + } + + with mock.patch.object(hook, "_is_oauth_token_valid", return_value=True): + token = hook._get_federated_databricks_token(resource) + + assert token == "cached_token" + + @mock.patch("builtins.open", side_effect=FileNotFoundError()) + def test_get_federated_token_k8s_not_available(self, mock_file): + """Test error when Kubernetes service account token is not available.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + with pytest.raises(AirflowException, match="Kubernetes service account token not found"): + hook._get_federated_databricks_token(resource) + + @mock.patch( + "builtins.open", + side_effect=[ + mock.mock_open(read_data="in_cluster_token").return_value, + mock.mock_open(read_data="default").return_value, + ], + ) + @mock.patch("requests.post") + def test_get_federated_token_databricks_error(self, mock_post, mock_file): + """Test error handling when Databricks token exchange fails.""" + # K8s call succeeds + k8s_response = mock.Mock() + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + k8s_response.raise_for_status.return_value = None + + # Databricks call fails + db_response = mock.Mock() + db_response.status_code = 401 + db_response.content.decode.return_value = "Unauthorized: Invalid JWT" + http_error = requests_exceptions.HTTPError() + http_error.response = db_response + + mock_post.side_effect = [k8s_response, http_error] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + 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"): + hook._get_federated_databricks_token(resource) + + @mock.patch("builtins.open") + @mock.patch("requests.post") + def test_get_k8s_jwt_token_with_custom_extra_settings(self, mock_post, mock_open): + """Test fetching K8s JWT with custom audience, expiration, and file paths via TokenRequest API.""" + # Mock file reads for custom paths + mock_token_file = mock.mock_open(read_data="custom_in_cluster_token") + mock_namespace_file = mock.mock_open(read_data="custom-namespace") + mock_open.side_effect = [ + mock_token_file.return_value, + mock_namespace_file.return_value, + ] + + mock_response = mock.Mock() + mock_response.json.return_value = {"status": {"token": "custom_settings_jwt"}} + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = { + "audience": "custom-audience", + "expiration_seconds": 7200, + "k8s_token_path": "/custom/path/token", + "k8s_namespace_path": "/custom/path/namespace", + } + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + jwt_token = hook._get_k8s_jwt_token() + + assert jwt_token == "custom_settings_jwt" + + # Verify custom file paths were opened + assert mock_open.call_count == 2 + mock_open.assert_any_call("/custom/path/token") + mock_open.assert_any_call("/custom/path/namespace") + + # Verify TokenRequest API was called with correct namespace + assert mock_post.call_count == 1 + call_args = mock_post.call_args + assert ( + "https://kubernetes.default.svc/api/v1/namespaces/custom-namespace/serviceaccounts/default/token" + in call_args[0] + ) + + # Verify custom audience and expiration were used + assert call_args[1]["json"]["spec"]["audiences"] == ["custom-audience"] + assert call_args[1]["json"]["spec"]["expirationSeconds"] == 7200 + + # Verify Authorization header uses custom token + assert call_args[1]["headers"]["Authorization"] == "Bearer custom_in_cluster_token" + + @mock.patch( + "builtins.open", + side_effect=[ + mock.mock_open(read_data="in_cluster_token").return_value, + mock.mock_open(read_data="default").return_value, + ], + ) + def test_get_token_with_federated_k8s_login(self, mock_file): + """Test _get_token with login='federated_k8s'.""" + with mock.patch("requests.post") as mock_post: + # Mock K8s TokenRequest API response + k8s_response = mock.Mock() + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + k8s_response.raise_for_status.return_value = None + + # Mock Databricks token exchange response + db_response = mock.Mock() + 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.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.password = None + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + token = hook._get_token() + + assert token == "databricks_token" + + @mock.patch( + "builtins.open", + side_effect=[ + mock.mock_open(read_data="in_cluster_token").return_value, + mock.mock_open(read_data="default").return_value, + ], + ) + def test_get_token_with_federated_k8s_extra(self, mock_file): + """Test _get_token with federated_k8s in extras.""" + with mock.patch("requests.post") as mock_post: + # Mock K8s TokenRequest API response + k8s_response = mock.Mock() + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + k8s_response.raise_for_status.return_value = None + + # Mock Databricks token exchange response + db_response = mock.Mock() + 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.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = None + mock_conn.password = None + mock_conn.extra_dejson = {"federated_k8s": True} # Use default audience + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + token = hook._get_token() + + assert token == "databricks_token" + + @pytest.mark.asyncio + @mock.patch("aiohttp.ClientSession.post") + @time_machine.travel("2025-07-12 12:00:00") + async def test_a_get_federated_token(self, mock_post): + """Test async version of federated token exchange.""" + expiry_date = int((datetime(2025, 7, 12, 12, 0, 0) + timedelta(minutes=60)).timestamp()) + + # Mock aiofiles.open for reading K8s token and namespace + mock_token_file = mock.AsyncMock() + mock_token_file.read = mock.AsyncMock(return_value="in_cluster_token") + mock_token_file.__aenter__ = mock.AsyncMock(return_value=mock_token_file) + mock_token_file.__aexit__ = mock.AsyncMock(return_value=None) + + mock_namespace_file = mock.AsyncMock() + mock_namespace_file.read = mock.AsyncMock(return_value="default") + mock_namespace_file.__aenter__ = mock.AsyncMock(return_value=mock_namespace_file) + mock_namespace_file.__aexit__ = mock.AsyncMock(return_value=None) + + # Mock K8s TokenRequest API response + k8s_response = mock.AsyncMock() + k8s_response.__aenter__.return_value = k8s_response + k8s_response.__aexit__.return_value = None + k8s_response.raise_for_status.return_value = None + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + + # Mock Databricks token exchange response + 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.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Use default extra settings + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + hook.token_timeout_seconds = 10 + + with mock.patch("aiofiles.open", side_effect=[mock_token_file, mock_namespace_file]): + async with aiohttp.ClientSession() as session: + hook._session = session + resource = f"https://{mock_conn.host}/oidc/v1/token" + token = await hook._a_get_federated_databricks_token(resource) + + assert token == "async_databricks_token" + assert resource in hook.oauth_tokens + assert hook.oauth_tokens[resource]["access_token"] == "async_databricks_token" + assert hook.oauth_tokens[resource]["expires_on"] == expiry_date + + # Verify both K8s and Databricks calls were made + assert mock_post.call_count == 2 From ff67dc9f5e177eaed9372bbef6ddc012ce6b31c0 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 29 Jan 2026 16:11:04 +0100 Subject: [PATCH 02/21] Added OIDC token federation as authentication option for Databricks provider --- .../src/airflow/providers/databricks/hooks/databricks_base.py | 3 +++ 1 file changed, 3 insertions(+) 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 6af9ecb80e174..9c33f739f455e 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -596,6 +596,8 @@ def _get_k8s_jwt_token(self) -> str: msg = f"Failed to get Kubernetes JWT token. Response: {e.response.content.decode()}, Status Code: {e.response.status_code}" raise AirflowException(msg) + raise RuntimeError("Failed to get JWT token") + async def _a_get_k8s_jwt_token(self) -> str: """Async version of _get_k8s_jwt_token().""" audience = self.databricks_conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) @@ -654,6 +656,7 @@ async def _a_get_k8s_jwt_token(self) -> str: raise AirflowException( f"Failed to get Kubernetes JWT token. Response: {err.message}, Status Code: {err.status}" ) + raise RuntimeError("Failed to get JWT token") def _get_federated_databricks_token(self, resource: str) -> str: """ From 32425334b08c4cdd156b7d91de3303bb556f89a0 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 29 Jan 2026 17:05:51 +0100 Subject: [PATCH 03/21] Increased test coverage --- .../databricks/hooks/test_databricks_base.py | 204 ++++++++++++++++++ 1 file changed, 204 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 121e63497890a..7ba585b921ce9 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -1262,3 +1262,207 @@ async def test_a_get_federated_token(self, mock_post): # Verify both K8s and Databricks calls were made assert mock_post.call_count == 2 + + @pytest.mark.asyncio + @time_machine.travel("2025-07-12 12:00:00") + async def test_a_get_federated_token_cached_valid(self): + """Test that async version returns cached valid token without fetching new one.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + # Set expiration far in the future + future_expiry = int(datetime(2025, 7, 12, 12, 0, 0).timestamp()) + 10000 + hook.oauth_tokens[resource] = { + "access_token": "cached_async_token", + "expires_on": future_expiry, + "token_type": "Bearer", + } + + with mock.patch.object(hook, "_is_oauth_token_valid", return_value=True): + token = await hook._a_get_federated_databricks_token(resource) + + assert token == "cached_async_token" + + @pytest.mark.asyncio + async def test_a_get_federated_token_k8s_not_available(self): + """Test async error when Kubernetes service account token is not available.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + with mock.patch("aiofiles.open", side_effect=FileNotFoundError()): + with pytest.raises(AirflowException, match="Kubernetes service account token not found"): + await hook._a_get_federated_databricks_token(resource) + + @pytest.mark.asyncio + @mock.patch("aiohttp.ClientSession.post") + async def test_a_get_federated_token_databricks_error(self, mock_post): + """Test async error handling when Databricks token exchange fails.""" + # Mock aiofiles.open for reading K8s token and namespace + mock_token_file = mock.AsyncMock() + mock_token_file.read = mock.AsyncMock(return_value="in_cluster_token") + mock_token_file.__aenter__ = mock.AsyncMock(return_value=mock_token_file) + mock_token_file.__aexit__ = mock.AsyncMock(return_value=None) + + mock_namespace_file = mock.AsyncMock() + mock_namespace_file.read = mock.AsyncMock(return_value="default") + mock_namespace_file.__aenter__ = mock.AsyncMock(return_value=mock_namespace_file) + mock_namespace_file.__aexit__ = mock.AsyncMock(return_value=None) + + # Mock K8s TokenRequest API response (succeeds) + k8s_response = mock.AsyncMock() + k8s_response.__aenter__.return_value = k8s_response + k8s_response.__aexit__.return_value = None + k8s_response.raise_for_status.return_value = None + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + + # Create ClientResponseError for Databricks call + request_info = mock.Mock() + request_info.url = "https://my-workspace.cloud.databricks.com/oidc/v1/token" + db_error = aiohttp.ClientResponseError( + request_info=request_info, history=(), status=401, message="Unauthorized: Invalid JWT" + ) + + # Mock Databricks response (fails) + db_response = mock.AsyncMock() + db_response.__aenter__.return_value = db_response + db_response.__aexit__.return_value = None + db_response.raise_for_status = mock.Mock(side_effect=db_error) + + mock_post.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + hook.token_timeout_seconds = 10 + + with mock.patch("aiofiles.open", side_effect=[mock_token_file, mock_namespace_file]): + async with aiohttp.ClientSession() as session: + 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" + ): + await hook._a_get_federated_databricks_token(resource) + + @pytest.mark.asyncio + @mock.patch("aiohttp.ClientSession.post") + async def test_a_get_token_with_federated_k8s_login(self, mock_post): + """Test async _a_get_token with login='federated_k8s'.""" + # Mock aiofiles.open for reading K8s token and namespace + mock_token_file = mock.AsyncMock() + mock_token_file.read = mock.AsyncMock(return_value="in_cluster_token") + mock_token_file.__aenter__ = mock.AsyncMock(return_value=mock_token_file) + mock_token_file.__aexit__ = mock.AsyncMock(return_value=None) + + mock_namespace_file = mock.AsyncMock() + mock_namespace_file.read = mock.AsyncMock(return_value="default") + mock_namespace_file.__aenter__ = mock.AsyncMock(return_value=mock_namespace_file) + mock_namespace_file.__aexit__ = mock.AsyncMock(return_value=None) + + # Mock K8s TokenRequest API response + k8s_response = mock.AsyncMock() + k8s_response.__aenter__.return_value = k8s_response + k8s_response.__aexit__.return_value = None + k8s_response.raise_for_status.return_value = None + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + + # Mock Databricks token exchange response + 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.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.password = None + mock_conn.extra_dejson = {} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + hook.token_timeout_seconds = 10 + + with mock.patch("aiofiles.open", side_effect=[mock_token_file, mock_namespace_file]): + async with aiohttp.ClientSession() as session: + hook._session = session + token = await hook._a_get_token() + + assert token == "async_databricks_token" + + @pytest.mark.asyncio + @mock.patch("aiohttp.ClientSession.post") + async def test_a_get_token_with_federated_k8s_extra(self, mock_post): + """Test async _a_get_token with federated_k8s in extras.""" + # Mock aiofiles.open for reading K8s token and namespace + mock_token_file = mock.AsyncMock() + mock_token_file.read = mock.AsyncMock(return_value="in_cluster_token") + mock_token_file.__aenter__ = mock.AsyncMock(return_value=mock_token_file) + mock_token_file.__aexit__ = mock.AsyncMock(return_value=None) + + mock_namespace_file = mock.AsyncMock() + mock_namespace_file.read = mock.AsyncMock(return_value="default") + mock_namespace_file.__aenter__ = mock.AsyncMock(return_value=mock_namespace_file) + mock_namespace_file.__aexit__ = mock.AsyncMock(return_value=None) + + # Mock K8s TokenRequest API response + k8s_response = mock.AsyncMock() + k8s_response.__aenter__.return_value = k8s_response + k8s_response.__aexit__.return_value = None + k8s_response.raise_for_status.return_value = None + k8s_response.json.return_value = {"status": {"token": "k8s_jwt_token"}} + + # Mock Databricks token exchange response + 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_extra", + "expires_in": 3600, + "token_type": "Bearer", + } + + mock_post.side_effect = [k8s_response, db_response] + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = None + mock_conn.password = None + mock_conn.extra_dejson = {"federated_k8s": True} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + hook.token_timeout_seconds = 10 + + with mock.patch("aiofiles.open", side_effect=[mock_token_file, mock_namespace_file]): + async with aiohttp.ClientSession() as session: + hook._session = session + token = await hook._a_get_token() + + assert token == "async_databricks_token_extra" From 79930b37561cc44b296ee57a159f1036db740ffa Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 29 Jan 2026 19:57:18 +0100 Subject: [PATCH 04/21] updated docs --- .../docs/connections/databricks.rst | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 4f021718cee29..de14c151f0696 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -141,6 +141,28 @@ Extra (optional) * **Subject:** ``system:serviceaccount::`` * **Token Signature Validation (JWKS JSON):** (Optional) Use your Kubernetes cluster's public keys or configure Databricks to fetch them from the well-known endpoint + **Complete Federation Policy Example:** + + .. code-block:: json + + { + "oidc_policy": { + "issuer": "https://kubernetes.default.svc", + "audiences": ["https://kubernetes.default.svc"], + "subject": "system:serviceaccount:airflow:airflow-worker", + "jwks_json": { + "keys": [{ + "kty": "RSA", + "e": "AQAB", + "use": "sig", + "kid": "your-cluster-key-id", + "alg": "RS256", + "n": "your-cluster-public-key-modulus..." + }] + } + } + } + Example matching JWT token that Kubernetes will provide: .. code-block:: json From 0c831a26c93e5d656fb4891268fa3077834b366d Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 29 Jan 2026 23:15:05 +0100 Subject: [PATCH 05/21] updated docs --- .../docs/connections/databricks.rst | 79 +++++++++++++++++-- 1 file changed, 71 insertions(+), 8 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index de14c151f0696..0ae75b4929c89 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -131,17 +131,56 @@ Extra (optional) Before using Kubernetes OIDC token federation, you must configure a federation policy in your Databricks account. The policy configuration is the same regardless of whether you create it at the account level or for a specific service principal. The difference is only where you create the policy: - * **Account-level policy:** All users and service principals in the account can use this authentication method - * **Service principal policy:** Only the specific service principal can use this authentication method (recommended for workloads) + * **Account-level policy:** Configured in the global authentication settings. All users and service principals in the account can use this authentication method. + * **Service principal policy:** Configured in the service principal settings. Only the specific service principal can use this authentication method (recommended for workloads). - **Federation Policy Configuration**: + Before configuring the federation policy, you need to understand the token issuer and JWKS requirements. + + **Understanding Token Issuer and JWKS Requirements:** + + The token issuer (``iss`` claim) is determined by your Kubernetes cluster's OIDC configuration, **not** by Airflow. The issuer affects whether JWKS JSON must be provided in the Databricks federation policy: + + * **Standard Kubernetes (self-managed or most managed clusters):** Issuer is ``https://kubernetes.default.svc`` (internal cluster address) + + - **JWKS JSON:** - (Required) You must manually provide your cluster's public keys in the federation policy + - **Why:** Databricks cannot reach the internal cluster address to fetch keys automatically + - This is the most common scenario for Airflow deployments + + * **Managed Kubernetes with public OIDC endpoints (AWS EKS with IRSA, Azure AKS with Workload Identity, GKE with Workload Identity):** Issuer is a public URL + + - **JWKS JSON:** - (Optional) Databricks can automatically fetch keys from the public endpoint + - **Example issuers:** + + - AWS EKS: ``https://oidc.eks.us-west-2.amazonaws.com/id/CLUSTER_ID`` + - Azure AKS: ``https://westus.oic.prod-aks.azure.com/TENANT_ID/CLUSTER_UUID/`` + - GKE: ``https://container.googleapis.com/v1/projects/PROJECT/locations/REGION/clusters/CLUSTER`` + + **How to Determine Your Cluster's Issuer:** + + To find out what issuer your cluster uses, inspect an actual JWT token: + + .. code-block:: bash + + # Request a token from your cluster + token=$(kubectl create token --audience="test" --duration=600s) + + # Decode the token + echo $token | cut -d. -f2 | base64 -d | jq . + + # Look for the "iss" (issuer) claim in the decoded token + + **DatabricksFederation Policy Configuration:** + + **Scenario 1: Standard Kubernetes (Most Common) - JWKS JSON Required** + + Federation policy configuration: * **Issuer:** ``https://kubernetes.default.svc`` - * **Audience:** ``https://kubernetes.default.svc`` (default) - * **Subject:** ``system:serviceaccount::`` - * **Token Signature Validation (JWKS JSON):** (Optional) Use your Kubernetes cluster's public keys or configure Databricks to fetch them from the well-known endpoint + * **Audience:** ``https://kubernetes.default.svc`` (default, customizable via connection extra) + * **Subject:** ``system:serviceaccount::`` + * **JWKS JSON:** - (Required) Your cluster's public signing keys - **Complete Federation Policy Example:** + Complete Databricks federation policy with JWKS JSON: .. code-block:: json @@ -163,7 +202,31 @@ Extra (optional) } } - Example matching JWT token that Kubernetes will provide: + **How to obtain your cluster's JWKS:** + + .. code-block:: bash + + kubectl get --raw /openid/v1/jwks + + Copy the entire output and use it as the value for ``jwks_json`` in your Databricks federation policy. Each Kubernetes cluster has unique signing keys, which is what provides the security - only tokens signed by your specific cluster can be validated. + + **Scenario 2: Managed Kubernetes with Public OIDC - JWKS JSON Optional** + + If your cluster uses a public OIDC issuer (check using the method presented above), you can omit JWKS JSON: + + .. code-block:: json + + { + "oidc_policy": { + "issuer": "https://oidc.eks.us-west-2.amazonaws.com/id/YOUR_CLUSTER_ID", + "audiences": ["sts.amazonaws.com"], + "subject": "system:serviceaccount:airflow:airflow-worker" + } + } + + In this case, Databricks will automatically fetch the public keys from ``/.well-known/openid-configuration``. + + Example JWT token from standard Kubernetes: .. code-block:: json From a21aae789d3c6f5848e7f44fa9c1bf54b1a987e2 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 09:24:30 +0100 Subject: [PATCH 06/21] updated docs --- .../docs/connections/databricks.rst | 40 ++++++++++++++----- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 0ae75b4929c89..a5b0825666f84 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -41,8 +41,8 @@ There are several ways to connect to Databricks using Airflow. `user inside workspace `_, or `outside of workspace having Owner or Contributor permissions `_ 4. Using Azure Active Directory (AAD) token obtained for `Azure managed identity `_, when Airflow runs on the VM with assigned managed identity (system-assigned or user-assigned) -5. Using Databricks Service Principal OAuth - i.e. use OAuth authentication with Service Principal client ID and secret. +5. Using Databricks-managed Service Principal OAuth + i.e. use OAuth authentication with Databricks-managed Service Principal client ID and secret. See `Authentication using OAuth for service principals `_. 6. Using Kubernetes `OIDC token federation `_ (applicable only when Airflow runs in Kubernetes) i.e. automatically fetch JWT tokens from Kubernetes Service Account and exchange them for Databricks OAuth tokens. @@ -63,14 +63,14 @@ Login (optional) * If authentication with *Databricks login credentials* is used then specify the ``username`` used to login to Databricks. * If authentication with *Azure Service Principal* is used then specify the ID of the Azure Service Principal * If authentication with *PAT* is used then either leave this field empty or use 'token' as login (both work, the only difference is that if login is empty then token will be sent in request header as Bearer token, if login is 'token' then it will be sent using Basic Auth which is allowed by Databricks API, this may be useful if you plan to reuse this connection with e.g. HttpOperator) - * If authentication with *Databricks Service Principal OAuth* is used then specify the ID of the Service Principal + * If authentication with *Databricks-managed Service Principal OAuth* is used then specify the ID of the Databricks-managed Service Principal * If authentication with *Kubernetes OIDC token federation* is used then specify ``federated_k8s`` as the login Password (optional) * If authentication with *Databricks login credentials* is used then specify the ``password`` used to login to Databricks. * If authentication with *Azure Service Principal* is used then specify the secret of the Azure Service Principal * If authentication with *PAT* is used, then specify PAT (recommended) - * If authentication with *Databricks Service Principal OAuth* is used then specify the secret of the Service Principal + * If authentication with *Databricks-managed Service Principal OAuth* is used then specify the secret of the Databricks-managed Service Principal * If authentication with *Kubernetes OIDC token federation* is used then leave this field empty (not required) Extra (optional) @@ -80,7 +80,7 @@ Extra (optional) * ``token``: Specify PAT to use. Consider to switch to specification of PAT in the Password field as it's more secure. - Following parameters are necessary if using authentication with OAuth token for Databricks Service Principal: + Following parameters are necessary if using authentication with OAuth token for Databricks-managed Service Principal: * ``service_principal_oauth``: required boolean flag. If specified as ``true``, use the Client ID and Client Secret as the Username and Password. See `Authentication using OAuth for service principals `_. @@ -102,7 +102,7 @@ Extra (optional) The following parameters are necessary if using authentication with Kubernetes OIDC token federation: * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as extra parameter. When enabled, the hook will fetch a JWT token from the Kubernetes Service Account TokenRequest API 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). - * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). The hook will call the Kubernetes TokenRequest API to obtain a token with this audience. Note: The default audience works for most use cases. Only customize if required by your specific Databricks federation policy configuration. + * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). **Important:** For production deployments, especially when using multiple Kubernetes clusters, it is recommended to use a unique audience per cluster/environment (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) to allow separate Databricks federation policies and proper access control. The default generic audience is only suitable for single-cluster development setups. * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). * ``k8s_token_path``: (optional) path to the Kubernetes service account token file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). Override this for custom Kubernetes configurations. * ``k8s_namespace_path``: (optional) path to the Kubernetes namespace file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/namespace``). Override this for custom Kubernetes configurations. @@ -120,6 +120,28 @@ Extra (optional) * Set ``Host`` to your Databricks workspace URL * Add optional parameters as necessary in ``Extra`` field, e.g. ``{"audience": "custom-audience", "expiration_seconds": 7200, "k8s_token_path": "/custom/path/token", "k8s_namespace_path": "/custom/path/namespace"}`` + **Best Practices for Multi-Cluster Environments:** + + If you operate multiple Kubernetes clusters (e.g., production, staging, development) accessing the same Databricks workspace, it is recommended to use unique audiences for each cluster: + + * **Production:** ``{"audience": "databricks-prod-airflow"}`` + * **Staging:** ``{"audience": "databricks-staging-airflow"}`` + * **Development:** ``{"audience": "databricks-dev-airflow"}`` + + Or include cluster identifiers for even better isolation: + + * ``databricks-prod-eks-us-west-2-airflow`` + * ``databricks-staging-aks-westeurope-airflow`` + + **Why unique audiences are important:** + + * **Separate policies:** Create different Databricks federation policies for each cluster with appropriate permissions + * **Selective revocation:** Revoke access for one cluster without affecting others + * **Audit trail:** Identify which cluster made specific API calls + * **Security isolation:** Prevent dev/staging clusters from accessing production resources + + Each cluster requires its own federation policy in Databricks with its corresponding audience value. + **Important Notes:** * This method requires no secrets to be stored in the Airflow connection @@ -175,9 +197,9 @@ Extra (optional) Federation policy configuration: - * **Issuer:** ``https://kubernetes.default.svc`` - * **Audience:** ``https://kubernetes.default.svc`` (default, customizable via connection extra) - * **Subject:** ``system:serviceaccount::`` + * **Issuer:** (Required) ``https://kubernetes.default.svc`` + * **Audience:** (Required) ``https://kubernetes.default.svc`` (if using default). Recommended to use unique audience per cluster (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) + * **Subject:** (Required) ``system:serviceaccount::`` * **JWKS JSON:** - (Required) Your cluster's public signing keys Complete Databricks federation policy with JWKS JSON: From 60b00011d6f6e7b6b083737a6de6054a2abd7a1c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 09:35:39 +0100 Subject: [PATCH 07/21] refactor --- .../databricks/hooks/databricks_base.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 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 9c33f739f455e..e5dc878b8d93f 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -78,6 +78,13 @@ DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token" DEFAULT_K8S_NAMESPACE_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/namespace" +# RFC 8693 token exchange data template +TOKEN_EXCHANGE_DATA = { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "scope": "all-apis", +} + class BaseDatabricksHook(BaseHook): """ @@ -681,12 +688,7 @@ def _get_federated_databricks_token(self, resource: str) -> str: # Prepare token exchange request following RFC 8693 token_exchange_url = resource - data = { - "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", - "subject_token": jwt_token, - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - "scope": "all-apis", - } + data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token} try: for attempt in self._get_retry_object(): @@ -731,12 +733,7 @@ async def _a_get_federated_databricks_token(self, resource: str) -> str: # Prepare token exchange request following RFC 8693 token_exchange_url = resource - data = { - "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", - "subject_token": jwt_token, - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - "scope": "all-apis", - } + data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token} try: async for attempt in self._a_get_retry_object(): From fe6445aa6906942ee3b4b6008d33d9b9d99fcecd Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 12:04:55 +0100 Subject: [PATCH 08/21] Added support for projected volume token paths --- .../docs/connections/databricks.rst | 76 ++++++++++++++-- .../databricks/hooks/databricks_base.py | 86 ++++++++++++++++++- 2 files changed, 154 insertions(+), 8 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index a5b0825666f84..e281cc63d2943 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -101,14 +101,61 @@ Extra (optional) The following parameters are necessary if using authentication with Kubernetes OIDC token federation: - * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as extra parameter. When enabled, the hook will fetch a JWT token from the Kubernetes Service Account TokenRequest API 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). + * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as extra parameter. 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). + + **Two methods are supported for obtaining the Kubernetes JWT token:** + + **Method 1: Projected Volume** + + * ``k8s_projected_volume_token_path``: (optional) path to a Kubernetes projected volume service account token. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. + + **Method 2: TokenRequest API** + + If ``k8s_projected_volume_token_path`` is not configured, the hook will use the TokenRequest API method (dynamic token generation): + * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). **Important:** For production deployments, especially when using multiple Kubernetes clusters, it is recommended to use a unique audience per cluster/environment (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) to allow separate Databricks federation policies and proper access control. The default generic audience is only suitable for single-cluster development setups. * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). - * ``k8s_token_path``: (optional) path to the Kubernetes service account token file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). Override this for custom Kubernetes configurations. + * ``k8s_token_path``: (optional) path to the Kubernetes service account token file used to authenticate to the TokenRequest API (default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). Override this for custom Kubernetes configurations. * ``k8s_namespace_path``: (optional) path to the Kubernetes namespace file (default: ``/var/run/secrets/kubernetes.io/serviceaccount/namespace``). Override this for custom Kubernetes configurations. **Example configuration for Kubernetes OIDC token federation:** + **Option A: Using Projected Volume** + + 1. Configure your Pod spec with a projected volume: + + .. code-block:: yaml + + apiVersion: v1 + kind: Pod + metadata: + name: airflow-worker + spec: + serviceAccountName: airflow-worker + containers: + - name: airflow + image: apache/airflow:latest + volumeMounts: + - name: databricks-token + mountPath: /var/run/secrets/databricks + readOnly: true + volumes: + - name: databricks-token + projected: + sources: + - serviceAccountToken: + path: token + expirationSeconds: 3600 + audience: databricks-prod-airflow + + 2. Configure the Airflow connection: + + * Set ``Login`` to ``federated_k8s`` + * Set ``Host`` to your Databricks workspace URL + * Set ``Extra`` to your Pod spec volume mount path: ``{"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"}`` + + **Option B: Using TokenRequest API** + Minimal configuration (uses defaults): * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. @@ -142,10 +189,26 @@ Extra (optional) Each cluster requires its own federation policy in Databricks with its corresponding audience value. + **Choosing Between Methods:** + + * **Projected Volume (Method 1)**: Recommended when you control your Pod specs. Benefits include: + + - No Kubernetes API permissions needed (no TokenRequest API calls) + - Better performance (simple file read vs API call) + - Explicit configuration in Pod spec (clearer what audience/expiration is used) + - More Kubernetes-native approach + + * **TokenRequest API (Method 2)**: Use when: + + - You cannot modify Pod specs + - You need dynamic audience/expiration configuration + - You're working with existing deployments that use default service accounts + **Important Notes:** * This method requires no secrets to be stored in the Airflow connection - * The Kubernetes service account must be configured with appropriate permissions to create TokenRequest resources (typically granted by default in most Kubernetes clusters) + * For TokenRequest API method: The Kubernetes service account must be configured with appropriate permissions to create TokenRequest resources (typically granted by default in most Kubernetes clusters) + * For Projected Volume method: No special Kubernetes permissions needed, just standard service account token projection * The Databricks workspace must have federation policies configured in Databricks Account for the Kubernetes identity provider. This can be Account-level or Service principal-level policies. * Both Kubernetes JWT and Databricks OAuth tokens are short-lived and automatically refreshed. @@ -292,8 +355,11 @@ For example: export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://@host-url?token=yourtoken' - # For Kubernetes OIDC token federation (minimal): + # For Kubernetes OIDC token federation with projected volume (recommended): + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?k8s_projected_volume_token_path=%2Fvar%2Frun%2Fsecrets%2Fdatabricks%2Ftoken' + + # For Kubernetes OIDC token federation with TokenRequest API (minimal): export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com' - # For Kubernetes OIDC token federation with custom audience and expiration: + # For Kubernetes OIDC token federation with TokenRequest API and custom audience: export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?audience=custom-audience&expiration_seconds=7200' 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 e5dc878b8d93f..7299a73e65006 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -117,6 +117,7 @@ class BaseDatabricksHook(BaseHook): "federated_k8s", "k8s_token_path", "k8s_namespace_path", + "k8s_projected_volume_token_path", ] def __init__( @@ -544,7 +545,86 @@ def _is_oauth_token_valid(token: dict, time_key="expires_on") -> bool: def _get_k8s_jwt_token(self) -> str: """ - Get JWT token from Kubernetes Service Account Token service using TokenRequest API. + Get JWT token from Kubernetes. + + Supports two methods: + 1. Projected volume: reads token directly from configured path + 2. TokenRequest API: dynamically requests token from K8s API + + :return: JWT Service Account token string + """ + if "k8s_projected_volume_token_path" in self.databricks_conn.extra_dejson: + self.log.info("Using Kubernetes projected volume token") + return self._get_k8s_projected_volume_token() + + self.log.info("Using Kubernetes TokenRequest API") + return self._get_k8s_token_request_api() + + async def _a_get_k8s_jwt_token(self) -> str: + """Async version of _get_k8s_jwt_token().""" + if "k8s_projected_volume_token_path" in self.databricks_conn.extra_dejson: + self.log.info("Using Kubernetes projected volume token") + return await self._a_get_k8s_projected_volume_token() + + self.log.info("Using Kubernetes TokenRequest API") + return await self._a_get_k8s_token_request_api() + + def _get_k8s_projected_volume_token(self) -> str: + """ + Get JWT token from Kubernetes projected volume. + + Reads a pre-configured service account token from a projected volume. + The token should be configured in the Pod spec with the desired audience + and expiration settings. + + :return: JWT Service Account token string + """ + projected_token_path = self.databricks_conn.extra_dejson.get("k8s_projected_volume_token_path") + + try: + with open(projected_token_path) as f: + token = f.read().strip() + + if not token: + raise AirflowException(f"Token file at {projected_token_path} is empty") + + self.log.debug("Successfully read token from projected volume at %s", projected_token_path) + return token + except FileNotFoundError as e: + raise AirflowException( + f"Kubernetes projected volume token not found at {projected_token_path}. " + "Ensure your Pod has a projected volume configured with serviceAccountToken." + ) from e + except PermissionError as e: + raise AirflowException(f"Permission denied reading token from {projected_token_path}") from e + + async def _a_get_k8s_projected_volume_token(self) -> str: + """Async version of _get_k8s_projected_volume_token().""" + projected_token_path = self.databricks_conn.extra_dejson.get("k8s_projected_volume_token_path") + + try: + async with aiofiles.open(projected_token_path) as f: + token = (await f.read()).strip() + + if not token: + raise AirflowException(f"Token file at {projected_token_path} is empty") + + self.log.debug("Successfully read token from projected volume at %s", projected_token_path) + return token + except FileNotFoundError as e: + raise AirflowException( + f"Kubernetes projected volume token not found at {projected_token_path}. " + "Ensure your Pod has a projected volume configured with serviceAccountToken." + ) from e + except PermissionError as e: + raise AirflowException(f"Permission denied reading token from {projected_token_path}") from e + + def _get_k8s_token_request_api(self) -> str: + """ + Get JWT token using Kubernetes TokenRequest API. + + Dynamically requests a service account token from the Kubernetes API server + with custom audience and expiration settings. :return: JWT Service Account token string """ @@ -605,8 +685,8 @@ def _get_k8s_jwt_token(self) -> str: raise RuntimeError("Failed to get JWT token") - async def _a_get_k8s_jwt_token(self) -> str: - """Async version of _get_k8s_jwt_token().""" + async def _a_get_k8s_token_request_api(self) -> str: + """Async version of _get_k8s_token_request_api().""" audience = self.databricks_conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) expiration_seconds = self.databricks_conn.extra_dejson.get("expiration_seconds", 3600) token_path = self.databricks_conn.extra_dejson.get( From 17e42a9cf389dca3ca05698e2d359f394c038c7f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 12:09:09 +0100 Subject: [PATCH 09/21] refactor --- .../providers/databricks/hooks/databricks_base.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 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 7299a73e65006..7bf98cff8c38a 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -579,7 +579,12 @@ def _get_k8s_projected_volume_token(self) -> str: :return: JWT Service Account token string """ - projected_token_path = self.databricks_conn.extra_dejson.get("k8s_projected_volume_token_path") + projected_token_path: str | None = self.databricks_conn.extra_dejson.get( + "k8s_projected_volume_token_path" + ) + + if not projected_token_path: + raise AirflowException("k8s_projected_volume_token_path is not configured in connection extras") try: with open(projected_token_path) as f: @@ -600,7 +605,12 @@ def _get_k8s_projected_volume_token(self) -> str: async def _a_get_k8s_projected_volume_token(self) -> str: """Async version of _get_k8s_projected_volume_token().""" - projected_token_path = self.databricks_conn.extra_dejson.get("k8s_projected_volume_token_path") + projected_token_path: str | None = self.databricks_conn.extra_dejson.get( + "k8s_projected_volume_token_path" + ) + + if not projected_token_path: + raise AirflowException("k8s_projected_volume_token_path is not configured in connection extras") try: async with aiofiles.open(projected_token_path) as f: From 25b1d974c6581892d2d75da7478ffd8d971a56ed Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 12:19:43 +0100 Subject: [PATCH 10/21] added unit tests for the projected volume path method --- .../databricks/hooks/test_databricks_base.py | 312 ++++++++++++++++++ 1 file changed, 312 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 7ba585b921ce9..29033b35d42cc 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -1121,6 +1121,140 @@ def test_get_k8s_jwt_token_with_custom_extra_settings(self, mock_post, mock_open # Verify Authorization header uses custom token assert call_args[1]["headers"]["Authorization"] == "Bearer custom_in_cluster_token" + @mock.patch("builtins.open", mock.mock_open(read_data="projected_token_content")) + def test_get_k8s_projected_volume_token_success(self): + """Test successfully reading token from Kubernetes projected volume.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + token = hook._get_k8s_projected_volume_token() + + assert token == "projected_token_content" + + @mock.patch("builtins.open", side_effect=FileNotFoundError()) + def test_get_k8s_projected_volume_token_file_not_found(self, mock_open): + """Test error when projected volume token file is not found.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with pytest.raises( + AirflowException, + match="Kubernetes projected volume token not found at /var/run/secrets/databricks/token", + ): + hook._get_k8s_projected_volume_token() + + @mock.patch("builtins.open", side_effect=PermissionError()) + def test_get_k8s_projected_volume_token_permission_denied(self, mock_open): + """Test error when permission denied reading projected volume token.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with pytest.raises( + AirflowException, match="Permission denied reading token from /var/run/secrets/databricks/token" + ): + hook._get_k8s_projected_volume_token() + + @mock.patch("builtins.open", mock.mock_open(read_data="")) + def test_get_k8s_projected_volume_token_empty_file(self): + """Test error when projected volume token file is empty.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with pytest.raises( + AirflowException, match="Token file at /var/run/secrets/databricks/token is empty" + ): + hook._get_k8s_projected_volume_token() + + def test_get_k8s_jwt_token_uses_projected_volume_when_configured(self): + """Test that _get_k8s_jwt_token delegates to projected volume method when configured.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch.object( + hook, "_get_k8s_projected_volume_token", return_value="projected_volume_token" + ) as mock_projected: + with mock.patch.object(hook, "_get_k8s_token_request_api") as mock_token_request: + token = hook._get_k8s_jwt_token() + + assert token == "projected_volume_token" + mock_projected.assert_called_once() + mock_token_request.assert_not_called() + + def test_get_k8s_jwt_token_uses_token_request_api_when_no_projected_path(self): + """Test that _get_k8s_jwt_token delegates to TokenRequest API when no projected path configured.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {} # No projected volume path + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch.object(hook, "_get_k8s_projected_volume_token") as mock_projected: + with mock.patch.object( + hook, "_get_k8s_token_request_api", return_value="token_request_api_token" + ) as mock_token_request: + token = hook._get_k8s_jwt_token() + + assert token == "token_request_api_token" + mock_projected.assert_not_called() + mock_token_request.assert_called_once() + + @mock.patch("builtins.open", mock.mock_open(read_data="projected_k8s_jwt_token")) + @mock.patch("requests.post") + @time_machine.travel("2025-07-12 12:00:00") + def test_get_federated_token_with_projected_volume(self, mock_post): + """Test end-to-end federated token flow using projected volume.""" + # Mock Databricks token exchange response + expiry_date = int((datetime(2025, 7, 12, 12, 0, 0) + timedelta(minutes=60)).timestamp()) + db_response = mock.Mock() + 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 + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + resource = f"https://{mock_conn.host}/oidc/v1/token" + token = hook._get_federated_databricks_token(resource) + + assert token == "databricks_token" + assert resource in hook.oauth_tokens + assert hook.oauth_tokens[resource]["access_token"] == "databricks_token" + assert hook.oauth_tokens[resource]["expires_on"] == expiry_date + + # Verify Databricks token exchange was called + assert mock_post.call_count == 1 + db_call_args = mock_post.call_args + assert "https://my-workspace.cloud.databricks.com/oidc/v1/token" in db_call_args[0] + assert db_call_args[1]["data"]["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert db_call_args[1]["data"]["subject_token"] == "projected_k8s_jwt_token" + assert db_call_args[1]["data"]["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" + assert db_call_args[1]["data"]["scope"] == "all-apis" + @mock.patch( "builtins.open", side_effect=[ @@ -1361,6 +1495,184 @@ async def test_a_get_federated_token_databricks_error(self, mock_post): ): await hook._a_get_federated_databricks_token(resource) + @pytest.mark.asyncio + async def test_a_get_k8s_projected_volume_token_success(self): + """Test async successfully reading token from Kubernetes projected volume.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + # Mock aiofiles.open + mock_file = mock.AsyncMock() + mock_file.read = mock.AsyncMock(return_value="projected_token_content") + + with mock.patch( + "aiofiles.open", return_value=mock.MagicMock(__aenter__=mock.AsyncMock(return_value=mock_file)) + ): + token = await hook._a_get_k8s_projected_volume_token() + + assert token == "projected_token_content" + + @pytest.mark.asyncio + async def test_a_get_k8s_projected_volume_token_file_not_found(self): + """Test async error when projected volume token file is not found.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch("aiofiles.open", side_effect=FileNotFoundError()): + with pytest.raises( + AirflowException, + match="Kubernetes projected volume token not found at /var/run/secrets/databricks/token", + ): + await hook._a_get_k8s_projected_volume_token() + + @pytest.mark.asyncio + async def test_a_get_k8s_projected_volume_token_permission_denied(self): + """Test async error when permission denied reading projected volume token.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch("aiofiles.open", side_effect=PermissionError()): + with pytest.raises( + AirflowException, + match="Permission denied reading token from /var/run/secrets/databricks/token", + ): + await hook._a_get_k8s_projected_volume_token() + + @pytest.mark.asyncio + async def test_a_get_k8s_projected_volume_token_empty_file(self): + """Test async error when projected volume token file is empty.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + # Mock aiofiles.open with empty content + mock_file = mock.AsyncMock() + mock_file.read = mock.AsyncMock(return_value="") + + with mock.patch( + "aiofiles.open", return_value=mock.MagicMock(__aenter__=mock.AsyncMock(return_value=mock_file)) + ): + with pytest.raises( + AirflowException, match="Token file at /var/run/secrets/databricks/token is empty" + ): + await hook._a_get_k8s_projected_volume_token() + + @pytest.mark.asyncio + async def test_a_get_k8s_jwt_token_uses_projected_volume_when_configured(self): + """Test that async _a_get_k8s_jwt_token delegates to projected volume method when configured.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch.object( + hook, + "_a_get_k8s_projected_volume_token", + new_callable=mock.AsyncMock, + return_value="projected_volume_token", + ) as mock_projected: + with mock.patch.object( + hook, "_a_get_k8s_token_request_api", new_callable=mock.AsyncMock + ) as mock_token_request: + token = await hook._a_get_k8s_jwt_token() + + assert token == "projected_volume_token" + mock_projected.assert_called_once() + mock_token_request.assert_not_called() + + @pytest.mark.asyncio + async def test_a_get_k8s_jwt_token_uses_token_request_api_when_no_projected_path(self): + """Test that async _a_get_k8s_jwt_token delegates to TokenRequest API when no projected path.""" + mock_conn = mock.Mock() + mock_conn.extra_dejson = {} # No projected volume path + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + with mock.patch.object( + hook, "_a_get_k8s_projected_volume_token", new_callable=mock.AsyncMock + ) as mock_projected: + with mock.patch.object( + hook, + "_a_get_k8s_token_request_api", + new_callable=mock.AsyncMock, + return_value="token_request_api_token", + ) as mock_token_request: + token = await hook._a_get_k8s_jwt_token() + + assert token == "token_request_api_token" + mock_projected.assert_not_called() + mock_token_request.assert_called_once() + + @pytest.mark.asyncio + @time_machine.travel("2025-07-12 12:00:00") + async def test_a_get_federated_token_with_projected_volume(self): + """Test async end-to-end federated token flow using projected volume.""" + # Mock Databricks token exchange response + expiry_date = int((datetime(2025, 7, 12, 12, 0, 0) + timedelta(minutes=60)).timestamp()) + db_response_data = { + "access_token": "databricks_token", + "expires_in": 3600, + "token_type": "Bearer", + } + + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + hook.user_agent_header = {"User-Agent": "test-agent"} + + # Mock aiofiles.open for projected volume + mock_file = mock.AsyncMock() + mock_file.read = mock.AsyncMock(return_value="projected_k8s_jwt_token") + + # Mock aiohttp session + mock_response = mock.AsyncMock() + mock_response.json = mock.AsyncMock(return_value=db_response_data) + mock_response.raise_for_status = mock.Mock() + + mock_session = mock.AsyncMock() + mock_session.post = mock.MagicMock( + return_value=mock.MagicMock(__aenter__=mock.AsyncMock(return_value=mock_response)) + ) + + async with hook: + hook._session = mock_session + + with mock.patch( + "aiofiles.open", + return_value=mock.MagicMock(__aenter__=mock.AsyncMock(return_value=mock_file)), + ): + resource = f"https://{mock_conn.host}/oidc/v1/token" + token = await hook._a_get_federated_databricks_token(resource) + + assert token == "databricks_token" + assert resource in hook.oauth_tokens + assert hook.oauth_tokens[resource]["access_token"] == "databricks_token" + assert hook.oauth_tokens[resource]["expires_on"] == expiry_date + + # Verify Databricks token exchange was called + mock_session.post.assert_called_once() + call_args = mock_session.post.call_args + assert resource in call_args[0] + assert call_args[1]["data"]["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert call_args[1]["data"]["subject_token"] == "projected_k8s_jwt_token" + @pytest.mark.asyncio @mock.patch("aiohttp.ClientSession.post") async def test_a_get_token_with_federated_k8s_login(self, mock_post): From aa6723d86409def3b309ccbc3c2bc76d0f58c57f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 14:25:10 +0100 Subject: [PATCH 11/21] updated docs --- providers/databricks/docs/connections/databricks.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index e281cc63d2943..1cb2ea7cf4aca 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -107,11 +107,11 @@ Extra (optional) **Method 1: Projected Volume** - * ``k8s_projected_volume_token_path``: (optional) path to a Kubernetes projected volume service account token. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. + * ``k8s_projected_volume_token_path``: (optional) path to a [Kubernetes projected volume service account token](https://kubernetes.io/docs/concepts/configuration/secret/#projected-volume). When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. **Method 2: TokenRequest API** - If ``k8s_projected_volume_token_path`` is not configured, the hook will use the TokenRequest API method (dynamic token generation): + If ``k8s_projected_volume_token_path`` is not configured, the hook will use the [TokenRequest API](https://kubernetes.io/docs/reference/kubernetes-api/authentication-resources/token-request-v1/) method (dynamic token generation): * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). **Important:** For production deployments, especially when using multiple Kubernetes clusters, it is recommended to use a unique audience per cluster/environment (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) to allow separate Databricks federation policies and proper access control. The default generic audience is only suitable for single-cluster development setups. * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). From 0958e26bbf7d20fdd96a7e61816aa69bbd56ed08 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 4 Feb 2026 14:32:40 +0100 Subject: [PATCH 12/21] updated docs --- providers/databricks/docs/connections/databricks.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 1cb2ea7cf4aca..1f44fa492f7a1 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -45,8 +45,9 @@ There are several ways to connect to Databricks using Airflow. i.e. use OAuth authentication with Databricks-managed Service Principal client ID and secret. See `Authentication using OAuth for service principals `_. 6. Using Kubernetes `OIDC token federation `_ (applicable only when Airflow runs in Kubernetes) - i.e. automatically fetch JWT tokens from Kubernetes Service Account 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. + 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). Default Connection IDs ---------------------- From f064e4b6cf227393d25216add8982c42db853cf6 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 5 Feb 2026 19:46:38 +0100 Subject: [PATCH 13/21] * Added exception handling for json response * Corrected links in docs --- .../docs/connections/databricks.rst | 4 +-- .../databricks/hooks/databricks_base.py | 33 +++++++++++++++++-- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 1f44fa492f7a1..d49e265861fda 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -108,11 +108,11 @@ Extra (optional) **Method 1: Projected Volume** - * ``k8s_projected_volume_token_path``: (optional) path to a [Kubernetes projected volume service account token](https://kubernetes.io/docs/concepts/configuration/secret/#projected-volume). When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. + * ``k8s_projected_volume_token_path``: (optional) path to a `Kubernetes projected volume service account token `_. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. **Method 2: TokenRequest API** - If ``k8s_projected_volume_token_path`` is not configured, the hook will use the [TokenRequest API](https://kubernetes.io/docs/reference/kubernetes-api/authentication-resources/token-request-v1/) method (dynamic token generation): + If ``k8s_projected_volume_token_path`` is not configured, the hook will use the `TokenRequest API `_ method (dynamic token generation): * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). **Important:** For production deployments, especially when using multiple Kubernetes clusters, it is recommended to use a unique audience per cluster/environment (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) to allow separate Databricks federation policies and proper access control. The default generic audience is only suitable for single-cluster development setups. * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). 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 7bf98cff8c38a..fbbfffd49869d 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -679,7 +679,18 @@ def _get_k8s_token_request_api(self) -> str: timeout=self.token_timeout_seconds, ) resp.raise_for_status() - return resp.json()["status"]["token"] + try: + response_json = resp.json() + return response_json["status"]["token"] + except (JSONDecodeError, ValueError) as e: + raise AirflowException( + f"Invalid JSON response from Kubernetes API. Response: {resp.text[:500]}" + ) from e + except KeyError as e: + raise AirflowException( + f"Malformed Kubernetes token response: missing key '{e}'. " + f"Response: {resp.text[:500]}" + ) from e except FileNotFoundError as e: raise AirflowException( "Kubernetes service account token not found. " @@ -738,8 +749,24 @@ async def _a_get_k8s_token_request_api(self) -> str: timeout=self.token_timeout_seconds, ) as resp: resp.raise_for_status() - jsn = await resp.json() - return jsn["status"]["token"] + try: + jsn = await resp.json() + return jsn["status"]["token"] + except (aiohttp.ContentTypeError, ValueError) as e: + # Try to read response text if JSON parsing failed + try: + response_text = await resp.text() + except Exception: + response_text = "unable to read response" + raise AirflowException( + f"Invalid JSON response from Kubernetes API. Response: {response_text[:500]}" + ) from e + except KeyError as e: + # Response body already consumed, use the parsed JSON for error message + raise AirflowException( + f"Malformed Kubernetes token response: missing key '{e}'. " + f"Response structure: {jsn}" + ) from e except FileNotFoundError as e: raise AirflowException( "Kubernetes service account token not found. " From 87518a03d9deec4af853cf760c596c1bcc447199 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 6 Feb 2026 09:51:34 +0100 Subject: [PATCH 14/21] * Updated docs * Minor code refactor --- .../databricks/docs/connections/databricks.rst | 6 +++--- .../providers/databricks/hooks/databricks_base.py | 14 ++------------ 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index d49e265861fda..1d97abb934b22 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -41,7 +41,7 @@ There are several ways to connect to Databricks using Airflow. `user inside workspace `_, or `outside of workspace having Owner or Contributor permissions `_ 4. Using Azure Active Directory (AAD) token obtained for `Azure managed identity `_, when Airflow runs on the VM with assigned managed identity (system-assigned or user-assigned) -5. Using Databricks-managed Service Principal OAuth +5. Using Databricks-managed Service Principal OAuth (available on all supported clouds) i.e. use OAuth authentication with Databricks-managed Service Principal client ID and secret. See `Authentication using OAuth for service principals `_. 6. Using Kubernetes `OIDC token federation `_ (applicable only when Airflow runs in Kubernetes) @@ -102,7 +102,7 @@ Extra (optional) The following parameters are necessary if using authentication with Kubernetes OIDC token federation: - * ``federated_k8s``: set ``login`` to ``"federated_k8s"`` or add this as extra parameter. 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). + * ``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). **Two methods are supported for obtaining the Kubernetes JWT token:** @@ -164,7 +164,7 @@ Extra (optional) With custom configuration: - * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. + * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter (``{"federated_k8s": true}``). * Set ``Host`` to your Databricks workspace URL * Add optional parameters as necessary in ``Extra`` field, e.g. ``{"audience": "custom-audience", "expiration_seconds": 7200, "k8s_token_path": "/custom/path/token", "k8s_namespace_path": "/custom/path/namespace"}`` 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 fbbfffd49869d..3090a3be6982d 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -579,12 +579,7 @@ def _get_k8s_projected_volume_token(self) -> str: :return: JWT Service Account token string """ - projected_token_path: str | None = self.databricks_conn.extra_dejson.get( - "k8s_projected_volume_token_path" - ) - - if not projected_token_path: - raise AirflowException("k8s_projected_volume_token_path is not configured in connection extras") + projected_token_path: str = self.databricks_conn.extra_dejson["k8s_projected_volume_token_path"] try: with open(projected_token_path) as f: @@ -605,12 +600,7 @@ def _get_k8s_projected_volume_token(self) -> str: async def _a_get_k8s_projected_volume_token(self) -> str: """Async version of _get_k8s_projected_volume_token().""" - projected_token_path: str | None = self.databricks_conn.extra_dejson.get( - "k8s_projected_volume_token_path" - ) - - if not projected_token_path: - raise AirflowException("k8s_projected_volume_token_path is not configured in connection extras") + projected_token_path: str = self.databricks_conn.extra_dejson["k8s_projected_volume_token_path"] try: async with aiofiles.open(projected_token_path) as f: From 434612d20d04f70254637c687df32e3c76f3e189 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 6 Feb 2026 09:56:43 +0100 Subject: [PATCH 15/21] * Minor code refactor --- .../databricks/hooks/databricks_base.py | 36 ++++++++++--------- 1 file changed, 20 insertions(+), 16 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 3090a3be6982d..2c5df402bb74a 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -619,6 +619,24 @@ async def _a_get_k8s_projected_volume_token(self) -> str: except PermissionError as e: raise AirflowException(f"Permission denied reading token from {projected_token_path}") from e + @staticmethod + def _build_k8s_token_request_payload(audience: str, expiration_seconds: int) -> dict[str, Any]: + """ + Build the JSON payload for Kubernetes TokenRequest API. + + :param audience: The audience value for the JWT token + :param expiration_seconds: Token expiration in seconds + :return: TokenRequest API payload dictionary + """ + return { + "apiVersion": "authentication.k8s.io/v1", + "kind": "TokenRequest", + "spec": { + "audiences": [audience], + "expirationSeconds": expiration_seconds, + }, + } + def _get_k8s_token_request_api(self) -> str: """ Get JWT token using Kubernetes TokenRequest API. @@ -657,14 +675,7 @@ def _get_k8s_token_request_api(self) -> str: "Authorization": f"Bearer {in_cluster_token}", "Content-Type": "application/json", }, - json={ - "apiVersion": "authentication.k8s.io/v1", - "kind": "TokenRequest", - "spec": { - "audiences": [audience], - "expirationSeconds": expiration_seconds, - }, - }, + json=self._build_k8s_token_request_payload(audience, expiration_seconds), verify=False, # K8s in-cluster uses self-signed certs timeout=self.token_timeout_seconds, ) @@ -727,14 +738,7 @@ async def _a_get_k8s_token_request_api(self) -> str: "Authorization": f"Bearer {in_cluster_token}", "Content-Type": "application/json", }, - json={ - "apiVersion": "authentication.k8s.io/v1", - "kind": "TokenRequest", - "spec": { - "audiences": [audience], - "expirationSeconds": expiration_seconds, - }, - }, + json=self._build_k8s_token_request_payload(audience, expiration_seconds), ssl=False, # K8s in-cluster uses self-signed certs timeout=self.token_timeout_seconds, ) as resp: From 6c17e5426d3d78a1d6339babb4f19e506b75780c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 12 Feb 2026 16:43:30 +0100 Subject: [PATCH 16/21] Added client id to support service-principal token federation. Remove option to use account-wide token federation. --- .../docs/connections/databricks.rst | 80 +++++++++++++------ .../databricks/hooks/databricks_base.py | 50 +++++++++--- 2 files changed, 96 insertions(+), 34 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index 1d97abb934b22..c79ac9268d4d6 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -103,6 +103,7 @@ Extra (optional) 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). + * ``client_id``: (required) Databricks service principal client UUID. The service principal must exist in your Databricks account and be assigned to the workspace specified in the Host field. **Two methods are supported for obtaining the Kubernetes JWT token:** @@ -153,20 +154,27 @@ Extra (optional) * Set ``Login`` to ``federated_k8s`` * Set ``Host`` to your Databricks workspace URL - * Set ``Extra`` to your Pod spec volume mount path: ``{"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"}`` + * Set ``Extra`` to your Pod spec volume mount path: ``{"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", "client_id": "your-service-principal-client-uuid"}`` **Option B: Using TokenRequest API** Minimal configuration (uses defaults): - * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. + **Option 1: Using Login field** + * Set ``Login`` to ``federated_k8s`` + * Set ``Host`` to your Databricks workspace URL + * Add ``client_id`` in ``Extra`` field, e.g. ``{"client_id": "your-service-principal-client-uuid"}`` + + **Option 2: Using Extra field only** + * Leave ``Login`` empty * Set ``Host`` to your Databricks workspace URL + * Add both ``federated_k8s`` and ``client_id`` in ``Extra`` field, e.g. ``{"federated_k8s": true, "client_id": "your-service-principal-client-uuid"}`` With custom configuration: - * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter (``{"federated_k8s": true}``). + * Set ``Login`` to ``federated_k8s`` or leave empty and add this as extra parameter. * Set ``Host`` to your Databricks workspace URL - * Add optional parameters as necessary in ``Extra`` field, e.g. ``{"audience": "custom-audience", "expiration_seconds": 7200, "k8s_token_path": "/custom/path/token", "k8s_namespace_path": "/custom/path/namespace"}`` + * Add required and optional parameters in ``Extra`` field, e.g. ``{"client_id": "your-service-principal-client-uuid", "audience": "custom-audience", "expiration_seconds": 7200, "k8s_token_path": "/custom/path/token", "k8s_namespace_path": "/custom/path/namespace"}`` **Best Practices for Multi-Cluster Environments:** @@ -208,17 +216,19 @@ Extra (optional) **Important Notes:** * This method requires no secrets to be stored in the Airflow connection - * For TokenRequest API method: The Kubernetes service account must be configured with appropriate permissions to create TokenRequest resources (typically granted by default in most Kubernetes clusters) - * For Projected Volume method: No special Kubernetes permissions needed, just standard service account token projection - * The Databricks workspace must have federation policies configured in Databricks Account for the Kubernetes identity provider. This can be Account-level or Service principal-level policies. + * For TokenRequest API method: The TokenRequest API is called to request a token for the ``default`` service account, regardless of which service account the pod is running with. The pod's service account must have appropriate RBAC permissions to create TokenRequest resources for the ``default`` service account. + * For Projected Volume method: No special Kubernetes permissions needed, just standard service account token projection. + * The Databricks workspace must have federation policy configured in Databricks Account for the Kubernetes identity provider. **Only service principal-level policies are supported** for Kubernetes OIDC token federation. + * ``client_id`` is required in the connection extra parameters. Service principal-level Databricks federation must be used because Kubernetes service account tokens do not support custom claims, which are required for account-wide federation. + * The federation policy must be configured for the specific service principal identified by the ``client_id``. Account-wide federation policies are not supported for Kubernetes OIDC token federation. * Both Kubernetes JWT and Databricks OAuth tokens are short-lived and automatically refreshed. **Databricks Federation Policy Configuration:** - Before using Kubernetes OIDC token federation, you must configure a federation policy in your Databricks account. The policy configuration is the same regardless of whether you create it at the account level or for a specific service principal. The difference is only where you create the policy: + Before using Kubernetes OIDC token federation, you must configure a federation policy in your Databricks account: - * **Account-level policy:** Configured in the global authentication settings. All users and service principals in the account can use this authentication method. - * **Service principal policy:** Configured in the service principal settings. Only the specific service principal can use this authentication method (recommended for workloads). + * **Service principal policy (required):** Configured in the service principal settings. Only the specific service principal can use this authentication method. The ``client_id`` value must be provided in the connection extra parameters and must match the service principal UUID for which the federation policy is configured. + * **Account-level policy:** Not supported for Kubernetes OIDC token federation. Before configuring the federation policy, you need to understand the token issuer and JWKS requirements. @@ -264,9 +274,12 @@ Extra (optional) * **Issuer:** (Required) ``https://kubernetes.default.svc`` * **Audience:** (Required) ``https://kubernetes.default.svc`` (if using default). Recommended to use unique audience per cluster (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) * **Subject:** (Required) ``system:serviceaccount::`` + + * **For Projected Volume method:** Use the actual service account name from your airflow worker pod (e.g., ``system:serviceaccount:airflow:my-sa``) + * **For TokenRequest API method:** Always use ``system:serviceaccount::default`` because the TokenRequest API is called to request a token for the ``default`` service account, regardless of which service account the airflow worker pod is running with. * **JWKS JSON:** - (Required) Your cluster's public signing keys - Complete Databricks federation policy with JWKS JSON: + Complete example of a Databricks federation policy with JWKS JSON: .. code-block:: json @@ -274,7 +287,7 @@ Extra (optional) "oidc_policy": { "issuer": "https://kubernetes.default.svc", "audiences": ["https://kubernetes.default.svc"], - "subject": "system:serviceaccount:airflow:airflow-worker", + "subject": "system:serviceaccount:airflow:default", "jwks_json": { "keys": [{ "kty": "RSA", @@ -296,39 +309,58 @@ Extra (optional) Copy the entire output and use it as the value for ``jwks_json`` in your Databricks federation policy. Each Kubernetes cluster has unique signing keys, which is what provides the security - only tokens signed by your specific cluster can be validated. + Example JWT token from standard Kubernetes matching the policy: + + .. code-block:: json + + { + "iss": "https://kubernetes.default.svc", + "aud": ["https://kubernetes.default.svc"], + "sub": "system:serviceaccount:airflow:default" + } + **Scenario 2: Managed Kubernetes with Public OIDC - JWKS JSON Optional** - If your cluster uses a public OIDC issuer (check using the method presented above), you can omit JWKS JSON: + Federation policy configuration: + + * **Issuer:** (Required) Your cluster's public OIDC issuer URL (e.g., ``https://oidc.eks.us-west-2.amazonaws.com/id/YOUR_CLUSTER_ID`` for EKS) + * **Audience:** (Required) The audience value used when requesting tokens. For standard Kubernetes, use ``https://kubernetes.default.svc`` (if using default) or a unique audience per cluster (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) + * **Subject:** (Required) ``system:serviceaccount::`` + + * **For Projected Volume method:** Use the actual service account name from your airflow worker pod (e.g., ``system:serviceaccount:airflow:my-sa``) + * **For TokenRequest API method:** Always use ``system:serviceaccount::default`` because the TokenRequest API is called to request a token for the ``default`` service account, regardless of which service account the airflow worker pod is running with. + + Complete example of a Databricks federation policy without JWKS JSON (recommended for public OIDC issuers): .. code-block:: json { "oidc_policy": { "issuer": "https://oidc.eks.us-west-2.amazonaws.com/id/YOUR_CLUSTER_ID", - "audiences": ["sts.amazonaws.com"], - "subject": "system:serviceaccount:airflow:airflow-worker" + "audiences": ["https://kubernetes.default.svc"], + "subject": "system:serviceaccount:airflow:default" } } In this case, Databricks will automatically fetch the public keys from ``/.well-known/openid-configuration``. - Example JWT token from standard Kubernetes: + **Note:** If your cluster's OIDC issuer is not publicly accessible or you prefer to provide JWKS JSON explicitly, you can include it in the same format as Scenario 1. + + Example JWT token from EKS matching the policy: .. code-block:: json { - "iss": "https://kubernetes.default.svc", + "iss": "https://oidc.eks.us-west-2.amazonaws.com/id/YOUR_CLUSTER_ID", "aud": ["https://kubernetes.default.svc"], - "sub": "system:serviceaccount:airflow:airflow-worker" + "sub": "system:serviceaccount:airflow:default" } - (where ``airflow`` is your Kubernetes namespace and ``airflow-worker`` is your service account name) - **Troubleshooting Common Issues:** * **"Kubernetes service account token not found" error:** This authentication method only works when Airflow is running inside a Kubernetes cluster. Ensure your pods have the service account token mounted (default behavior in Kubernetes). - * **Permission denied errors:** Verify your Kubernetes service account has permission to create TokenRequest resources. In most clusters, this is granted by default to all service accounts. + * **Permission denied errors:** For the TokenRequest API method, verify your pod's service account has RBAC permissions to create TokenRequest resources for the ``default`` service account. By default, service accounts can only create TokenRequest resources for themselves, not for other service accounts. You must explicitly grant these permissions via RBAC policies (ClusterRole and ClusterRoleBinding or Role and RoleBinding). * **Token exchange failures:** Ensure your Databricks federation policy configuration matches the JWT token properties: @@ -357,10 +389,10 @@ For example: export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://@host-url?token=yourtoken' # For Kubernetes OIDC token federation with projected volume (recommended): - export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?k8s_projected_volume_token_path=%2Fvar%2Frun%2Fsecrets%2Fdatabricks%2Ftoken' + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?k8s_projected_volume_token_path=%2Fvar%2Frun%2Fsecrets%2Fdatabricks%2Ftoken&client_id=your-service-principal-client-uuid' # For Kubernetes OIDC token federation with TokenRequest API (minimal): - export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com' + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?client_id=your-service-principal-client-uuid' # For Kubernetes OIDC token federation with TokenRequest API and custom audience: - export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?audience=custom-audience&expiration_seconds=7200' + export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://federated_k8s@my-workspace.cloud.databricks.com?client_id=your-service-principal-client-uuid&audience=custom-audience&expiration_seconds=7200' 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 2c5df402bb74a..12c1031b8a112 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -218,6 +218,7 @@ def _parse_host(host: str) -> str: assert h._parse_host('xx.cloud.databricks.com') == 'xx.cloud.databricks.com' """ + host = host.rstrip('/') # Remove trailing slashes urlparse_host = urlsplit(host).hostname if urlparse_host: # In this case, host = https://xx.cloud.databricks.com @@ -776,13 +777,30 @@ async def _a_get_k8s_token_request_api(self) -> str: ) raise RuntimeError("Failed to get JWT token") + def _get_required_client_id(self) -> str: + """ + Get and validate client_id for Kubernetes OIDC token federation. + + :return: Service principal client ID + :raises AirflowException: If client_id is not provided + """ + client_id = self.databricks_conn.extra_dejson.get("client_id", None) + if not client_id: + # see: https://github.com/kubernetes/kubernetes/issues/116638 + raise AirflowException( + "client_id is required for Kubernetes OIDC token federation. " + "Kubernetes service account tokens do not support custom claims, " + "so service principal-level federation must be used. " + "Please provide client_id in the connection extra parameters." + ) + return client_id + def _get_federated_databricks_token(self, resource: str) -> str: """ Get Databricks OAuth token by exchanging Kubernetes JWT token. Uses RFC 8693 token exchange to convert a Kubernetes service account JWT - into a Databricks OAuth token. Supports both account-level and service principal - federation policies. + into a Databricks OAuth token. Requires service principal-level federation. :param resource: Databricks OIDC token exchange URL :return: Databricks OAuth access token @@ -793,13 +811,15 @@ 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) # Prepare token exchange request following RFC 8693 token_exchange_url = resource - - data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token} + data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token, "client_id": client_id} try: for attempt in self._get_retry_object(): @@ -838,13 +858,15 @@ 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) # Prepare token exchange request following RFC 8693 token_exchange_url = resource - - data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token} + data = {**TOKEN_EXCHANGE_DATA, "subject_token": jwt_token, "client_id": client_id} try: async for attempt in self._a_get_retry_object(): @@ -960,13 +982,13 @@ def _get_token(self, raise_error: bool = False) -> str | None: if self.databricks_conn.login == "" or self.databricks_conn.password == "": raise AirflowException("Service Principal credentials aren't provided") self.log.debug("Using Service Principal Token.") - return self._get_sp_token(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) + return self._get_sp_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( - OIDC_TOKEN_SERVICE_URL.format(f"https://{self.databricks_conn.host}") + self._get_oidc_token_service_url() ) if raise_error: raise AirflowException("Token authentication isn't configured") @@ -999,13 +1021,13 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None: if self.databricks_conn.login == "" or self.databricks_conn.password == "": raise AirflowException("Service Principal credentials aren't provided") self.log.debug("Using Service Principal Token.") - return await self._a_get_sp_token(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) + return await self._a_get_sp_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( - OIDC_TOKEN_SERVICE_URL.format(f"https://{self.databricks_conn.host}") + self._get_oidc_token_service_url() ) if raise_error: raise AirflowException("Token authentication isn't configured") @@ -1015,6 +1037,14 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None: def _log_request_error(self, attempt_num: int, error: str) -> None: self.log.error("Attempt %s API Request to Databricks failed with reason: %s", attempt_num, error) + def _get_oidc_token_service_url(self) -> str: + """ + Construct the OIDC token service URL for Databricks. + + :return: Full URL to the OIDC token service endpoint + """ + return OIDC_TOKEN_SERVICE_URL.format(f"https://{self.host}") + def _endpoint_url(self, endpoint): port = f":{self.databricks_conn.port}" if self.databricks_conn.port else "" schema = self.databricks_conn.schema or "https" From c21647969fd9e8951a5c0b2ec8fe475e0b7e7498 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 12 Feb 2026 17:25:14 +0100 Subject: [PATCH 17/21] Updated docs with clear explanation on how to configure projected volume token path --- providers/databricks/docs/connections/databricks.rst | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index c79ac9268d4d6..aba11cea4eca2 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -109,7 +109,7 @@ Extra (optional) **Method 1: Projected Volume** - * ``k8s_projected_volume_token_path``: (optional) path to a `Kubernetes projected volume service account token `_. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. + * ``k8s_projected_volume_token_path``: (optional) path to a `Kubernetes projected volume service account token `_. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. **Important:** The ``k8s_projected_volume_token_path`` must match the full path constructed from your Pod spec: ``{mountPath}/{path}``, where ``mountPath`` is the volume mount path and ``path`` is the serviceAccountToken path field. For example, if your Pod spec has ``mountPath: /var/run/secrets/databricks`` and ``path: token``, then ``k8s_projected_volume_token_path`` should be ``/var/run/secrets/databricks/token``. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. **Method 2: TokenRequest API** @@ -154,7 +154,12 @@ Extra (optional) * Set ``Login`` to ``federated_k8s`` * Set ``Host`` to your Databricks workspace URL - * Set ``Extra`` to your Pod spec volume mount path: ``{"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", "client_id": "your-service-principal-client-uuid"}`` + * Set ``Extra`` to include the full token path matching your Pod spec: ``{"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", "client_id": "your-service-principal-client-uuid"}`` + + **Note:** The ``k8s_projected_volume_token_path`` value must match ``{mountPath}/{path}`` from your Pod spec. In the example above: + - ``mountPath`` is ``/var/run/secrets/databricks`` (from volumeMounts) + - ``path`` is ``token`` (from serviceAccountToken) + - Full path is ``/var/run/secrets/databricks/token`` **Option B: Using TokenRequest API** From 2de4f2d2da8514df0895308af7797d5a1b4226b9 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 12 Feb 2026 18:28:52 +0100 Subject: [PATCH 18/21] Updated unit tests --- .../unit/databricks/hooks/test_databricks.py | 2 +- .../databricks/hooks/test_databricks_base.py | 42 +++++++++++-------- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py index 3971a096fab39..b761972f613c9 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py @@ -2191,7 +2191,7 @@ def test_submit_run(self, mock_requests): run_id = self.hook.submit_run(data) ad_call_args = mock_requests.method_calls[0] - assert ad_call_args[1][0] == OIDC_TOKEN_SERVICE_URL.format(HOST) + assert ad_call_args[1][0] == OIDC_TOKEN_SERVICE_URL.format(f"https://{HOST}") assert ad_call_args[2]["data"] == "grant_type=client_credentials&scope=all-apis" assert run_id == "1" 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 29033b35d42cc..07b4aceb5563d 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -417,7 +417,7 @@ def test_get_token_service_principal_oauth_success(self, mock_get_sp_token, mock with mock.patch.object(hook.log, "debug") as mock_log_debug: token = hook._get_token() assert token == "spn_token" - mock_get_sp_token.assert_called_once_with("example.databricks.com/oidc/v1/token") + mock_get_sp_token.assert_called_once_with("https://example.databricks.com/oidc/v1/token") mock_log_debug.assert_called_once_with("Using Service Principal Token.") @mock.patch( @@ -553,7 +553,7 @@ async def test_a_get_token_service_principal_oauth_success(self, mock_conn, mock with mock.patch.object(hook.log, "debug") as mock_log_debug: token = await hook._a_get_token() assert token == "spn_token" - mock_get_sp_token.assert_awaited_once_with("example.databricks.com/oidc/v1/token") + mock_get_sp_token.assert_awaited_once_with("https://example.databricks.com/oidc/v1/token") mock_log_debug.assert_called_once_with("Using Service Principal Token.") @pytest.mark.asyncio @@ -942,7 +942,7 @@ def test_get_federated_token(self, mock_post, mock_file): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -987,7 +987,7 @@ def test_get_federated_token(self, mock_post, mock_file): assert db_call_args[1]["data"]["subject_token"] == "k8s_jwt_token" assert db_call_args[1]["data"]["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" assert db_call_args[1]["data"]["scope"] == "all-apis" - assert "client_id" not in db_call_args[1]["data"] + assert db_call_args[1]["data"]["client_id"] == "test-client-id" @time_machine.travel("2025-07-12 12:00:00") def test_get_federated_token_cached_valid(self): @@ -995,7 +995,7 @@ def test_get_federated_token_cached_valid(self): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1020,7 +1020,7 @@ def test_get_federated_token_k8s_not_available(self, mock_file): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1056,7 +1056,7 @@ def test_get_federated_token_databricks_error(self, mock_post, mock_file): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1232,7 +1232,10 @@ def test_get_federated_token_with_projected_volume(self, mock_post): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + mock_conn.extra_dejson = { + "k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", + "client_id": "test-client-id" + } hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1254,6 +1257,7 @@ def test_get_federated_token_with_projected_volume(self, mock_post): assert db_call_args[1]["data"]["subject_token"] == "projected_k8s_jwt_token" assert db_call_args[1]["data"]["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" assert db_call_args[1]["data"]["scope"] == "all-apis" + assert db_call_args[1]["data"]["client_id"] == "test-client-id" @mock.patch( "builtins.open", @@ -1285,7 +1289,7 @@ def test_get_token_with_federated_k8s_login(self, mock_file): mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" mock_conn.password = None - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1325,7 +1329,7 @@ def test_get_token_with_federated_k8s_extra(self, mock_file): mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = None mock_conn.password = None - mock_conn.extra_dejson = {"federated_k8s": True} # Use default audience + mock_conn.extra_dejson = {"federated_k8s": True, "client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1376,7 +1380,7 @@ async def test_a_get_federated_token(self, mock_post): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} # Use default extra settings + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1404,7 +1408,7 @@ async def test_a_get_federated_token_cached_valid(self): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1429,7 +1433,7 @@ async def test_a_get_federated_token_k8s_not_available(self): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1479,7 +1483,7 @@ async def test_a_get_federated_token_databricks_error(self, mock_post): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {} + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1631,7 +1635,10 @@ async def test_a_get_federated_token_with_projected_volume(self): mock_conn = mock.Mock() mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" - mock_conn.extra_dejson = {"k8s_projected_volume_token_path": "/var/run/secrets/databricks/token"} + mock_conn.extra_dejson = { + "k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", + "client_id": "test-client-id" + } hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1672,6 +1679,7 @@ async def test_a_get_federated_token_with_projected_volume(self): assert resource in call_args[0] assert call_args[1]["data"]["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" assert call_args[1]["data"]["subject_token"] == "projected_k8s_jwt_token" + assert call_args[1]["data"]["client_id"] == "test-client-id" @pytest.mark.asyncio @mock.patch("aiohttp.ClientSession.post") @@ -1712,7 +1720,7 @@ async def test_a_get_token_with_federated_k8s_login(self, mock_post): mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = "federated_k8s" mock_conn.password = None - mock_conn.extra_dejson = {} + mock_conn.extra_dejson = {"client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn @@ -1765,7 +1773,7 @@ async def test_a_get_token_with_federated_k8s_extra(self, mock_post): mock_conn.host = "my-workspace.cloud.databricks.com" mock_conn.login = None mock_conn.password = None - mock_conn.extra_dejson = {"federated_k8s": True} + mock_conn.extra_dejson = {"federated_k8s": True, "client_id": "test-client-id"} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn From f28d6779c3c0a465c2b1369544ee6dc0d832479f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 12 Feb 2026 18:36:11 +0100 Subject: [PATCH 19/21] updated docs --- providers/databricks/docs/connections/databricks.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index aba11cea4eca2..c71f2e640b3e6 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -113,7 +113,7 @@ Extra (optional) **Method 2: TokenRequest API** - If ``k8s_projected_volume_token_path`` is not configured, the hook will use the `TokenRequest API `_ method (dynamic token generation): + If ``k8s_projected_volume_token_path`` is not configured, the hook will use the `TokenRequest API `_ method (dynamic token generation with in-cluster authentication): * ``audience``: (optional) the audience value for the Kubernetes JWT token (default: ``https://kubernetes.default.svc``). **Important:** For production deployments, especially when using multiple Kubernetes clusters, it is recommended to use a unique audience per cluster/environment (e.g., ``databricks-prod-airflow``, ``databricks-staging-airflow``) to allow separate Databricks federation policies and proper access control. The default generic audience is only suitable for single-cluster development setups. * ``expiration_seconds``: (optional) token expiration in seconds for the Kubernetes JWT (default: 3600). From 97a768609775f0d67cd65adda59687327671169b Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Fri, 13 Feb 2026 11:30:20 +0100 Subject: [PATCH 20/21] * Minor code refactor, added unit tests --- .../databricks/hooks/databricks_base.py | 12 ++---- .../databricks/hooks/test_databricks_base.py | 37 ++++++++++++++++++- 2 files changed, 39 insertions(+), 10 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 12c1031b8a112..4c6bdb5339458 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -218,7 +218,7 @@ def _parse_host(host: str) -> str: assert h._parse_host('xx.cloud.databricks.com') == 'xx.cloud.databricks.com' """ - host = host.rstrip('/') # Remove trailing slashes + host = host.rstrip("/") # Remove trailing slashes urlparse_host = urlsplit(host).hostname if urlparse_host: # In this case, host = https://xx.cloud.databricks.com @@ -784,7 +784,7 @@ def _get_required_client_id(self) -> str: :return: Service principal client ID :raises AirflowException: If client_id is not provided """ - client_id = self.databricks_conn.extra_dejson.get("client_id", None) + client_id = self.databricks_conn.extra_dejson.get("client_id") if not client_id: # see: https://github.com/kubernetes/kubernetes/issues/116638 raise AirflowException( @@ -987,9 +987,7 @@ def _get_token(self, raise_error: bool = False) -> str | None: "federated_k8s", False ): self.log.debug("Using Kubernetes OIDC token federation.") - return self._get_federated_databricks_token( - self._get_oidc_token_service_url() - ) + return self._get_federated_databricks_token(self._get_oidc_token_service_url()) if raise_error: raise AirflowException("Token authentication isn't configured") @@ -1026,9 +1024,7 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None: "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() - ) + 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 07b4aceb5563d..4c764cb2256ab 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -1029,6 +1029,22 @@ def test_get_federated_token_k8s_not_available(self, mock_file): with pytest.raises(AirflowException, match="Kubernetes service account token not found"): hook._get_federated_databricks_token(resource) + def test_get_federated_token_missing_client_id(self): + """Test error when client_id is missing from connection extra.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Missing client_id + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + with pytest.raises( + AirflowException, match="client_id is required for Kubernetes OIDC token federation" + ): + hook._get_federated_databricks_token(resource) + @mock.patch( "builtins.open", side_effect=[ @@ -1234,7 +1250,7 @@ def test_get_federated_token_with_projected_volume(self, mock_post): mock_conn.login = "federated_k8s" mock_conn.extra_dejson = { "k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", - "client_id": "test-client-id" + "client_id": "test-client-id", } hook = BaseDatabricksHook() @@ -1443,6 +1459,23 @@ async def test_a_get_federated_token_k8s_not_available(self): with pytest.raises(AirflowException, match="Kubernetes service account token not found"): await hook._a_get_federated_databricks_token(resource) + @pytest.mark.asyncio + async def test_a_get_federated_token_missing_client_id(self): + """Test async error when client_id is missing from connection extra.""" + mock_conn = mock.Mock() + mock_conn.host = "my-workspace.cloud.databricks.com" + mock_conn.login = "federated_k8s" + mock_conn.extra_dejson = {} # Missing client_id + + hook = BaseDatabricksHook() + hook.databricks_conn = mock_conn + + resource = f"https://{mock_conn.host}/oidc/v1/token" + with pytest.raises( + AirflowException, match="client_id is required for Kubernetes OIDC token federation" + ): + await hook._a_get_federated_databricks_token(resource) + @pytest.mark.asyncio @mock.patch("aiohttp.ClientSession.post") async def test_a_get_federated_token_databricks_error(self, mock_post): @@ -1637,7 +1670,7 @@ async def test_a_get_federated_token_with_projected_volume(self): mock_conn.login = "federated_k8s" mock_conn.extra_dejson = { "k8s_projected_volume_token_path": "/var/run/secrets/databricks/token", - "client_id": "test-client-id" + "client_id": "test-client-id", } hook = BaseDatabricksHook() From 7d0049d15ca537833976dcd4dbb0238ca9835f49 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Mon, 16 Feb 2026 11:48:36 +0100 Subject: [PATCH 21/21] updated docs to fix ci check --- providers/databricks/docs/connections/databricks.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/databricks/docs/connections/databricks.rst b/providers/databricks/docs/connections/databricks.rst index c71f2e640b3e6..084743b49834d 100644 --- a/providers/databricks/docs/connections/databricks.rst +++ b/providers/databricks/docs/connections/databricks.rst @@ -109,7 +109,7 @@ Extra (optional) **Method 1: Projected Volume** - * ``k8s_projected_volume_token_path``: (optional) path to a `Kubernetes projected volume service account token `_. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. **Important:** The ``k8s_projected_volume_token_path`` must match the full path constructed from your Pod spec: ``{mountPath}/{path}``, where ``mountPath`` is the volume mount path and ``path`` is the serviceAccountToken path field. For example, if your Pod spec has ``mountPath: /var/run/secrets/databricks`` and ``path: token``, then ``k8s_projected_volume_token_path`` should be ``/var/run/secrets/databricks/token``. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. + * ``k8s_projected_volume_token_path``: (optional) path to a `Kubernetes projected volume service account token `_. When configured, the hook will read the token directly from this file. The token must be configured in your Pod spec with the appropriate audience and expiration. **Important:** The ``k8s_projected_volume_token_path`` must match the full path constructed from your Pod spec: ``{mountPath}/{path}``, where ``mountPath`` is the volume mount path and ``path`` is the ``serviceAccountToken`` path field. For example, if your Pod spec has ``mountPath: /var/run/secrets/databricks`` and ``path: token``, then ``k8s_projected_volume_token_path`` should be ``/var/run/secrets/databricks/token``. This is the recommended method as it's simpler and more efficient (no API calls). See the example Pod configuration below. **Method 2: TokenRequest API** @@ -158,7 +158,7 @@ Extra (optional) **Note:** The ``k8s_projected_volume_token_path`` value must match ``{mountPath}/{path}`` from your Pod spec. In the example above: - ``mountPath`` is ``/var/run/secrets/databricks`` (from volumeMounts) - - ``path`` is ``token`` (from serviceAccountToken) + - ``path`` is ``token`` (from ``serviceAccountToken``) - Full path is ``/var/run/secrets/databricks/token`` **Option B: Using TokenRequest API**