From ecfd1423af1e3a733b41ad49ae74bd00a1d75991 Mon Sep 17 00:00:00 2001 From: Rohit Sharma Date: Thu, 30 Jul 2026 14:01:27 +0100 Subject: [PATCH 1/3] refactor(spark): isolate Databricks PAT authentication Introduce an authentication strategy boundary while preserving the existing PAT resolution and connector behavior, enabling subsequent auth methods to be reviewed independently. Signed-off-by: Rohit Sharma --- .../flytekitplugins/spark/connector.py | 17 ++-- .../flytekitplugins/spark/databricks_auth.py | 83 +++++++++++++++++++ .../tests/test_databricks_auth.py | 66 +++++++++++++++ 3 files changed, 154 insertions(+), 12 deletions(-) create mode 100644 plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py create mode 100644 plugins/flytekit-spark/tests/test_databricks_auth.py diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 06b02048d1..7b47f1b3f5 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/connector.py @@ -262,6 +262,8 @@ async def create( task_execution_metadata: Optional[TaskExecutionMetadata] = None, **kwargs, ) -> DatabricksJobMetadata: + from .databricks_auth import select_auth + data = json.dumps(_get_databricks_job_spec(task_template)) databricks_instance = task_template.custom.get( "databricksInstance", os.getenv(DEFAULT_DATABRICKS_INSTANCE_ENV_KEY) @@ -272,22 +274,13 @@ async def create( f"Missing databricks instance. Please set the value through the task config or set the {DEFAULT_DATABRICKS_INSTANCE_ENV_KEY} environment variable in the connector." ) - # Get workflow-specific token or fall back to default namespace = task_execution_metadata.namespace if task_execution_metadata else None - - # Extract custom secret name from task template (if provided) - custom_secret_name = task_template.custom.get("databricksTokenSecret") - - logger.info(f"Creating Databricks job for namespace: {namespace or 'unknown'}") - if custom_secret_name: - logger.info(f"Using custom secret name: {custom_secret_name}") - - auth_token = get_databricks_token( - namespace=namespace, task_template=task_template, secret_name=custom_secret_name - ) + auth = await select_auth(task_template=task_template, namespace=namespace) + logger.info("Databricks auth resolved: %s", auth.describe()) databricks_url = f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/submit" async with aiohttp.ClientSession() as session: + auth_token = await auth.get_bearer_token(session) async with session.post(databricks_url, headers=get_header(auth_token=auth_token), data=data) as resp: response = await resp.json() if resp.status != http.HTTPStatus.OK: diff --git a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py new file mode 100644 index 0000000000..bda5f2a19d --- /dev/null +++ b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py @@ -0,0 +1,83 @@ +"""Authentication strategies for the Databricks connector. + +This module introduces a small strategy boundary around the connector's +existing Personal Access Token (PAT) flow. It intentionally preserves the +current token resolution behavior so additional authentication methods can be +added independently. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Optional + +from flytekit import lazy_module +from flytekit.models.task import TaskTemplate + +aiohttp = lazy_module("aiohttp") + + +@dataclass +class _Settings: + """Authentication settings resolved for one Databricks task.""" + + task_template: Optional[TaskTemplate] + token_secret_name: Optional[str] + namespace: Optional[str] + + @staticmethod + def from_task(task_template: Optional[TaskTemplate], namespace: Optional[str]) -> "_Settings": + custom = task_template.custom if task_template is not None else {} + return _Settings( + task_template=task_template, + token_secret_name=custom.get("databricksTokenSecret"), + namespace=namespace, + ) + + +class DatabricksAuth(ABC): + """Interface for obtaining a bearer token for Databricks API calls.""" + + auth_type = "unknown" + strategy_name = "DatabricksAuth" + + def __init__(self, settings: _Settings): + self.settings = settings + + @abstractmethod + async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] + """Return a bearer token for a Databricks API request.""" + + def describe(self) -> str: + """Return a description that is safe to write to connector logs.""" + return ( + f"strategy={self.strategy_name} auth_type={self.auth_type} " f"namespace={self.settings.namespace or 'N/A'}" + ) + + +class PATAuth(DatabricksAuth): + """Delegate to the connector's existing multi-tenant PAT lookup.""" + + auth_type = "pat" + strategy_name = "PATAuth" + + async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] + from .connector import get_databricks_token + + return get_databricks_token( + namespace=self.settings.namespace, + task_template=self.settings.task_template, + secret_name=self.settings.token_secret_name, + ) + + +async def select_auth( + task_template: Optional[TaskTemplate], + namespace: Optional[str], +) -> DatabricksAuth: + """Select authentication for a task. + + PAT remains the only strategy and therefore the unconditional default in + this refactor. Later authentication methods can extend this dispatcher + without changing the connector request lifecycle. + """ + return PATAuth(_Settings.from_task(task_template, namespace)) diff --git a/plugins/flytekit-spark/tests/test_databricks_auth.py b/plugins/flytekit-spark/tests/test_databricks_auth.py new file mode 100644 index 0000000000..f1cf1bd9a9 --- /dev/null +++ b/plugins/flytekit-spark/tests/test_databricks_auth.py @@ -0,0 +1,66 @@ +"""Tests for the Databricks authentication strategy boundary.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from flytekitplugins.spark.databricks_auth import ( + PATAuth, + _Settings, + select_auth, +) + + +def _task_template(**custom): + task_template = MagicMock() + task_template.custom = custom + return task_template + + +def test_settings_use_default_secret_name(): + settings = _Settings.from_task(task_template=None, namespace="project-a") + + assert settings.token_secret_name is None + assert settings.namespace == "project-a" + + +def test_settings_use_task_secret_name(): + task_template = _task_template() + task_template.custom["databricksTokenSecret"] = "custom-token" + + settings = _Settings.from_task(task_template=task_template, namespace="project-a") + + assert settings.token_secret_name == "custom-token" + + +@pytest.mark.asyncio +async def test_select_auth_returns_pat_by_default(): + task_template = _task_template() + auth = await select_auth(task_template=task_template, namespace="project-a") + + assert isinstance(auth, PATAuth) + assert auth.auth_type == "pat" + + +@pytest.mark.asyncio +async def test_pat_auth_delegates_to_existing_token_lookup(): + task_template = _task_template() + settings = _Settings( + task_template=task_template, + token_secret_name="custom-token", + namespace="project-a", + ) + auth = PATAuth(settings) + + with patch( + "flytekitplugins.spark.connector.get_databricks_token", + return_value="example-token", + ) as get_token: + token = await auth.get_bearer_token(AsyncMock()) + + assert token == "example-token" + get_token.assert_called_once_with( + namespace="project-a", + task_template=task_template, + secret_name="custom-token", + ) From e94cbb2a2acb5cd595d4e83c9fce1989dcafa42c Mon Sep 17 00:00:00 2001 From: Rohit Sharma Date: Thu, 30 Jul 2026 14:06:34 +0100 Subject: [PATCH 2/3] feat(spark): add Databricks OAuth M2M authentication Allow connectors and individual tasks to opt into short-lived service-principal tokens while keeping PAT as the unchanged default authentication mode. Signed-off-by: Rohit Sharma --- plugins/flytekit-spark/README.md | 40 +++ .../flytekitplugins/spark/connector.py | 112 ++++-- .../flytekitplugins/spark/databricks_auth.py | 274 ++++++++++++++- .../flytekitplugins/spark/task.py | 15 + .../flytekit-spark/tests/test_connector.py | 6 + .../tests/test_databricks_auth.py | 323 +++++++++++++++++- .../tests/test_databricks_token.py | 6 + 7 files changed, 741 insertions(+), 35 deletions(-) diff --git a/plugins/flytekit-spark/README.md b/plugins/flytekit-spark/README.md index 9cc7c7cf9d..f5609e57ce 100644 --- a/plugins/flytekit-spark/README.md +++ b/plugins/flytekit-spark/README.md @@ -11,3 +11,43 @@ pip install flytekitplugins-spark To configure Spark in the Flyte deployment's backend, follow [Step 1](https://docs.flyte.org/en/latest/deployment/plugins/k8s/index.html#deployment-plugin-setup-k8s), [2](https://docs.flyte.org/en/latest/flytesnacks/examples/k8s_spark_plugin/index.html). All [examples](https://docs.flyte.org/en/latest/flytesnacks/examples/k8s_spark_plugin/index.html) showcasing execution of Spark jobs using the plugin can be found in the documentation. + +## Databricks authentication + +The Databricks connector uses PAT authentication by default, preserving the +existing `databricks-token` namespace Secret and +`FLYTE_DATABRICKS_ACCESS_TOKEN` fallback. + +OAuth machine-to-machine (M2M) authentication can be enabled on the connector: + +```yaml +env: + - name: FLYTE_DATABRICKS_AUTH_TYPE + value: oauth_m2m + - name: DATABRICKS_CLIENT_ID + value: "" + - name: DATABRICKS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: databricks-connector-oauth + key: client_secret +``` + +For per-namespace identities, create a `databricks-oauth` Secret in each +workflow namespace: + +```yaml +apiVersion: v1 +kind: Secret +metadata: + name: databricks-oauth + namespace: "" +type: Opaque +stringData: + client_id: "" + client_secret: "" +``` + +The namespace Secret takes precedence over connector-level credentials. +`get` and `delete` operations cache short-lived OAuth tokens and retry once +with a refreshed token when the Databricks API returns HTTP 401. diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 7b47f1b3f5..6c8252f495 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/connector.py @@ -29,9 +29,20 @@ @dataclass class DatabricksJobMetadata(ResourceMeta): + """Metadata persisted for a Databricks run. + + OAuth metadata allows ``get`` and ``delete`` to obtain fresh short-lived + tokens. ``auth_token`` remains populated for PAT jobs and for metadata + written by older connector versions. + """ + databricks_instance: str run_id: str - auth_token: Optional[str] = None # Store auth token for get/delete operations + auth_token: Optional[str] = None + auth_type: Optional[str] = None + client_id: Optional[str] = None + oauth_secret_name: Optional[str] = None + namespace: Optional[str] = None def _configure_serverless(databricks_job: dict, envs: dict) -> str: @@ -254,6 +265,9 @@ class DatabricksConnector(AsyncConnectorBase): def __init__(self): super().__init__(task_type_name="spark", metadata_type=DatabricksJobMetadata) + from .databricks_auth import validate_connector_config + + validate_connector_config() async def create( self, @@ -275,7 +289,11 @@ async def create( ) namespace = task_execution_metadata.namespace if task_execution_metadata else None - auth = await select_auth(task_template=task_template, namespace=namespace) + auth = await select_auth( + task_template=task_template, + workspace_url=databricks_instance, + namespace=namespace, + ) logger.info("Databricks auth resolved: %s", auth.describe()) databricks_url = f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/submit" @@ -288,7 +306,13 @@ async def create( logger.info(f"Successfully created Databricks job with run_id: {response['run_id']}") return DatabricksJobMetadata( - databricks_instance=databricks_instance, run_id=str(response["run_id"]), auth_token=auth_token + databricks_instance=databricks_instance, + run_id=str(response["run_id"]), + auth_token=auth_token if auth.auth_type == "pat" else None, + auth_type=auth.auth_type, + client_id=auth.settings.client_id, + oauth_secret_name=auth.settings.oauth_secret_name, + namespace=namespace, ) async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource: @@ -297,14 +321,14 @@ async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource: f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/get?run_id={resource_meta.run_id}" ) - # Use the stored auth token if available, otherwise fall back to default - headers = get_header(auth_token=resource_meta.auth_token) - async with aiohttp.ClientSession() as session: - async with session.get(databricks_url, headers=headers) as resp: - if resp.status != http.HTTPStatus.OK: - raise RuntimeError(f"Failed to get databricks job {resource_meta.run_id} with error: {resp.reason}") - response = await resp.json() + response = await self._request_with_auth( + session=session, + method="GET", + url=databricks_url, + resource_meta=resource_meta, + action_label=f"get databricks job {resource_meta.run_id}", + ) cur_phase = TaskExecution.UNDEFINED message = "" @@ -332,16 +356,66 @@ async def delete(self, resource_meta: DatabricksJobMetadata, **kwargs): databricks_url = f"https://{resource_meta.databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/cancel" data = json.dumps({"run_id": resource_meta.run_id}) - # Use the stored auth token if available, otherwise fall back to default - headers = get_header(auth_token=resource_meta.auth_token) - async with aiohttp.ClientSession() as session: - async with session.post(databricks_url, headers=headers, data=data) as resp: - if resp.status != http.HTTPStatus.OK: - raise RuntimeError( - f"Failed to cancel databricks job {resource_meta.run_id} with error: {resp.reason}" - ) - await resp.json() + await self._request_with_auth( + session=session, + method="POST", + url=databricks_url, + resource_meta=resource_meta, + data=data, + action_label=f"cancel databricks job {resource_meta.run_id}", + ) + + async def _request_with_auth( + self, + session: "aiohttp.ClientSession", # type: ignore[name-defined] + method: str, + url: str, + resource_meta: DatabricksJobMetadata, + action_label: str, + data: Optional[str] = None, + ) -> dict: + """Call the Jobs API and retry once after refreshing OAuth on 401.""" + from .databricks_auth import DatabricksAuthError, build_auth + + auth = None + if resource_meta.auth_type == "oauth_m2m": + auth = build_auth( + workspace_url=resource_meta.databricks_instance, + auth_type=resource_meta.auth_type, + namespace=resource_meta.namespace, + client_id=resource_meta.client_id, + oauth_secret_name=resource_meta.oauth_secret_name, + ) + + token = resource_meta.auth_token + if auth is not None: + try: + token = await auth.get_bearer_token(session) + except DatabricksAuthError as error: + raise RuntimeError(f"Failed to {action_label}: could not obtain Databricks auth: {error}") from error + + def _request(bearer: Optional[str]): + headers = get_header(auth_token=bearer) + if method.upper() == "GET": + return session.get(url, headers=headers) + return session.post(url, headers=headers, data=data) + + async with _request(token) as response: + if response.status == http.HTTPStatus.UNAUTHORIZED and auth is not None: + await auth.invalidate_cache() + try: + refreshed_token = await auth.get_bearer_token(session) + except DatabricksAuthError as error: + raise RuntimeError(f"Failed to {action_label}: auth refresh failed after 401: {error}") from error + async with _request(refreshed_token) as retry_response: + if retry_response.status != http.HTTPStatus.OK: + raise RuntimeError(f"Failed to {action_label} with error: {retry_response.reason}") + return await retry_response.json() + + if response.status != http.HTTPStatus.OK: + raise RuntimeError(f"Failed to {action_label} with error: {response.reason}") + return await response.json() class DatabricksConnectorV2(DatabricksConnector): diff --git a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py index bda5f2a19d..dd3ed8730d 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py @@ -1,52 +1,183 @@ """Authentication strategies for the Databricks connector. -This module introduces a small strategy boundary around the connector's -existing Personal Access Token (PAT) flow. It intentionally preserves the -current token resolution behavior so additional authentication methods can be -added independently. +PAT remains the default for backward compatibility. OAuth machine-to-machine +(M2M) authentication is opt-in through ``FLYTE_DATABRICKS_AUTH_TYPE`` or the +equivalent per-task setting. """ +import asyncio +import json +import logging +import os +import random +import time from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Optional +from typing import Any, Dict, Optional, Tuple from flytekit import lazy_module from flytekit.models.task import TaskTemplate aiohttp = lazy_module("aiohttp") +logger = logging.getLogger(__name__) + +FLYTE_DATABRICKS_AUTH_TYPE_ENV = "FLYTE_DATABRICKS_AUTH_TYPE" +FLYTE_DATABRICKS_OAUTH_SECRET_NAME_ENV = "FLYTE_DATABRICKS_OAUTH_SECRET_NAME" +DATABRICKS_CLIENT_ID_ENV = "DATABRICKS_CLIENT_ID" +DATABRICKS_CLIENT_SECRET_ENV = "DATABRICKS_CLIENT_SECRET" + +DEFAULT_OAUTH_SECRET_NAME = "databricks-oauth" +TOKEN_REFRESH_BUFFER_SECONDS = 60 +TOKEN_ENDPOINT_MAX_RETRIES = 3 +TOKEN_ENDPOINT_BACKOFF_BASE_SECONDS = 0.2 +VALID_AUTH_TYPES = {"pat", "oauth_m2m"} + + +class DatabricksAuthError(Exception): + """Raised when Databricks authentication cannot be obtained.""" + @dataclass class _Settings: """Authentication settings resolved for one Databricks task.""" task_template: Optional[TaskTemplate] + auth_type: Optional[str] + client_id: Optional[str] + oauth_secret_name: str token_secret_name: Optional[str] namespace: Optional[str] @staticmethod def from_task(task_template: Optional[TaskTemplate], namespace: Optional[str]) -> "_Settings": - custom = task_template.custom if task_template is not None else {} + custom: Dict[str, Any] = task_template.custom if task_template is not None else {} + + def _pick(task_key: str, env_key: Optional[str], default: Optional[str] = None) -> Optional[str]: + task_value = custom.get(task_key) + if task_value: + return task_value + if env_key: + env_value = os.getenv(env_key) + if env_value: + return env_value + return default + return _Settings( task_template=task_template, + auth_type=_pick("databricksAuthType", FLYTE_DATABRICKS_AUTH_TYPE_ENV), + client_id=_pick("databricksClientId", DATABRICKS_CLIENT_ID_ENV), + oauth_secret_name=_pick( + "databricksOauthSecret", + FLYTE_DATABRICKS_OAUTH_SECRET_NAME_ENV, + DEFAULT_OAUTH_SECRET_NAME, + ) + or DEFAULT_OAUTH_SECRET_NAME, token_secret_name=custom.get("databricksTokenSecret"), namespace=namespace, ) +@dataclass +class _CachedToken: + access_token: str + refresh_at: float + + +class _TokenCache: + """Async-safe in-memory cache for short-lived OAuth tokens.""" + + def __init__(self) -> None: + self._store: Dict[Tuple[str, str, str], _CachedToken] = {} + self._lock = asyncio.Lock() + + async def get(self, key: Tuple[str, str, str]) -> Optional[str]: + async with self._lock: + entry = self._store.get(key) + if entry is None: + return None + if entry.refresh_at <= time.time(): + self._store.pop(key, None) + return None + return entry.access_token + + async def put(self, key: Tuple[str, str, str], access_token: str, expires_in: int) -> None: + async with self._lock: + refresh_after = max(int(expires_in) - TOKEN_REFRESH_BUFFER_SECONDS, 0) + self._store[key] = _CachedToken( + access_token=access_token, + refresh_at=time.time() + refresh_after, + ) + + async def invalidate(self, key: Tuple[str, str, str]) -> None: + async with self._lock: + self._store.pop(key, None) + + +_TOKEN_CACHE = _TokenCache() + + +async def _post_token( + session: "aiohttp.ClientSession", # type: ignore[name-defined] + workspace_url: str, + form: Dict[str, str], +) -> Dict[str, Any]: + """Request a Databricks OAuth token with bounded transient retries.""" + url = f"https://{workspace_url.rstrip('/')}/oidc/v1/token" + last_error: Optional[str] = None + + for attempt in range(TOKEN_ENDPOINT_MAX_RETRIES): + try: + async with session.post( + url, + data=form, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) as response: + body = await response.text() + if response.status == 200: + try: + return json.loads(body) + except json.JSONDecodeError as error: + raise DatabricksAuthError( + f"Databricks token endpoint returned invalid JSON: {error}" + ) from error + if response.status in (429, 500, 502, 503, 504): + last_error = f"HTTP {response.status}: {body[:500]}" + else: + raise DatabricksAuthError( + f"Databricks token endpoint returned HTTP {response.status}: {body[:500]}" + ) + except aiohttp.ClientError as error: # type: ignore[attr-defined] + last_error = f"network error: {error}" + except asyncio.TimeoutError: + last_error = "timeout" + + if attempt < TOKEN_ENDPOINT_MAX_RETRIES - 1: + delay = TOKEN_ENDPOINT_BACKOFF_BASE_SECONDS * (2**attempt) + random.uniform(0, 0.1) + await asyncio.sleep(delay) + + raise DatabricksAuthError( + f"Databricks token endpoint failed after {TOKEN_ENDPOINT_MAX_RETRIES} attempts: {last_error}" + ) + + class DatabricksAuth(ABC): """Interface for obtaining a bearer token for Databricks API calls.""" auth_type = "unknown" strategy_name = "DatabricksAuth" - def __init__(self, settings: _Settings): + def __init__(self, workspace_url: str, settings: _Settings): + self.workspace_url = workspace_url self.settings = settings @abstractmethod async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] """Return a bearer token for a Databricks API request.""" + async def invalidate_cache(self) -> None: + """Invalidate cached authentication state, if any.""" + def describe(self) -> str: """Return a description that is safe to write to connector logs.""" return ( @@ -70,14 +201,131 @@ async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # ty ) +class OAuthM2MAuth(DatabricksAuth): + """Authenticate a Databricks service principal with client credentials.""" + + auth_type = "oauth_m2m" + strategy_name = "OAuthM2MAuth" + + @property + def cache_key(self) -> Tuple[str, str, str]: + return ( + self.workspace_url, + self.settings.client_id or "", + self.settings.namespace or "_", + ) + + def _resolve_credentials(self) -> Tuple[str, str]: + from .connector import get_secret_from_k8s + + client_id = self.settings.client_id + client_secret: Optional[str] = None + + if self.settings.namespace: + secret_client_id = get_secret_from_k8s( + secret_name=self.settings.oauth_secret_name, + secret_key="client_id", + namespace=self.settings.namespace, + ) + secret_client_secret = get_secret_from_k8s( + secret_name=self.settings.oauth_secret_name, + secret_key="client_secret", + namespace=self.settings.namespace, + ) + if secret_client_id: + client_id = secret_client_id + if secret_client_secret: + client_secret = secret_client_secret + + client_id = client_id or os.getenv(DATABRICKS_CLIENT_ID_ENV) + client_secret = client_secret or os.getenv(DATABRICKS_CLIENT_SECRET_ENV) + + if not client_id: + raise DatabricksAuthError( + "OAuth M2M requires a client ID. Configure databricks_client_id, " + "DATABRICKS_CLIENT_ID, or the namespace OAuth secret." + ) + if not client_secret: + raise DatabricksAuthError( + "OAuth M2M requires a client secret. Configure DATABRICKS_CLIENT_SECRET " + "or the namespace OAuth secret." + ) + return client_id, client_secret + + async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] + client_id, client_secret = self._resolve_credentials() + key = (self.workspace_url, client_id, self.settings.namespace or "_") + cached = await _TOKEN_CACHE.get(key) + if cached: + return cached + + payload = await _post_token( + session, + self.workspace_url, + { + "grant_type": "client_credentials", + "scope": "all-apis", + "client_id": client_id, + "client_secret": client_secret, + }, + ) + access_token = payload.get("access_token") + if not access_token: + raise DatabricksAuthError("Databricks OAuth response did not contain access_token") + await _TOKEN_CACHE.put(key, access_token, int(payload.get("expires_in", 3600))) + return access_token + + async def invalidate_cache(self) -> None: + client_id, _ = self._resolve_credentials() + await _TOKEN_CACHE.invalidate((self.workspace_url, client_id, self.settings.namespace or "_")) + + async def select_auth( task_template: Optional[TaskTemplate], + workspace_url: str, namespace: Optional[str], ) -> DatabricksAuth: - """Select authentication for a task. + """Select an explicitly configured strategy, defaulting to PAT.""" + settings = _Settings.from_task(task_template, namespace) + auth_type = settings.auth_type or "pat" + if auth_type not in VALID_AUTH_TYPES: + raise DatabricksAuthError( + f"Invalid Databricks auth type '{auth_type}'. Expected one of {sorted(VALID_AUTH_TYPES)}." + ) + if auth_type == "oauth_m2m": + return OAuthM2MAuth(workspace_url, settings) + return PATAuth(workspace_url, settings) - PAT remains the only strategy and therefore the unconditional default in - this refactor. Later authentication methods can extend this dispatcher - without changing the connector request lifecycle. - """ - return PATAuth(_Settings.from_task(task_template, namespace)) + +def build_auth( + workspace_url: str, + auth_type: str, + namespace: Optional[str] = None, + client_id: Optional[str] = None, + oauth_secret_name: Optional[str] = None, +) -> DatabricksAuth: + """Rebuild an auth strategy from persisted connector metadata.""" + settings = _Settings( + task_template=None, + auth_type=auth_type, + client_id=client_id, + oauth_secret_name=oauth_secret_name or DEFAULT_OAUTH_SECRET_NAME, + token_secret_name=None, + namespace=namespace, + ) + if auth_type == "oauth_m2m": + return OAuthM2MAuth(workspace_url, settings) + if auth_type == "pat": + return PATAuth(workspace_url, settings) + raise DatabricksAuthError( + f"Invalid Databricks auth type '{auth_type}'. Expected one of {sorted(VALID_AUTH_TYPES)}." + ) + + +def validate_connector_config() -> None: + """Validate an explicitly configured connector-wide auth type.""" + auth_type = os.getenv(FLYTE_DATABRICKS_AUTH_TYPE_ENV) + if auth_type and auth_type not in VALID_AUTH_TYPES: + raise DatabricksAuthError( + f"Invalid {FLYTE_DATABRICKS_AUTH_TYPE_ENV}='{auth_type}'. " f"Expected one of {sorted(VALID_AUTH_TYPES)}." + ) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index 6c447e8bd4..afea0e1470 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -89,6 +89,12 @@ class DatabricksV2(Spark): Service Credentials for S3 access. Falls back to FLYTE_DATABRICKS_SERVICE_CREDENTIAL_PROVIDER env var. databricks_token_secret (Optional[str]): Custom name for the K8s secret containing the Databricks token. Defaults to 'databricks-token' if not specified. + databricks_auth_type (Optional[str]): Authentication mode. Supported values are + ``"pat"`` and ``"oauth_m2m"``. When unset, PAT remains the default. + databricks_client_id (Optional[str]): Databricks service-principal client ID used + for OAuth M2M. Falls back to ``DATABRICKS_CLIENT_ID`` on the connector. + databricks_oauth_secret (Optional[str]): Name of the namespace K8s secret containing + ``client_id`` and ``client_secret``. Defaults to ``databricks-oauth``. notebook_path (Optional[str]): Path to Databricks notebook (e.g., "/Users/user@example.com/notebook"). notebook_base_parameters (Optional[Dict[str, str]]): Parameters to pass to the notebook. @@ -199,6 +205,9 @@ class DatabricksV2(Spark): databricks_instance: Optional[str] = None databricks_service_credential_provider: Optional[str] = None databricks_token_secret: Optional[str] = None + databricks_auth_type: Optional[str] = None + databricks_client_id: Optional[str] = None + databricks_oauth_secret: Optional[str] = None notebook_path: Optional[str] = None notebook_base_parameters: Optional[Dict[str, str]] = None @@ -314,6 +323,12 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: custom_dict["databricksServiceCredentialProvider"] = cfg.databricks_service_credential_provider if cfg.databricks_token_secret: custom_dict["databricksTokenSecret"] = cfg.databricks_token_secret + if cfg.databricks_auth_type: + custom_dict["databricksAuthType"] = cfg.databricks_auth_type + if cfg.databricks_client_id: + custom_dict["databricksClientId"] = cfg.databricks_client_id + if cfg.databricks_oauth_secret: + custom_dict["databricksOauthSecret"] = cfg.databricks_oauth_secret if cfg.notebook_path: custom_dict["notebookPath"] = cfg.notebook_path if cfg.notebook_base_parameters: diff --git a/plugins/flytekit-spark/tests/test_connector.py b/plugins/flytekit-spark/tests/test_connector.py index 0fc4effbb9..803d609c8b 100644 --- a/plugins/flytekit-spark/tests/test_connector.py +++ b/plugins/flytekit-spark/tests/test_connector.py @@ -125,6 +125,8 @@ async def test_databricks_agent(task_template: TaskTemplate): databricks_instance="test-account.cloud.databricks.com", run_id="123", auth_token=mocked_token, + auth_type="pat", + oauth_secret_name="databricks-oauth", ) mock_create_response = {"run_id": "123"} @@ -187,6 +189,8 @@ async def test_agent_create_with_default_instance(task_template: TaskTemplate): databricks_instance="test-account.cloud.databricks.com", run_id="123", auth_token=mocked_token, + auth_type="pat", + oauth_secret_name="databricks-oauth", ) mock_create_response = {"run_id": "123"} @@ -609,6 +613,8 @@ async def test_databricks_agent_serverless(serverless_task_template_with_env_key databricks_instance="test-account.cloud.databricks.com", run_id="456", auth_token=mocked_token, + auth_type="pat", + oauth_secret_name="databricks-oauth", ) mock_create_response = {"run_id": "456"} diff --git a/plugins/flytekit-spark/tests/test_databricks_auth.py b/plugins/flytekit-spark/tests/test_databricks_auth.py index f1cf1bd9a9..2695f32247 100644 --- a/plugins/flytekit-spark/tests/test_databricks_auth.py +++ b/plugins/flytekit-spark/tests/test_databricks_auth.py @@ -1,12 +1,21 @@ -"""Tests for the Databricks authentication strategy boundary.""" +"""Tests for Databricks PAT and OAuth M2M authentication.""" +import http from unittest.mock import AsyncMock, MagicMock, patch import pytest +from aiohttp import ClientSession +from aioresponses import aioresponses from flytekitplugins.spark.databricks_auth import ( + DEFAULT_OAUTH_SECRET_NAME, + DatabricksAuthError, + OAuthM2MAuth, PATAuth, + _TokenCache, _Settings, + _post_token, + build_auth, select_auth, ) @@ -17,10 +26,22 @@ def _task_template(**custom): return task_template +@pytest.fixture(autouse=True) +def _clear_auth_environment(monkeypatch): + for name in ( + "FLYTE_DATABRICKS_AUTH_TYPE", + "FLYTE_DATABRICKS_OAUTH_SECRET_NAME", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + ): + monkeypatch.delenv(name, raising=False) + + def test_settings_use_default_secret_name(): settings = _Settings.from_task(task_template=None, namespace="project-a") assert settings.token_secret_name is None + assert settings.oauth_secret_name == DEFAULT_OAUTH_SECRET_NAME assert settings.namespace == "project-a" @@ -36,7 +57,11 @@ def test_settings_use_task_secret_name(): @pytest.mark.asyncio async def test_select_auth_returns_pat_by_default(): task_template = _task_template() - auth = await select_auth(task_template=task_template, namespace="project-a") + auth = await select_auth( + task_template=task_template, + workspace_url="example.cloud.databricks.com", + namespace="project-a", + ) assert isinstance(auth, PATAuth) assert auth.auth_type == "pat" @@ -47,10 +72,13 @@ async def test_pat_auth_delegates_to_existing_token_lookup(): task_template = _task_template() settings = _Settings( task_template=task_template, + auth_type="pat", + client_id=None, + oauth_secret_name=DEFAULT_OAUTH_SECRET_NAME, token_secret_name="custom-token", namespace="project-a", ) - auth = PATAuth(settings) + auth = PATAuth("example.cloud.databricks.com", settings) with patch( "flytekitplugins.spark.connector.get_databricks_token", @@ -64,3 +92,292 @@ async def test_pat_auth_delegates_to_existing_token_lookup(): task_template=task_template, secret_name="custom-token", ) + + +@pytest.mark.asyncio +async def test_select_auth_uses_explicit_m2m(): + auth = await select_auth( + task_template=_task_template(databricksAuthType="oauth_m2m"), + workspace_url="example.cloud.databricks.com", + namespace="project-a", + ) + + assert isinstance(auth, OAuthM2MAuth) + + +@pytest.mark.asyncio +async def test_select_auth_rejects_unknown_type(): + with pytest.raises(DatabricksAuthError, match="Invalid Databricks auth type"): + await select_auth( + task_template=_task_template(databricksAuthType="unknown"), + workspace_url="example.cloud.databricks.com", + namespace="project-a", + ) + + +@pytest.mark.asyncio +async def test_m2m_uses_environment_credentials(monkeypatch): + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "example-client") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "example-secret") + auth = build_auth("example.cloud.databricks.com", "oauth_m2m") + + with aioresponses() as mocked: + mocked.post( + "https://example.cloud.databricks.com/oidc/v1/token", + status=200, + payload={"access_token": "example-access-token", "expires_in": 3600}, + ) + async with ClientSession() as session: + token = await auth.get_bearer_token(session) + + assert token == "example-access-token" + + +@pytest.mark.asyncio +async def test_namespace_secret_overrides_environment(monkeypatch): + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "environment-client") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "environment-secret") + auth = build_auth( + "namespace.cloud.databricks.com", + "oauth_m2m", + namespace="project-a", + ) + + def _secret(secret_name, secret_key, namespace): + assert secret_name == DEFAULT_OAUTH_SECRET_NAME + assert namespace == "project-a" + return { + "client_id": "namespace-client", + "client_secret": "namespace-secret", + }[secret_key] + + posted = {} + + async def _capture(session, workspace_url, form): + posted.update(form) + return {"access_token": "namespace-access-token", "expires_in": 3600} + + with patch( + "flytekitplugins.spark.connector.get_secret_from_k8s", + side_effect=_secret, + ), patch( + "flytekitplugins.spark.databricks_auth._post_token", + side_effect=_capture, + ): + async with ClientSession() as session: + await auth.get_bearer_token(session) + + assert posted["client_id"] == "namespace-client" + assert posted["client_secret"] == "namespace-secret" + + +@pytest.mark.asyncio +async def test_m2m_requires_client_id(monkeypatch): + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "example-secret") + auth = build_auth("missing-id.cloud.databricks.com", "oauth_m2m") + + async with ClientSession() as session: + with pytest.raises(DatabricksAuthError, match="requires a client ID"): + await auth.get_bearer_token(session) + + +@pytest.mark.asyncio +async def test_m2m_requires_client_secret(monkeypatch): + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "example-client") + auth = build_auth("missing-secret.cloud.databricks.com", "oauth_m2m") + + async with ClientSession() as session: + with pytest.raises(DatabricksAuthError, match="requires a client secret"): + await auth.get_bearer_token(session) + + +@pytest.mark.asyncio +async def test_token_cache_expires(monkeypatch): + cache = _TokenCache() + now = {"value": 1000.0} + monkeypatch.setattr( + "flytekitplugins.spark.databricks_auth.time.time", + lambda: now["value"], + ) + key = ("example.cloud.databricks.com", "client", "project-a") + + await cache.put(key, "cached-token", expires_in=120) + assert await cache.get(key) == "cached-token" + now["value"] = 1061.0 + assert await cache.get(key) is None + + +@pytest.mark.asyncio +async def test_token_endpoint_fails_fast_on_401(): + async with ClientSession() as session: + with aioresponses() as mocked: + mocked.post( + "https://example.cloud.databricks.com/oidc/v1/token", + status=401, + body="unauthorized", + ) + with pytest.raises(DatabricksAuthError, match="HTTP 401"): + await _post_token( + session, + "example.cloud.databricks.com", + {"grant_type": "client_credentials"}, + ) + + +@pytest.mark.asyncio +async def test_token_endpoint_retries_transient_failure(monkeypatch): + monkeypatch.setattr( + "flytekitplugins.spark.databricks_auth.asyncio.sleep", + AsyncMock(), + ) + async with ClientSession() as session: + with aioresponses() as mocked: + url = "https://retry.cloud.databricks.com/oidc/v1/token" + mocked.post(url, status=503, body="busy") + mocked.post( + url, + status=200, + payload={"access_token": "retried-token", "expires_in": 3600}, + ) + result = await _post_token( + session, + "retry.cloud.databricks.com", + {"grant_type": "client_credentials"}, + ) + + assert result["access_token"] == "retried-token" + + +@pytest.mark.asyncio +async def test_connector_create_persists_m2m_metadata(monkeypatch): + from flytekit.extend.backend.base_agent import AgentRegistry + from flytekitplugins.spark.connector import DATABRICKS_API_ENDPOINT + + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "create-client") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "create-secret") + task_template = _task_template( + databricksInstance="create.cloud.databricks.com", + databricksAuthType="oauth_m2m", + ) + execution_metadata = MagicMock(namespace="project-a") + connector = AgentRegistry.get_agent("spark") + + with patch( + "flytekitplugins.spark.connector._get_databricks_job_spec", + return_value={"run_name": "example"}, + ), patch( + "flytekitplugins.spark.connector.get_secret_from_k8s", + return_value=None, + ): + with aioresponses() as mocked: + mocked.post( + "https://create.cloud.databricks.com/oidc/v1/token", + status=200, + payload={"access_token": "create-access-token", "expires_in": 3600}, + ) + mocked.post( + f"https://create.cloud.databricks.com{DATABRICKS_API_ENDPOINT}/runs/submit", + status=200, + payload={"run_id": 42}, + ) + result = await connector.create( + task_template, + task_execution_metadata=execution_metadata, + ) + + assert result.auth_type == "oauth_m2m" + assert result.auth_token is None + assert result.client_id == "create-client" + assert result.oauth_secret_name == DEFAULT_OAUTH_SECRET_NAME + assert result.namespace == "project-a" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["get", "delete"]) +async def test_connector_refreshes_m2m_once_after_401(monkeypatch, operation): + from flytekit.extend.backend.base_agent import AgentRegistry + from flytekitplugins.spark.connector import DATABRICKS_API_ENDPOINT, DatabricksJobMetadata + + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "refresh-secret") + metadata = DatabricksJobMetadata( + databricks_instance=f"{operation}-refresh.cloud.databricks.com", + run_id="42", + auth_type="oauth_m2m", + client_id="refresh-client", + oauth_secret_name=DEFAULT_OAUTH_SECRET_NAME, + namespace="project-a", + ) + connector = AgentRegistry.get_agent("spark") + issued_tokens = iter(("stale-token", "fresh-token")) + + async def _token_response(session, workspace_url, form): + return {"access_token": next(issued_tokens), "expires_in": 3600} + + if operation == "get": + url = ( + f"https://{metadata.databricks_instance}" + f"{DATABRICKS_API_ENDPOINT}/runs/get?run_id={metadata.run_id}" + ) + else: + url = f"https://{metadata.databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/cancel" + + with patch( + "flytekitplugins.spark.databricks_auth._post_token", + side_effect=_token_response, + ), patch( + "flytekitplugins.spark.connector.get_secret_from_k8s", + return_value=None, + ): + with aioresponses() as mocked: + request = mocked.get if operation == "get" else mocked.post + request(url, status=http.HTTPStatus.UNAUTHORIZED) + request( + url, + status=http.HTTPStatus.OK, + payload={ + "job_id": "1", + "state": {"life_cycle_state": "RUNNING"}, + }, + ) + if operation == "get": + await connector.get(metadata) + else: + await connector.delete(metadata) + + requests = [call for calls in mocked.requests.values() for call in calls] + + assert len(requests) == 2 + assert requests[0].kwargs["headers"]["Authorization"] == "Bearer stale-token" + assert requests[1].kwargs["headers"]["Authorization"] == "Bearer fresh-token" + + +@pytest.mark.asyncio +async def test_connector_does_not_retry_non_401(monkeypatch): + from flytekit.extend.backend.base_agent import AgentRegistry + from flytekitplugins.spark.connector import DATABRICKS_API_ENDPOINT, DatabricksJobMetadata + + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "failure-secret") + metadata = DatabricksJobMetadata( + databricks_instance="failure.cloud.databricks.com", + run_id="42", + auth_type="oauth_m2m", + client_id="failure-client", + ) + connector = AgentRegistry.get_agent("spark") + url = ( + f"https://{metadata.databricks_instance}" + f"{DATABRICKS_API_ENDPOINT}/runs/get?run_id={metadata.run_id}" + ) + + with patch( + "flytekitplugins.spark.databricks_auth._post_token", + return_value={"access_token": "failure-token", "expires_in": 3600}, + ): + with aioresponses() as mocked: + mocked.get(url, status=http.HTTPStatus.FORBIDDEN) + with pytest.raises(RuntimeError, match="Failed to get"): + await connector.get(metadata) + + requests = [call for calls in mocked.requests.values() for call in calls] + + assert len(requests) == 1 diff --git a/plugins/flytekit-spark/tests/test_databricks_token.py b/plugins/flytekit-spark/tests/test_databricks_token.py index 9bed13bb56..ca20a7aa30 100644 --- a/plugins/flytekit-spark/tests/test_databricks_token.py +++ b/plugins/flytekit-spark/tests/test_databricks_token.py @@ -597,6 +597,9 @@ def test_get_custom_includes_token_secret(self): databricks_conf=databricks_conf, databricks_instance="test.cloud.databricks.com", databricks_token_secret="project-x-token", + databricks_auth_type="oauth_m2m", + databricks_client_id="example-client", + databricks_oauth_secret="project-x-oauth", ) ) def my_task(x: int) -> int: @@ -614,6 +617,9 @@ def my_task(x: int) -> int: custom = my_task.get_custom(settings) assert "databricksTokenSecret" in custom assert custom["databricksTokenSecret"] == "project-x-token" + assert custom["databricksAuthType"] == "oauth_m2m" + assert custom["databricksClientId"] == "example-client" + assert custom["databricksOauthSecret"] == "project-x-oauth" def test_get_custom_excludes_token_secret_when_none(self): """get_custom() does NOT include databricksTokenSecret when None.""" From e8ba14e95c570d5e87c827e5af3bd540f01284fd Mon Sep 17 00:00:00 2001 From: Rohit Sharma Date: Thu, 30 Jul 2026 14:10:26 +0100 Subject: [PATCH 3/3] feat(spark): add Databricks OIDC connector authentication Enable connectors to exchange their projected workload JWT for short-lived Databricks tokens without adding namespace discovery or Kubernetes RBAC requirements. Signed-off-by: Rohit Sharma --- plugins/flytekit-spark/README.md | 27 +++ .../flytekitplugins/spark/connector.py | 8 +- .../flytekitplugins/spark/databricks_auth.py | 105 +++++++++- .../flytekitplugins/spark/task.py | 13 +- .../tests/test_databricks_auth.py | 198 ++++++++++++++++++ .../tests/test_databricks_token.py | 4 + 6 files changed, 352 insertions(+), 3 deletions(-) diff --git a/plugins/flytekit-spark/README.md b/plugins/flytekit-spark/README.md index f5609e57ce..40bdf54cda 100644 --- a/plugins/flytekit-spark/README.md +++ b/plugins/flytekit-spark/README.md @@ -51,3 +51,30 @@ stringData: The namespace Secret takes precedence over connector-level credentials. `get` and `delete` operations cache short-lived OAuth tokens and retry once with a refreshed token when the Databricks API returns HTTP 401. + +### OIDC workload identity federation + +The connector can exchange its own projected workload JWT for a short-lived +Databricks token without storing a client secret: + +```yaml +env: + - name: FLYTE_DATABRICKS_AUTH_TYPE + value: oidc_federation + - name: DATABRICKS_CLIENT_ID + value: "" + - name: FLYTE_DATABRICKS_OIDC_TOKEN_FILE + value: /var/run/secrets/databricks/token +``` + +The token file is resolved in this order: + +1. `databricks_oidc_token_file` task override or + `FLYTE_DATABRICKS_OIDC_TOKEN_FILE` +2. `AWS_WEB_IDENTITY_TOKEN_FILE` +3. `/var/run/secrets/databricks/token` + +The connector deployment is responsible for projecting a JWT at one of these +paths and configuring a matching federation policy for the Databricks service +principal. PAT remains the default unless `oidc_federation` is selected +explicitly. diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 6c8252f495..1f23e1ae5d 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/connector.py @@ -42,6 +42,8 @@ class DatabricksJobMetadata(ResourceMeta): auth_type: Optional[str] = None client_id: Optional[str] = None oauth_secret_name: Optional[str] = None + oidc_token_file: Optional[str] = None + oidc_audience: Optional[str] = None namespace: Optional[str] = None @@ -312,6 +314,8 @@ async def create( auth_type=auth.auth_type, client_id=auth.settings.client_id, oauth_secret_name=auth.settings.oauth_secret_name, + oidc_token_file=(auth.settings.oidc_token_file if auth.auth_type == "oidc_federation" else None), + oidc_audience=(auth.settings.oidc_audience if auth.auth_type == "oidc_federation" else None), namespace=namespace, ) @@ -379,13 +383,15 @@ async def _request_with_auth( from .databricks_auth import DatabricksAuthError, build_auth auth = None - if resource_meta.auth_type == "oauth_m2m": + if resource_meta.auth_type in {"oauth_m2m", "oidc_federation"}: auth = build_auth( workspace_url=resource_meta.databricks_instance, auth_type=resource_meta.auth_type, namespace=resource_meta.namespace, client_id=resource_meta.client_id, oauth_secret_name=resource_meta.oauth_secret_name, + oidc_token_file=resource_meta.oidc_token_file, + oidc_audience=resource_meta.oidc_audience, ) token = resource_meta.auth_token diff --git a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py index dd3ed8730d..67b8fd5a82 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py @@ -24,14 +24,19 @@ FLYTE_DATABRICKS_AUTH_TYPE_ENV = "FLYTE_DATABRICKS_AUTH_TYPE" FLYTE_DATABRICKS_OAUTH_SECRET_NAME_ENV = "FLYTE_DATABRICKS_OAUTH_SECRET_NAME" +FLYTE_DATABRICKS_OIDC_TOKEN_FILE_ENV = "FLYTE_DATABRICKS_OIDC_TOKEN_FILE" +FLYTE_DATABRICKS_OIDC_AUDIENCE_ENV = "FLYTE_DATABRICKS_OIDC_AUDIENCE" DATABRICKS_CLIENT_ID_ENV = "DATABRICKS_CLIENT_ID" DATABRICKS_CLIENT_SECRET_ENV = "DATABRICKS_CLIENT_SECRET" +AWS_WEB_IDENTITY_TOKEN_FILE_ENV = "AWS_WEB_IDENTITY_TOKEN_FILE" DEFAULT_OAUTH_SECRET_NAME = "databricks-oauth" +DEFAULT_OIDC_AUDIENCE = "databricks" +DEFAULT_PROJECTED_SA_TOKEN_PATH = "/var/run/secrets/databricks/token" TOKEN_REFRESH_BUFFER_SECONDS = 60 TOKEN_ENDPOINT_MAX_RETRIES = 3 TOKEN_ENDPOINT_BACKOFF_BASE_SECONDS = 0.2 -VALID_AUTH_TYPES = {"pat", "oauth_m2m"} +VALID_AUTH_TYPES = {"pat", "oauth_m2m", "oidc_federation"} class DatabricksAuthError(Exception): @@ -47,6 +52,8 @@ class _Settings: client_id: Optional[str] oauth_secret_name: str token_secret_name: Optional[str] + oidc_token_file: Optional[str] + oidc_audience: str namespace: Optional[str] @staticmethod @@ -74,10 +81,30 @@ def _pick(task_key: str, env_key: Optional[str], default: Optional[str] = None) ) or DEFAULT_OAUTH_SECRET_NAME, token_secret_name=custom.get("databricksTokenSecret"), + oidc_token_file=_pick( + "databricksOidcTokenFile", + FLYTE_DATABRICKS_OIDC_TOKEN_FILE_ENV, + ), + oidc_audience=_pick( + "databricksOidcAudience", + FLYTE_DATABRICKS_OIDC_AUDIENCE_ENV, + DEFAULT_OIDC_AUDIENCE, + ) + or DEFAULT_OIDC_AUDIENCE, namespace=namespace, ) +def _resolve_oidc_token_file(settings: _Settings) -> Optional[str]: + """Resolve the first configured projected JWT file that exists.""" + candidates = ( + settings.oidc_token_file, + os.getenv(AWS_WEB_IDENTITY_TOKEN_FILE_ENV), + DEFAULT_PROJECTED_SA_TOKEN_PATH, + ) + return next((path for path in candidates if path and os.path.exists(path)), None) + + @dataclass class _CachedToken: access_token: str @@ -280,6 +307,74 @@ async def invalidate_cache(self) -> None: await _TOKEN_CACHE.invalidate((self.workspace_url, client_id, self.settings.namespace or "_")) +class OIDCConnectorAuth(DatabricksAuth): + """Exchange the connector workload's projected JWT for a Databricks token.""" + + auth_type = "oidc_federation" + strategy_name = "OIDCConnectorAuth" + + def _client_id(self) -> str: + client_id = self.settings.client_id or os.getenv(DATABRICKS_CLIENT_ID_ENV) + if not client_id: + raise DatabricksAuthError( + "OIDC federation requires a client ID. Configure databricks_client_id " "or DATABRICKS_CLIENT_ID." + ) + return client_id + + def _subject_token_file(self) -> str: + token_file = _resolve_oidc_token_file(self.settings) + if not token_file: + raise DatabricksAuthError( + "OIDC federation requires a projected JWT file. Configure " + "databricks_oidc_token_file, FLYTE_DATABRICKS_OIDC_TOKEN_FILE, " + "or AWS_WEB_IDENTITY_TOKEN_FILE." + ) + return token_file + + def _cache_key(self, client_id: str) -> Tuple[str, str, str]: + return ( + self.workspace_url, + client_id, + f"oidc:{self.settings.oidc_audience}", + ) + + async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] + client_id = self._client_id() + key = self._cache_key(client_id) + cached = await _TOKEN_CACHE.get(key) + if cached: + return cached + + token_file = self._subject_token_file() + try: + with open(token_file) as token_stream: + subject_token = token_stream.read().strip() + except OSError as error: + raise DatabricksAuthError(f"Unable to read projected OIDC token file '{token_file}': {error}") from error + if not subject_token: + raise DatabricksAuthError(f"Projected OIDC token file '{token_file}' is empty") + + payload = await _post_token( + session, + self.workspace_url, + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "scope": "all-apis", + "client_id": client_id, + "subject_token": subject_token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + }, + ) + access_token = payload.get("access_token") + if not access_token: + raise DatabricksAuthError("Databricks OIDC response did not contain access_token") + await _TOKEN_CACHE.put(key, access_token, int(payload.get("expires_in", 3600))) + return access_token + + async def invalidate_cache(self) -> None: + await _TOKEN_CACHE.invalidate(self._cache_key(self._client_id())) + + async def select_auth( task_template: Optional[TaskTemplate], workspace_url: str, @@ -294,6 +389,8 @@ async def select_auth( ) if auth_type == "oauth_m2m": return OAuthM2MAuth(workspace_url, settings) + if auth_type == "oidc_federation": + return OIDCConnectorAuth(workspace_url, settings) return PATAuth(workspace_url, settings) @@ -303,6 +400,8 @@ def build_auth( namespace: Optional[str] = None, client_id: Optional[str] = None, oauth_secret_name: Optional[str] = None, + oidc_token_file: Optional[str] = None, + oidc_audience: Optional[str] = None, ) -> DatabricksAuth: """Rebuild an auth strategy from persisted connector metadata.""" settings = _Settings( @@ -311,12 +410,16 @@ def build_auth( client_id=client_id, oauth_secret_name=oauth_secret_name or DEFAULT_OAUTH_SECRET_NAME, token_secret_name=None, + oidc_token_file=oidc_token_file, + oidc_audience=oidc_audience or DEFAULT_OIDC_AUDIENCE, namespace=namespace, ) if auth_type == "oauth_m2m": return OAuthM2MAuth(workspace_url, settings) if auth_type == "pat": return PATAuth(workspace_url, settings) + if auth_type == "oidc_federation": + return OIDCConnectorAuth(workspace_url, settings) raise DatabricksAuthError( f"Invalid Databricks auth type '{auth_type}'. Expected one of {sorted(VALID_AUTH_TYPES)}." ) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index afea0e1470..8d8506ae07 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -90,11 +90,16 @@ class DatabricksV2(Spark): databricks_token_secret (Optional[str]): Custom name for the K8s secret containing the Databricks token. Defaults to 'databricks-token' if not specified. databricks_auth_type (Optional[str]): Authentication mode. Supported values are - ``"pat"`` and ``"oauth_m2m"``. When unset, PAT remains the default. + ``"pat"``, ``"oauth_m2m"``, and ``"oidc_federation"``. When unset, PAT + remains the default. databricks_client_id (Optional[str]): Databricks service-principal client ID used for OAuth M2M. Falls back to ``DATABRICKS_CLIENT_ID`` on the connector. databricks_oauth_secret (Optional[str]): Name of the namespace K8s secret containing ``client_id`` and ``client_secret``. Defaults to ``databricks-oauth``. + databricks_oidc_token_file (Optional[str]): Path to the connector workload's + projected OIDC JWT. Falls back to ``AWS_WEB_IDENTITY_TOKEN_FILE``. + databricks_oidc_audience (Optional[str]): Audience associated with the projected + JWT. Defaults to ``databricks``. notebook_path (Optional[str]): Path to Databricks notebook (e.g., "/Users/user@example.com/notebook"). notebook_base_parameters (Optional[Dict[str, str]]): Parameters to pass to the notebook. @@ -208,6 +213,8 @@ class DatabricksV2(Spark): databricks_auth_type: Optional[str] = None databricks_client_id: Optional[str] = None databricks_oauth_secret: Optional[str] = None + databricks_oidc_token_file: Optional[str] = None + databricks_oidc_audience: Optional[str] = None notebook_path: Optional[str] = None notebook_base_parameters: Optional[Dict[str, str]] = None @@ -329,6 +336,10 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: custom_dict["databricksClientId"] = cfg.databricks_client_id if cfg.databricks_oauth_secret: custom_dict["databricksOauthSecret"] = cfg.databricks_oauth_secret + if cfg.databricks_oidc_token_file: + custom_dict["databricksOidcTokenFile"] = cfg.databricks_oidc_token_file + if cfg.databricks_oidc_audience: + custom_dict["databricksOidcAudience"] = cfg.databricks_oidc_audience if cfg.notebook_path: custom_dict["notebookPath"] = cfg.notebook_path if cfg.notebook_base_parameters: diff --git a/plugins/flytekit-spark/tests/test_databricks_auth.py b/plugins/flytekit-spark/tests/test_databricks_auth.py index 2695f32247..afaf58996a 100644 --- a/plugins/flytekit-spark/tests/test_databricks_auth.py +++ b/plugins/flytekit-spark/tests/test_databricks_auth.py @@ -8,13 +8,16 @@ from aioresponses import aioresponses from flytekitplugins.spark.databricks_auth import ( + DEFAULT_OIDC_AUDIENCE, DEFAULT_OAUTH_SECRET_NAME, DatabricksAuthError, + OIDCConnectorAuth, OAuthM2MAuth, PATAuth, _TokenCache, _Settings, _post_token, + _resolve_oidc_token_file, build_auth, select_auth, ) @@ -31,8 +34,11 @@ def _clear_auth_environment(monkeypatch): for name in ( "FLYTE_DATABRICKS_AUTH_TYPE", "FLYTE_DATABRICKS_OAUTH_SECRET_NAME", + "FLYTE_DATABRICKS_OIDC_TOKEN_FILE", + "FLYTE_DATABRICKS_OIDC_AUDIENCE", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET", + "AWS_WEB_IDENTITY_TOKEN_FILE", ): monkeypatch.delenv(name, raising=False) @@ -76,6 +82,8 @@ async def test_pat_auth_delegates_to_existing_token_lookup(): client_id=None, oauth_secret_name=DEFAULT_OAUTH_SECRET_NAME, token_secret_name="custom-token", + oidc_token_file=None, + oidc_audience=DEFAULT_OIDC_AUDIENCE, namespace="project-a", ) auth = PATAuth("example.cloud.databricks.com", settings) @@ -381,3 +389,193 @@ async def test_connector_does_not_retry_non_401(monkeypatch): requests = [call for calls in mocked.requests.values() for call in calls] assert len(requests) == 1 + + +def test_oidc_token_file_precedence(tmp_path, monkeypatch): + explicit = tmp_path / "explicit.jwt" + fallback = tmp_path / "fallback.jwt" + explicit.write_text("explicit") + fallback.write_text("fallback") + monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN_FILE", str(fallback)) + settings = _Settings.from_task( + _task_template(databricksOidcTokenFile=str(explicit)), + namespace=None, + ) + + assert _resolve_oidc_token_file(settings) == str(explicit) + + +@pytest.mark.asyncio +async def test_select_auth_uses_explicit_oidc(): + auth = await select_auth( + task_template=_task_template(databricksAuthType="oidc_federation"), + workspace_url="oidc-select.cloud.databricks.com", + namespace=None, + ) + + assert isinstance(auth, OIDCConnectorAuth) + + +@pytest.mark.asyncio +async def test_oidc_exchanges_projected_jwt(tmp_path): + token_file = tmp_path / "workload.jwt" + token_file.write_text("example-subject-token") + auth = build_auth( + "oidc-exchange.cloud.databricks.com", + "oidc_federation", + client_id="oidc-client", + oidc_token_file=str(token_file), + ) + posted = {} + + async def _capture(session, workspace_url, form): + posted.update(form) + return {"access_token": "oidc-access-token", "expires_in": 3600} + + with patch( + "flytekitplugins.spark.databricks_auth._post_token", + side_effect=_capture, + ): + async with ClientSession() as session: + token = await auth.get_bearer_token(session) + + assert token == "oidc-access-token" + assert posted["client_id"] == "oidc-client" + assert posted["subject_token"] == "example-subject-token" + assert posted["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + + +@pytest.mark.asyncio +async def test_oidc_rereads_projected_jwt_after_refresh(tmp_path): + token_file = tmp_path / "rotating.jwt" + token_file.write_text("version-one") + auth = build_auth( + "oidc-refresh.cloud.databricks.com", + "oidc_federation", + client_id="refresh-client", + oidc_token_file=str(token_file), + ) + observed = [] + + async def _capture(session, workspace_url, form): + observed.append(form["subject_token"]) + return {"access_token": f"token-{len(observed)}", "expires_in": 3600} + + with patch( + "flytekitplugins.spark.databricks_auth._post_token", + side_effect=_capture, + ): + async with ClientSession() as session: + await auth.get_bearer_token(session) + await auth.invalidate_cache() + token_file.write_text("version-two") + await auth.get_bearer_token(session) + + assert observed == ["version-one", "version-two"] + + +@pytest.mark.asyncio +async def test_oidc_requires_projected_token_file(): + auth = build_auth( + "oidc-missing-file.cloud.databricks.com", + "oidc_federation", + client_id="oidc-client", + oidc_token_file="/does/not/exist", + ) + + with patch("os.path.exists", return_value=False): + async with ClientSession() as session: + with pytest.raises(DatabricksAuthError, match="requires a projected JWT"): + await auth.get_bearer_token(session) + + +@pytest.mark.asyncio +async def test_oidc_requires_client_id(tmp_path): + token_file = tmp_path / "workload.jwt" + token_file.write_text("example-subject-token") + auth = build_auth( + "oidc-missing-client.cloud.databricks.com", + "oidc_federation", + oidc_token_file=str(token_file), + ) + + async with ClientSession() as session: + with pytest.raises(DatabricksAuthError, match="requires a client ID"): + await auth.get_bearer_token(session) + + +@pytest.mark.asyncio +async def test_oidc_rejects_missing_access_token(tmp_path): + token_file = tmp_path / "workload.jwt" + token_file.write_text("example-subject-token") + auth = build_auth( + "oidc-missing-response.cloud.databricks.com", + "oidc_federation", + client_id="oidc-client", + oidc_token_file=str(token_file), + ) + + with patch( + "flytekitplugins.spark.databricks_auth._post_token", + return_value={"expires_in": 3600}, + ): + async with ClientSession() as session: + with pytest.raises(DatabricksAuthError, match="did not contain access_token"): + await auth.get_bearer_token(session) + + +@pytest.mark.asyncio +async def test_token_endpoint_rejects_invalid_json(): + async with ClientSession() as session: + with aioresponses() as mocked: + mocked.post( + "https://invalid-json.cloud.databricks.com/oidc/v1/token", + status=200, + body="not-json", + ) + with pytest.raises(DatabricksAuthError, match="invalid JSON"): + await _post_token( + session, + "invalid-json.cloud.databricks.com", + {"grant_type": "example"}, + ) + + +@pytest.mark.asyncio +async def test_connector_create_persists_oidc_metadata(tmp_path): + from flytekit.extend.backend.base_agent import AgentRegistry + from flytekitplugins.spark.connector import DATABRICKS_API_ENDPOINT + + token_file = tmp_path / "connector.jwt" + token_file.write_text("connector-subject-token") + task_template = _task_template( + databricksInstance="oidc-create.cloud.databricks.com", + databricksAuthType="oidc_federation", + databricksClientId="connector-client", + databricksOidcTokenFile=str(token_file), + databricksOidcAudience="example-audience", + ) + connector = AgentRegistry.get_agent("spark") + + with patch( + "flytekitplugins.spark.connector._get_databricks_job_spec", + return_value={"run_name": "example"}, + ): + with aioresponses() as mocked: + mocked.post( + "https://oidc-create.cloud.databricks.com/oidc/v1/token", + status=200, + payload={"access_token": "connector-access-token", "expires_in": 3600}, + ) + mocked.post( + f"https://oidc-create.cloud.databricks.com{DATABRICKS_API_ENDPOINT}/runs/submit", + status=200, + payload={"run_id": 42}, + ) + result = await connector.create(task_template) + + assert result.auth_type == "oidc_federation" + assert result.auth_token is None + assert result.client_id == "connector-client" + assert result.oidc_token_file == str(token_file) + assert result.oidc_audience == "example-audience" diff --git a/plugins/flytekit-spark/tests/test_databricks_token.py b/plugins/flytekit-spark/tests/test_databricks_token.py index ca20a7aa30..eeaf41df53 100644 --- a/plugins/flytekit-spark/tests/test_databricks_token.py +++ b/plugins/flytekit-spark/tests/test_databricks_token.py @@ -600,6 +600,8 @@ def test_get_custom_includes_token_secret(self): databricks_auth_type="oauth_m2m", databricks_client_id="example-client", databricks_oauth_secret="project-x-oauth", + databricks_oidc_token_file="/var/run/secrets/example/token", + databricks_oidc_audience="example-audience", ) ) def my_task(x: int) -> int: @@ -620,6 +622,8 @@ def my_task(x: int) -> int: assert custom["databricksAuthType"] == "oauth_m2m" assert custom["databricksClientId"] == "example-client" assert custom["databricksOauthSecret"] == "project-x-oauth" + assert custom["databricksOidcTokenFile"] == "/var/run/secrets/example/token" + assert custom["databricksOidcAudience"] == "example-audience" def test_get_custom_excludes_token_secret_when_none(self): """get_custom() does NOT include databricksTokenSecret when None."""