feat(spark): add Databricks OAuth M2M authentication

Allow connectors and individual tasks to opt into short-lived service-principal tokens while keeping PAT as the unchanged default authentication mode.

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

Author: rohitrsh

plugins/flytekit-spark/README.md

@@ -11,3 +11,43 @@ pip install flytekitplugins-spark
 To configure Spark in the Flyte deployment's backend, follow [Step 1](https://docs.flyte.org/en/latest/deployment/plugins/k8s/index.html#deployment-plugin-setup-k8s), [2](https://docs.flyte.org/en/latest/flytesnacks/examples/k8s_spark_plugin/index.html).
 
 All [examples](https://docs.flyte.org/en/latest/flytesnacks/examples/k8s_spark_plugin/index.html) showcasing execution of Spark jobs using the plugin can be found in the documentation.
+
+## Databricks authentication
+
+The Databricks connector uses PAT authentication by default, preserving the
+existing `databricks-token` namespace Secret and
+`FLYTE_DATABRICKS_ACCESS_TOKEN` fallback.
+
+OAuth machine-to-machine (M2M) authentication can be enabled on the connector:
+
+```yaml
+env:
+  - name: FLYTE_DATABRICKS_AUTH_TYPE
+    value: oauth_m2m
+  - name: DATABRICKS_CLIENT_ID
+    value: "<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.

plugins/flytekit-spark/flytekitplugins/spark/connector.py

@@ -29,9 +29,20 @@
 
 @dataclass
 class DatabricksJobMetadata(ResourceMeta):
+    """Metadata persisted for a Databricks run.
+
+    OAuth metadata allows ``get`` and ``delete`` to obtain fresh short-lived
+    tokens. ``auth_token`` remains populated for PAT jobs and for metadata
+    written by older connector versions.
+    """
+
     databricks_instance: str
     run_id: str
-    auth_token: Optional[str] = None  # Store auth token for get/delete operations
+    auth_token: Optional[str] = None
+    auth_type: Optional[str] = None
+    client_id: Optional[str] = None
+    oauth_secret_name: Optional[str] = None
+    namespace: Optional[str] = None
 
 
 def _configure_serverless(databricks_job: dict, envs: dict) -> str:
@@ -254,6 +265,9 @@ class DatabricksConnector(AsyncConnectorBase):
 
     def __init__(self):
         super().__init__(task_type_name="spark", metadata_type=DatabricksJobMetadata)
+        from .databricks_auth import validate_connector_config
+
+        validate_connector_config()
 
     async def create(
         self,
@@ -275,7 +289,11 @@ async def create(
             )
 
         namespace = task_execution_metadata.namespace if task_execution_metadata else None
-        auth = await select_auth(task_template=task_template, namespace=namespace)
+        auth = await select_auth(
+            task_template=task_template,
+            workspace_url=databricks_instance,
+            namespace=namespace,
+        )
         logger.info("Databricks auth resolved: %s", auth.describe())
         databricks_url = f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/submit"
 
@@ -288,7 +306,13 @@ async def create(
 
         logger.info(f"Successfully created Databricks job with run_id: {response['run_id']}")
         return DatabricksJobMetadata(
-            databricks_instance=databricks_instance, run_id=str(response["run_id"]), auth_token=auth_token
+            databricks_instance=databricks_instance,
+            run_id=str(response["run_id"]),
+            auth_token=auth_token if auth.auth_type == "pat" else None,
+            auth_type=auth.auth_type,
+            client_id=auth.settings.client_id,
+            oauth_secret_name=auth.settings.oauth_secret_name,
+            namespace=namespace,
         )
 
     async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource:
@@ -297,14 +321,14 @@ async def get(self, resource_meta: DatabricksJobMetadata, **kwargs) -> Resource:
             f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/get?run_id={resource_meta.run_id}"
         )
 
-        # Use the stored auth token if available, otherwise fall back to default
-        headers = get_header(auth_token=resource_meta.auth_token)
-
         async with aiohttp.ClientSession() as session:
-            async with session.get(databricks_url, headers=headers) as resp:
-                if resp.status != http.HTTPStatus.OK:
-                    raise RuntimeError(f"Failed to get databricks job {resource_meta.run_id} with error: {resp.reason}")
-                response = await resp.json()
+            response = await self._request_with_auth(
+                session=session,
+                method="GET",
+                url=databricks_url,
+                resource_meta=resource_meta,
+                action_label=f"get databricks job {resource_meta.run_id}",
+            )
 
         cur_phase = TaskExecution.UNDEFINED
         message = ""
@@ -332,16 +356,66 @@ async def delete(self, resource_meta: DatabricksJobMetadata, **kwargs):
         databricks_url = f"https://{resource_meta.databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/cancel"
         data = json.dumps({"run_id": resource_meta.run_id})
 
-        # Use the stored auth token if available, otherwise fall back to default
-        headers = get_header(auth_token=resource_meta.auth_token)
-
         async with aiohttp.ClientSession() as session:
-            async with session.post(databricks_url, headers=headers, data=data) as resp:
-                if resp.status != http.HTTPStatus.OK:
-                    raise RuntimeError(
-                        f"Failed to cancel databricks job {resource_meta.run_id} with error: {resp.reason}"
-                    )
-                await resp.json()
+            await self._request_with_auth(
+                session=session,
+                method="POST",
+                url=databricks_url,
+                resource_meta=resource_meta,
+                data=data,
+                action_label=f"cancel databricks job {resource_meta.run_id}",
+            )
+
+    async def _request_with_auth(
+        self,
+        session: "aiohttp.ClientSession",  # type: ignore[name-defined]
+        method: str,
+        url: str,
+        resource_meta: DatabricksJobMetadata,
+        action_label: str,
+        data: Optional[str] = None,
+    ) -> dict:
+        """Call the Jobs API and retry once after refreshing OAuth on 401."""
+        from .databricks_auth import DatabricksAuthError, build_auth
+
+        auth = None
+        if resource_meta.auth_type == "oauth_m2m":
+            auth = build_auth(
+                workspace_url=resource_meta.databricks_instance,
+                auth_type=resource_meta.auth_type,
+                namespace=resource_meta.namespace,
+                client_id=resource_meta.client_id,
+                oauth_secret_name=resource_meta.oauth_secret_name,
+            )
+
+        token = resource_meta.auth_token
+        if auth is not None:
+            try:
+                token = await auth.get_bearer_token(session)
+            except DatabricksAuthError as error:
+                raise RuntimeError(f"Failed to {action_label}: could not obtain Databricks auth: {error}") from error
+
+        def _request(bearer: Optional[str]):
+            headers = get_header(auth_token=bearer)
+            if method.upper() == "GET":
+                return session.get(url, headers=headers)
+            return session.post(url, headers=headers, data=data)
+
+        async with _request(token) as response:
+            if response.status == http.HTTPStatus.UNAUTHORIZED and auth is not None:
+                await auth.invalidate_cache()
+                try:
+                    refreshed_token = await auth.get_bearer_token(session)
+                except DatabricksAuthError as error:
+                    raise RuntimeError(f"Failed to {action_label}: auth refresh failed after 401: {error}") from error
+                async with _request(refreshed_token) as retry_response:
+                    if retry_response.status != http.HTTPStatus.OK:
+                        raise RuntimeError(f"Failed to {action_label} with error: {retry_response.reason}")
+                    return await retry_response.json()
+
+            if response.status != http.HTTPStatus.OK:
+                raise RuntimeError(f"Failed to {action_label} with error: {response.reason}")
+            return await response.json()
 
 
 class DatabricksConnectorV2(DatabricksConnector):