diff --git a/plugins/flytekit-spark/README.md b/plugins/flytekit-spark/README.md index 9cc7c7cf9d..40bdf54cda 100644 --- a/plugins/flytekit-spark/README.md +++ b/plugins/flytekit-spark/README.md @@ -11,3 +11,70 @@ 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. diff --git a/plugins/flytekit-spark/flytekitplugins/spark/connector.py b/plugins/flytekit-spark/flytekitplugins/spark/connector.py index 06b02048d1..1f23e1ae5d 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/connector.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/connector.py @@ -29,9 +29,22 @@ @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_audience: Optional[str] = None + namespace: Optional[str] = None def _configure_serverless(databricks_job: dict, envs: dict) -> str: @@ -254,6 +267,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 +278,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 +290,17 @@ 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: @@ -295,7 +308,15 @@ 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, + 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, ) async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource: @@ -304,14 +325,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 +360,68 @@ 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, + ) + + 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 new file mode 100644 index 0000000000..67b8fd5a82 --- /dev/null +++ b/plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py @@ -0,0 +1,434 @@ +"""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" + +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", "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() + + +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, 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())) + + +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": + 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, +) -> 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": + 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..afaf58996a --- /dev/null +++ b/plugins/flytekit-spark/tests/test_databricks_auth.py @@ -0,0 +1,581 @@ +"""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, + DatabricksAuthError, + OIDCConnectorAuth, + OAuthM2MAuth, + PATAuth, + _TokenCache, + _Settings, + _post_token, + _resolve_oidc_token_file, + build_auth, + select_auth, +) + + +def _task_template(**custom): + task_template = MagicMock() + task_template.custom = 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", + "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" 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."""