databricks-sdk 0.53.0__py3-none-any.whl → 0.55.0__py3-none-any.whl

databricks/sdk/oidc.py

@@ -0,0 +1,206 @@
+"""
+Package oidc provides utilities for working with OIDC ID tokens.
+
+This package is experimental and subject to change.
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from . import oauth
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class IdToken:
+    """Represents an OIDC ID token that can be exchanged for a Databricks access token.
+
+    Parameters
+    ----------
+    jwt : str
+        The signed JWT token string.
+    """
+
+    jwt: str
+
+
+class IdTokenSource(ABC):
+    """Abstract base class representing anything that returns an IDToken.
+
+    This class defines the interface for token sources that can provide OIDC ID tokens.
+    """
+
+    @abstractmethod
+    def id_token(self) -> IdToken:
+        """Get an ID token.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        Exception
+            Implementation specific exceptions.
+        """
+
+
+class EnvIdTokenSource(IdTokenSource):
+    """IDTokenSource that reads the ID token from an environment variable.
+
+    Parameters
+    ----------
+    env_var : str
+        The name of the environment variable containing the ID token.
+    """
+
+    def __init__(self, env_var: str):
+        self.env_var = env_var
+
+    def id_token(self) -> IdToken:
+        """Get an ID token from an environment variable.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        ValueError
+            If the environment variable is not set.
+        """
+        token = os.getenv(self.env_var)
+        if not token:
+            raise ValueError(f"Missing env var {self.env_var!r}")
+        return IdToken(jwt=token)
+
+
+class FileIdTokenSource(IdTokenSource):
+    """IDTokenSource that reads the ID token from a file.
+
+    Parameters
+    ----------
+    path : str
+        The path to the file containing the ID token.
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+
+    def id_token(self) -> IdToken:
+        """Get an ID token from a file.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        ValueError
+            If the file is empty, does not exist, or cannot be read.
+        """
+        if not self.path:
+            raise ValueError("Missing path")
+
+        token = None
+        try:
+            with open(self.path, "r") as f:
+                token = f.read().strip()
+        except FileNotFoundError:
+            raise ValueError(f"File {self.path!r} does not exist")
+        except Exception as e:
+            raise ValueError(f"Error reading token file: {str(e)}")
+
+        if not token:
+            raise ValueError(f"File {self.path!r} is empty")
+        return IdToken(jwt=token)
+
+
+class DatabricksOidcTokenSource(oauth.TokenSource):
+    """A TokenSource which exchanges a token using Workload Identity Federation.
+
+    Parameters
+    ----------
+    host : str
+        The host of the Databricks account or workspace.
+    id_token_source : IdTokenSource
+        IDTokenSource that returns the IDToken to be used for the token exchange.
+    token_endpoint_provider : Callable[[], dict]
+        Returns the token endpoint for the Databricks OIDC application.
+    client_id : Optional[str], optional
+        ClientID of the Databricks OIDC application. It corresponds to the
+        Application ID of the Databricks Service Principal. Only required for
+        Workload Identity Federation and should be empty for Account-wide token
+        federation.
+    account_id : Optional[str], optional
+        The account ID of the Databricks Account. Only required for
+        Account-wide token federation.
+    audience : Optional[str], optional
+        The audience of the Databricks OIDC application. Only used for
+        Workspace level tokens.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        token_endpoint: str,
+        id_token_source: IdTokenSource,
+        client_id: Optional[str] = None,
+        account_id: Optional[str] = None,
+        audience: Optional[str] = None,
+        disable_async: bool = False,
+    ):
+        self._host = host
+        self._id_token_source = id_token_source
+        self._token_endpoint = token_endpoint
+        self._client_id = client_id
+        self._account_id = account_id
+        self._audience = audience
+        self._disable_async = disable_async
+
+    def token(self) -> oauth.Token:
+        """Get a token by exchanging the ID token.
+
+        Returns
+        -------
+        dict
+            The exchanged token.
+
+        Raises
+        ------
+        ValueError
+            If the host is missing or other configuration errors occur.
+        """
+        if not self._host:
+            logger.debug("Missing Host")
+            raise ValueError("missing Host")
+
+        if not self._client_id:
+            logger.debug("No ClientID provided, authenticating with Account-wide token federation")
+        else:
+            logger.debug("Client ID provided, authenticating with Workload Identity Federation")
+
+        id_token = self._id_token_source.id_token()
+
+        client = oauth.ClientCredentials(
+            client_id=self._client_id,
+            client_secret="",  # we have no (rotatable) secrets in OIDC flow
+            token_url=self._token_endpoint,
+            endpoint_params={
+                "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+                "subject_token": id_token,
+                "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            },
+            scopes=["all-apis"],
+            use_params=True,
+            disable_async=self._disable_async,
+        )
+
+        return client.token()

databricks/sdk/service/apps.py

@@ -655,6 +655,8 @@ class AppResource:
 
     sql_warehouse: Optional[AppResourceSqlWarehouse] = None
 
+    uc_securable: Optional[AppResourceUcSecurable] = None
+
     def as_dict(self) -> dict:
         """Serializes the AppResource into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -670,6 +672,8 @@ class AppResource:
             body["serving_endpoint"] = self.serving_endpoint.as_dict()
         if self.sql_warehouse:
             body["sql_warehouse"] = self.sql_warehouse.as_dict()
