feat(spark): add Databricks OIDC connector authentication

Enable connectors to exchange their projected workload JWT for short-lived Databricks tokens without adding namespace discovery or Kubernetes RBAC requirements.

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

Author: rohitrsh

plugins/flytekit-spark/README.md

@@ -51,3 +51,30 @@ stringData:
 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.

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

@@ -42,6 +42,8 @@ class DatabricksJobMetadata(ResourceMeta):
     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
 
 
@@ -312,6 +314,8 @@ async def create(
             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,
         )
 
@@ -379,13 +383,15 @@ async def _request_with_auth(
         from .databricks_auth import DatabricksAuthError, build_auth
 
         auth = None
-        if resource_meta.auth_type == "oauth_m2m":
+        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

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

@@ -24,14 +24,19 @@
 
 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"}
+VALID_AUTH_TYPES = {"pat", "oauth_m2m", "oidc_federation"}
 
 
 class DatabricksAuthError(Exception):
@@ -47,6 +52,8 @@ class _Settings:
     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
@@ -74,10 +81,30 @@ def _pick(task_key: str, env_key: Optional[str], default: Optional[str] = None)
             )
             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
@@ -280,6 +307,74 @@ async def invalidate_cache(self) -> None:
         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,
@@ -294,6 +389,8 @@ async def select_auth(
         )
     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)
 
 
@@ -303,6 +400,8 @@ def build_auth(
     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(
@@ -311,12 +410,16 @@ def build_auth(
         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)}."
     )

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

@@ -90,11 +90,16 @@ class DatabricksV2(Spark):
         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"`` and ``"oauth_m2m"``. When unset, PAT remains the default.
+            ``"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.
@@ -208,6 +213,8 @@ class DatabricksV2(Spark):
     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
 
@@ -329,6 +336,10 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
                 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: