diff --git a/plugins/flytekit-spark/README.md b/plugins/flytekit-spark/README.md index 9cc7c7cf9d..47e699fe18 100644 --- a/plugins/flytekit-spark/README.md +++ b/plugins/flytekit-spark/README.md @@ -11,3 +11,108 @@ 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. + +### 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. + +#### Per-namespace ServiceAccount identity + +When OIDC federation is selected, the connector first looks for one +ServiceAccount in the workflow namespace with this configuration: + +```yaml +apiVersion: v1 +kind: ServiceAccount +metadata: + name: databricks-workload + namespace: "" + labels: + flyte.org/databricks-enabled: "true" + annotations: + flyte.org/databricks-client-id: "" + flyte.org/databricks-audience: "databricks" +``` + +If found, the connector creates a short-lived JWT for that ServiceAccount +through the Kubernetes TokenRequest API and exchanges it for a Databricks +token. If no matching ServiceAccount exists, connector-identity OIDC remains +the fallback. + +The connector ServiceAccount needs these additional Kubernetes permissions: + +```yaml +rules: + - apiGroups: [""] + resources: ["serviceaccounts"] + verbs: ["get", "list"] + - apiGroups: [""] + resources: ["serviceaccounts/token"] + verbs: ["create"] +``` + +Configure exactly one Databricks-enabled ServiceAccount per workflow +namespace. Multiple matching ServiceAccounts are treated as an error. diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 06b02048d1..697e3ea743 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/connector.py @@ -29,9 +29,23 @@ @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 + oidc_token_file: Optional[str] = None + oidc_service_account: Optional[str] = None + oidc_audience: Optional[str] = None + namespace: Optional[str] = None def _configure_serverless(databricks_job: dict, envs: dict) -> str: @@ -254,6 +268,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, @@ -262,6 +279,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,30 +291,42 @@ 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, + 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" 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: raise RuntimeError(f"Failed to create databricks job with error: {response}") logger.info(f"Successfully created Databricks job with run_id: {response['run_id']}") + discovered = getattr(auth, "discovered", None) + persisted_client_id = discovered.client_id if discovered is not None else auth.settings.client_id + persisted_service_account = discovered.service_account if discovered is not None else None + persisted_audience = ( + discovered.audience + if discovered is not None + else (auth.settings.oidc_audience if auth.auth_type == "oidc_federation" else None) + ) 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=persisted_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_service_account=persisted_service_account, + oidc_audience=persisted_audience, + namespace=namespace, ) async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource: @@ -304,14 +335,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 = "" @@ -339,16 +370,69 @@ 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 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, + oidc_service_account=resource_meta.oidc_service_account, + ) + + 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): @@ -364,6 +448,32 @@ def __init__(self): super(DatabricksConnector, self).__init__(task_type_name="databricks", metadata_type=DatabricksJobMetadata) +def list_serviceaccounts_in_k8s(namespace: str, label_selector: Optional[str] = None) -> list: + """List labelled ServiceAccounts in a workflow namespace.""" + try: + from kubernetes import client, config + + try: + config.load_incluster_config() + except config.ConfigException: + config.load_kube_config() + + arguments = {"namespace": namespace} + if label_selector: + arguments["label_selector"] = label_selector + response = client.CoreV1Api().list_namespaced_service_account(**arguments) + return list(response.items or []) + except ImportError: + logger.warning("Kubernetes Python client is unavailable; skipping namespace " "ServiceAccount discovery") + except Exception as error: + logger.warning( + "Unable to discover ServiceAccounts in namespace '%s': %s", + namespace, + error, + ) + return [] + + def get_secret_from_k8s(secret_name: str, secret_key: str, namespace: str) -> Optional[str]: """Read a secret from Kubernetes using the Kubernetes Python client. 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..c172d49338 --- /dev/null +++ b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py @@ -0,0 +1,688 @@ +"""Authentication strategies for the Databricks connector. + +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 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" +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" + +LABEL_DATABRICKS_ENABLED = "flyte.org/databricks-enabled" +ANNOTATION_DATABRICKS_CLIENT_ID = "flyte.org/databricks-client-id" +ANNOTATION_DATABRICKS_AUDIENCE = "flyte.org/databricks-audience" +DATABRICKS_ENABLED_LABEL_SELECTOR = f"{LABEL_DATABRICKS_ENABLED}=true" + +DEFAULT_OAUTH_SECRET_NAME = "databricks-oauth" +DEFAULT_OIDC_AUDIENCE = "databricks" +DEFAULT_PROJECTED_SA_TOKEN_PATH = "/var/run/secrets/databricks/token" +NAMESPACE_DISCOVERY_CACHE_TTL_SECONDS = 300 +TOKEN_REFRESH_BUFFER_SECONDS = 60 +TOKEN_ENDPOINT_MAX_RETRIES = 3 +TOKEN_ENDPOINT_BACKOFF_BASE_SECONDS = 0.2 +VALID_AUTH_TYPES = {"pat", "oauth_m2m", "oidc_federation"} + + +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] + oidc_token_file: Optional[str] + oidc_audience: str + namespace: Optional[str] + + @staticmethod + def from_task(task_template: Optional[TaskTemplate], namespace: Optional[str]) -> "_Settings": + 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"), + 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 + 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() + + +@dataclass +class _DiscoveredOIDCConfig: + service_account: str + client_id: str + audience: str + + +@dataclass +class _CachedDiscovery: + config: Optional[_DiscoveredOIDCConfig] + expires_at: float + + +class _NamespaceDiscoveryCache: + """Cache namespace ServiceAccount discovery results, including misses.""" + + def __init__(self) -> None: + self._store: Dict[str, _CachedDiscovery] = {} + self._lock = asyncio.Lock() + + async def get(self, namespace: str) -> Tuple[bool, Optional[_DiscoveredOIDCConfig]]: + async with self._lock: + entry = self._store.get(namespace) + if entry is None or entry.expires_at <= time.time(): + self._store.pop(namespace, None) + return False, None + return True, entry.config + + async def put( + self, + namespace: str, + config: Optional[_DiscoveredOIDCConfig], + ) -> None: + async with self._lock: + self._store[namespace] = _CachedDiscovery( + config=config, + expires_at=time.time() + NAMESPACE_DISCOVERY_CACHE_TTL_SECONDS, + ) + + async def invalidate(self, namespace: str) -> None: + async with self._lock: + self._store.pop(namespace, None) + + +_NAMESPACE_DISCOVERY_CACHE = _NamespaceDiscoveryCache() + + +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}" + ) + + +def _discover_namespace_oidc_sa_sync( + namespace: str, + default_audience: str, +) -> Optional[_DiscoveredOIDCConfig]: + """Discover exactly one labelled and annotated ServiceAccount.""" + from .connector import list_serviceaccounts_in_k8s + + service_accounts = list_serviceaccounts_in_k8s( + namespace=namespace, + label_selector=DATABRICKS_ENABLED_LABEL_SELECTOR, + ) + matches = [] + for service_account in service_accounts: + metadata = service_account.metadata + annotations = (metadata.annotations or {}) if metadata else {} + client_id = (annotations.get(ANNOTATION_DATABRICKS_CLIENT_ID) or "").strip() + if not client_id: + continue + audience = (annotations.get(ANNOTATION_DATABRICKS_AUDIENCE) or "").strip() or default_audience + matches.append( + _DiscoveredOIDCConfig( + service_account=metadata.name, + client_id=client_id, + audience=audience, + ) + ) + + if not matches: + return None + if len(matches) > 1: + names = ", ".join(sorted(match.service_account for match in matches)) + raise DatabricksAuthError( + f"Multiple Databricks-enabled ServiceAccounts found in namespace " + f"'{namespace}': {names}. Configure exactly one." + ) + return matches[0] + + +async def _discover_namespace_oidc_sa( + namespace: str, + default_audience: str, +) -> Optional[_DiscoveredOIDCConfig]: + cached, config = await _NAMESPACE_DISCOVERY_CACHE.get(namespace) + if cached: + return config + loop = asyncio.get_running_loop() + config = await loop.run_in_executor( + None, + _discover_namespace_oidc_sa_sync, + namespace, + default_audience, + ) + await _NAMESPACE_DISCOVERY_CACHE.put(namespace, config) + return config + + +def _request_service_account_token( + namespace: str, + service_account: str, + audience: str, + expiration_seconds: int = 3600, +) -> str: + """Mint a projected JWT through the Kubernetes TokenRequest API.""" + try: + from kubernetes import client, config + from kubernetes.client.exceptions import ApiException + except ImportError as error: + raise DatabricksAuthError( + "Kubernetes Python client is required for namespace ServiceAccount federation" + ) from error + + try: + config.load_incluster_config() + except config.ConfigException: + try: + config.load_kube_config() + except Exception as error: + raise DatabricksAuthError(f"Unable to load Kubernetes configuration: {error}") from error + + request = client.AuthenticationV1TokenRequest( + spec=client.V1TokenRequestSpec( + audiences=[audience], + expiration_seconds=expiration_seconds, + ) + ) + try: + response = client.CoreV1Api().create_namespaced_service_account_token( + name=service_account, + namespace=namespace, + body=request, + ) + except ApiException as error: + if error.status == 404: + raise DatabricksAuthError(f"ServiceAccount '{namespace}/{service_account}' was not found") from error + if error.status == 403: + raise DatabricksAuthError( + f"Connector is forbidden from creating a token for " f"'{namespace}/{service_account}'" + ) from error + raise DatabricksAuthError(f"TokenRequest failed for '{namespace}/{service_account}': {error}") from error + + token = response.status.token + if not token: + raise DatabricksAuthError(f"TokenRequest for '{namespace}/{service_account}' returned an empty token") + return token + + +class DatabricksAuth(ABC): + """Interface for obtaining a bearer token for Databricks API calls.""" + + auth_type = "unknown" + strategy_name = "DatabricksAuth" + + 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 ( + 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, + ) + + +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 "_")) + + +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())) + + +class OIDCNamespaceServiceAccountAuth(DatabricksAuth): + """Federate as a ServiceAccount discovered in the workflow namespace.""" + + auth_type = "oidc_federation" + strategy_name = "OIDCNamespaceServiceAccountAuth" + + def __init__( + self, + workspace_url: str, + settings: _Settings, + discovered: _DiscoveredOIDCConfig, + ): + super().__init__(workspace_url, settings) + self.discovered = discovered + + @property + def cache_key(self) -> Tuple[str, str, str]: + return ( + self.workspace_url, + self.discovered.client_id, + ( + f"serviceaccount:{self.settings.namespace}:" + f"{self.discovered.service_account}:{self.discovered.audience}" + ), + ) + + async def get_bearer_token(self, session: "aiohttp.ClientSession") -> str: # type: ignore[name-defined] + if not self.settings.namespace: + raise DatabricksAuthError("Namespace ServiceAccount federation requires a workflow namespace") + + cached = await _TOKEN_CACHE.get(self.cache_key) + if cached: + return cached + + loop = asyncio.get_running_loop() + subject_token = await loop.run_in_executor( + None, + _request_service_account_token, + self.settings.namespace, + self.discovered.service_account, + self.discovered.audience, + ) + payload = await _post_token( + session, + self.workspace_url, + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "scope": "all-apis", + "client_id": self.discovered.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( + self.cache_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) + if self.settings.namespace: + await _NAMESPACE_DISCOVERY_CACHE.invalidate(self.settings.namespace) + + +async def select_auth( + task_template: Optional[TaskTemplate], + workspace_url: str, + namespace: Optional[str], +) -> DatabricksAuth: + """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) + if auth_type == "oidc_federation": + if namespace: + discovered = await _discover_namespace_oidc_sa( + namespace, + settings.oidc_audience, + ) + if discovered: + return OIDCNamespaceServiceAccountAuth( + workspace_url, + settings, + discovered, + ) + return OIDCConnectorAuth(workspace_url, settings) + return PATAuth(workspace_url, settings) + + +def build_auth( + workspace_url: str, + auth_type: str, + 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, + oidc_service_account: 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, + 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": + if oidc_service_account: + if not namespace or not client_id: + raise DatabricksAuthError( + "Rebuilding namespace ServiceAccount federation requires " "namespace and client_id metadata" + ) + return OIDCNamespaceServiceAccountAuth( + workspace_url, + settings, + _DiscoveredOIDCConfig( + service_account=oidc_service_account, + client_id=client_id, + audience=oidc_audience or DEFAULT_OIDC_AUDIENCE, + ), + ) + return OIDCConnectorAuth(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..8d8506ae07 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -89,6 +89,17 @@ 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"``, ``"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. @@ -199,6 +210,11 @@ 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 + 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 @@ -314,6 +330,16 @@ 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.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_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 new file mode 100644 index 0000000000..cee04189b1 --- /dev/null +++ b/plugins/flytekit-spark/tests/test_databricks_auth.py @@ -0,0 +1,870 @@ +"""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_OIDC_AUDIENCE, + DEFAULT_OAUTH_SECRET_NAME, + DATABRICKS_ENABLED_LABEL_SELECTOR, + DatabricksAuthError, + OIDCConnectorAuth, + OIDCNamespaceServiceAccountAuth, + OAuthM2MAuth, + PATAuth, + _DiscoveredOIDCConfig, + _NAMESPACE_DISCOVERY_CACHE, + _TokenCache, + _Settings, + _discover_namespace_oidc_sa, + _discover_namespace_oidc_sa_sync, + _post_token, + _request_service_account_token, + _resolve_oidc_token_file, + build_auth, + select_auth, +) + + +def _task_template(**custom): + task_template = MagicMock() + task_template.custom = custom + return task_template + + +def _service_account(name, client_id=None, audience=None): + annotations = {} + if client_id is not None: + annotations["flyte.org/databricks-client-id"] = client_id + if audience is not None: + annotations["flyte.org/databricks-audience"] = audience + service_account = MagicMock() + service_account.metadata.name = name + service_account.metadata.annotations = annotations or None + return service_account + + +@pytest.fixture(autouse=True) +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) + + +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" + + +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, + workspace_url="example.cloud.databricks.com", + 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, + auth_type="pat", + 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) + + 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", + ) + + +@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 + + +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" + + +@pytest.mark.asyncio +async def test_select_oidc_uses_discovered_namespace_service_account(): + discovered = _DiscoveredOIDCConfig( + service_account="databricks-workload", + client_id="namespace-client", + audience="databricks", + ) + with patch( + "flytekitplugins.spark.databricks_auth._discover_namespace_oidc_sa", + AsyncMock(return_value=discovered), + ): + auth = await select_auth( + task_template=_task_template(databricksAuthType="oidc_federation"), + workspace_url="namespace-select.cloud.databricks.com", + namespace="project-namespace-select", + ) + + assert isinstance(auth, OIDCNamespaceServiceAccountAuth) + assert auth.discovered == discovered + + +@pytest.mark.asyncio +async def test_select_oidc_falls_back_to_connector_identity(): + with patch( + "flytekitplugins.spark.databricks_auth._discover_namespace_oidc_sa", + AsyncMock(return_value=None), + ): + auth = await select_auth( + task_template=_task_template(databricksAuthType="oidc_federation"), + workspace_url="namespace-fallback.cloud.databricks.com", + namespace="project-namespace-fallback", + ) + + assert isinstance(auth, OIDCConnectorAuth) + + +def test_namespace_discovery_returns_none_without_match(): + with patch( + "flytekitplugins.spark.connector.list_serviceaccounts_in_k8s", + return_value=[], + ) as list_serviceaccounts: + result = _discover_namespace_oidc_sa_sync("project-empty", "databricks") + + assert result is None + list_serviceaccounts.assert_called_once_with( + namespace="project-empty", + label_selector=DATABRICKS_ENABLED_LABEL_SELECTOR, + ) + + +def test_namespace_discovery_reads_annotations(): + service_account = _service_account( + "databricks-workload", + client_id="namespace-client", + audience="custom-audience", + ) + with patch( + "flytekitplugins.spark.connector.list_serviceaccounts_in_k8s", + return_value=[service_account], + ): + result = _discover_namespace_oidc_sa_sync("project-one", "databricks") + + assert result == _DiscoveredOIDCConfig( + service_account="databricks-workload", + client_id="namespace-client", + audience="custom-audience", + ) + + +def test_namespace_discovery_uses_default_audience(): + service_account = _service_account( + "databricks-workload", + client_id="namespace-client", + ) + with patch( + "flytekitplugins.spark.connector.list_serviceaccounts_in_k8s", + return_value=[service_account], + ): + result = _discover_namespace_oidc_sa_sync("project-default-audience", "default-audience") + + assert result.audience == "default-audience" + + +def test_namespace_discovery_skips_missing_client_id(): + service_account = _service_account("databricks-workload") + with patch( + "flytekitplugins.spark.connector.list_serviceaccounts_in_k8s", + return_value=[service_account], + ): + result = _discover_namespace_oidc_sa_sync("project-misconfigured", "databricks") + + assert result is None + + +def test_namespace_discovery_rejects_multiple_matches(): + service_accounts = [ + _service_account("first", client_id="first-client"), + _service_account("second", client_id="second-client"), + ] + with patch( + "flytekitplugins.spark.connector.list_serviceaccounts_in_k8s", + return_value=service_accounts, + ): + with pytest.raises(DatabricksAuthError, match="Multiple.*first, second"): + _discover_namespace_oidc_sa_sync("project-multiple", "databricks") + + +@pytest.mark.asyncio +async def test_namespace_discovery_cache_avoids_repeated_lookup(): + namespace = "project-cache-hit" + discovered = _DiscoveredOIDCConfig( + service_account="databricks-workload", + client_id="namespace-client", + audience="databricks", + ) + await _NAMESPACE_DISCOVERY_CACHE.invalidate(namespace) + + with patch( + "flytekitplugins.spark.databricks_auth._discover_namespace_oidc_sa_sync", + return_value=discovered, + ) as discover: + first = await _discover_namespace_oidc_sa(namespace, "databricks") + second = await _discover_namespace_oidc_sa(namespace, "databricks") + + assert first == second == discovered + assert discover.call_count == 1 + + +def test_token_request_returns_service_account_jwt(): + response = MagicMock() + response.status.token = "namespace-subject-token" + with patch("kubernetes.config.load_incluster_config"), patch( + "kubernetes.client.CoreV1Api" + ) as api: + api.return_value.create_namespaced_service_account_token.return_value = response + token = _request_service_account_token( + "project-token", + "databricks-workload", + "databricks", + ) + + assert token == "namespace-subject-token" + + +@pytest.mark.parametrize( + ("status", "message"), + [(403, "forbidden"), (404, "was not found")], +) +def test_token_request_translates_kubernetes_errors(status, message): + from kubernetes.client.exceptions import ApiException + + with patch("kubernetes.config.load_incluster_config"), patch( + "kubernetes.client.CoreV1Api" + ) as api: + api.return_value.create_namespaced_service_account_token.side_effect = ApiException( + status=status + ) + with pytest.raises(DatabricksAuthError, match=message): + _request_service_account_token( + "project-token-error", + "databricks-workload", + "databricks", + ) + + +@pytest.mark.asyncio +async def test_namespace_service_account_exchanges_minted_jwt(): + auth = build_auth( + "namespace-exchange.cloud.databricks.com", + "oidc_federation", + namespace="project-exchange", + client_id="namespace-client", + oidc_service_account="databricks-workload", + oidc_audience="databricks", + ) + posted = {} + + async def _capture(session, workspace_url, form): + posted.update(form) + return {"access_token": "namespace-access-token", "expires_in": 3600} + + with patch( + "flytekitplugins.spark.databricks_auth._request_service_account_token", + return_value="namespace-subject-token", + ), patch( + "flytekitplugins.spark.databricks_auth._post_token", + side_effect=_capture, + ): + async with ClientSession() as session: + token = await auth.get_bearer_token(session) + + assert token == "namespace-access-token" + assert posted["subject_token"] == "namespace-subject-token" + assert posted["client_id"] == "namespace-client" + + +@pytest.mark.asyncio +async def test_namespace_auth_invalidation_clears_discovery_cache(): + namespace = "project-cache-invalidation" + discovered = _DiscoveredOIDCConfig( + service_account="databricks-workload", + client_id="namespace-client", + audience="databricks", + ) + auth = OIDCNamespaceServiceAccountAuth( + "namespace-cache.cloud.databricks.com", + _Settings.from_task(None, namespace), + discovered, + ) + await _NAMESPACE_DISCOVERY_CACHE.put(namespace, discovered) + + await auth.invalidate_cache() + is_cached, _ = await _NAMESPACE_DISCOVERY_CACHE.get(namespace) + + assert is_cached is False + + +@pytest.mark.asyncio +async def test_connector_create_persists_namespace_oidc_metadata(): + from flytekit.extend.backend.base_agent import AgentRegistry + from flytekitplugins.spark.connector import DATABRICKS_API_ENDPOINT + + discovered = _DiscoveredOIDCConfig( + service_account="databricks-workload", + client_id="namespace-client", + audience="namespace-audience", + ) + task_template = _task_template( + databricksInstance="namespace-create.cloud.databricks.com", + databricksAuthType="oidc_federation", + ) + execution_metadata = MagicMock(namespace="project-create-namespace") + connector = AgentRegistry.get_agent("spark") + + with patch( + "flytekitplugins.spark.connector._get_databricks_job_spec", + return_value={"run_name": "example"}, + ), patch( + "flytekitplugins.spark.databricks_auth._discover_namespace_oidc_sa", + AsyncMock(return_value=discovered), + ), patch( + "flytekitplugins.spark.databricks_auth._request_service_account_token", + return_value="namespace-subject-token", + ): + with aioresponses() as mocked: + mocked.post( + "https://namespace-create.cloud.databricks.com/oidc/v1/token", + status=200, + payload={"access_token": "namespace-access-token", "expires_in": 3600}, + ) + mocked.post( + ( + "https://namespace-create.cloud.databricks.com" + f"{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 == "oidc_federation" + assert result.client_id == "namespace-client" + assert result.oidc_service_account == "databricks-workload" + assert result.oidc_audience == "namespace-audience" + assert result.namespace == "project-create-namespace" diff --git a/plugins/flytekit-spark/tests/test_databricks_token.py b/plugins/flytekit-spark/tests/test_databricks_token.py index 9bed13bb56..eeaf41df53 100644 --- a/plugins/flytekit-spark/tests/test_databricks_token.py +++ b/plugins/flytekit-spark/tests/test_databricks_token.py @@ -597,6 +597,11 @@ 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", + databricks_oidc_token_file="/var/run/secrets/example/token", + databricks_oidc_audience="example-audience", ) ) def my_task(x: int) -> int: @@ -614,6 +619,11 @@ 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" + 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."""