feat(spark): unify Databricks PAT, OAuth M2M, connector OIDC, and namespace OIDC auth (#3455)

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>

Author: rohitrsh

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: "<service-principal-client-id>"
+  - 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: "<workflow-namespace>"
+type: Opaque
+stringData:
+  client_id: "<service-principal-client-id>"
+  client_secret: "<service-principal-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: "<service-principal-client-id>"
+  - 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: "<workflow-namespace>"
+  labels:
+    flyte.org/databricks-enabled: "true"
+  annotations:
+    flyte.org/databricks-client-id: "<service-principal-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.

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.