+        if self.uc_securable:
+            body["uc_securable"] = self.uc_securable.as_dict()
         return body
 
     def as_shallow_dict(self) -> dict:
@@ -687,6 +691,8 @@ class AppResource:
             body["serving_endpoint"] = self.serving_endpoint
         if self.sql_warehouse:
             body["sql_warehouse"] = self.sql_warehouse
+        if self.uc_securable:
+            body["uc_securable"] = self.uc_securable
         return body
 
     @classmethod
@@ -699,6 +705,7 @@ class AppResource:
             secret=_from_dict(d, "secret", AppResourceSecret),
             serving_endpoint=_from_dict(d, "serving_endpoint", AppResourceServingEndpoint),
             sql_warehouse=_from_dict(d, "sql_warehouse", AppResourceSqlWarehouse),
+            uc_securable=_from_dict(d, "uc_securable", AppResourceUcSecurable),
         )
 
 
@@ -880,6 +887,57 @@ class AppResourceSqlWarehouseSqlWarehousePermission(Enum):
     IS_OWNER = "IS_OWNER"
 
 
+@dataclass
+class AppResourceUcSecurable:
+    securable_full_name: str
+
+    securable_type: AppResourceUcSecurableUcSecurableType
+
+    permission: AppResourceUcSecurableUcSecurablePermission
+
+    def as_dict(self) -> dict:
+        """Serializes the AppResourceUcSecurable into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.permission is not None:
+            body["permission"] = self.permission.value
+        if self.securable_full_name is not None:
+            body["securable_full_name"] = self.securable_full_name
+        if self.securable_type is not None:
+            body["securable_type"] = self.securable_type.value
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the AppResourceUcSecurable into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.permission is not None:
+            body["permission"] = self.permission
+        if self.securable_full_name is not None:
+            body["securable_full_name"] = self.securable_full_name
+        if self.securable_type is not None:
+            body["securable_type"] = self.securable_type
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> AppResourceUcSecurable:
+        """Deserializes the AppResourceUcSecurable from a dictionary."""
+        return cls(
+            permission=_enum(d, "permission", AppResourceUcSecurableUcSecurablePermission),
+            securable_full_name=d.get("securable_full_name", None),
+            securable_type=_enum(d, "securable_type", AppResourceUcSecurableUcSecurableType),
+        )
+
+
+class AppResourceUcSecurableUcSecurablePermission(Enum):
+
+    READ_VOLUME = "READ_VOLUME"
+    WRITE_VOLUME = "WRITE_VOLUME"
+
+
+class AppResourceUcSecurableUcSecurableType(Enum):
+
+    VOLUME = "VOLUME"
+
+
 class ApplicationState(Enum):
 
     CRASHED = "CRASHED"