feat(spark): add namespace ServiceAccount OIDC auth

Allow each workflow namespace to select a Databricks service principal through an annotated ServiceAccount while retaining connector identity as the OIDC fallback.

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

Author: rohitrsh

plugins/flytekit-spark/README.md

@@ -78,3 +78,41 @@ 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

@@ -43,6 +43,7 @@ class DatabricksJobMetadata(ResourceMeta):
     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
 
@@ -307,15 +308,24 @@ async def create(
                     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 if auth.auth_type == "pat" else None,
             auth_type=auth.auth_type,
-            client_id=auth.settings.client_id,
+            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_audience=(auth.settings.oidc_audience if auth.auth_type == "oidc_federation" else None),
+            oidc_service_account=persisted_service_account,
+            oidc_audience=persisted_audience,
             namespace=namespace,
         )
 
@@ -392,6 +402,7 @@ async def _request_with_auth(
                 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
@@ -437,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.