databricks-sdk 0.54.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/version.py

@@ -1 +1 @@
-__version__ = "0.54.0"
+__version__ = "0.55.0"

{databricks_sdk-0.54.0.dist-info → databricks_sdk-0.55.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-sdk
-Version: 0.54.0
+Version: 0.55.0
 Summary: Databricks SDK for Python (Beta)
 Project-URL: Documentation, https://databricks-sdk-py.readthedocs.io
 Keywords: databricks,sdk

{databricks_sdk-0.54.0.dist-info → databricks_sdk-0.55.0.dist-info}/RECORD

@@ -1,22 +1,23 @@
 databricks/__init__.py,sha256=CF2MJcZFwbpn9TwQER8qnCDhkPooBGQNVkX4v7g6p3g,537
-databricks/sdk/__init__.py,sha256=gEriDe8PhOc8T_htILvkBJe91u0PwVDdMZ9sgJ9N8lw,57808
+databricks/sdk/__init__.py,sha256=lliePy2zXmwEuf6KYDV6LjEl3ehc-SwCCZNvipK0xtg,57869
 databricks/sdk/_base_client.py,sha256=IMHtzC5BhWt-lBVjifewR1Ah5fegGDMv0__-O1hCxWI,15850
 databricks/sdk/_property.py,sha256=ccbxhkXZmZOxbx2sqKMTzhVZDuvWXG0WPHFRgac6JAM,1701
 databricks/sdk/azure.py,sha256=sN_ARpmP9h1JovtiHIsDLtrVQP_K11eNDDtHS6PD19k,1015
 databricks/sdk/casing.py,sha256=gZy-FlI7og5WNVX88Vb_7S1WeInwJLGws80CGj_9s48,1137
 databricks/sdk/clock.py,sha256=Ivlow0r_TkXcTJ8UXkxSA0czKrY0GvwHAeOvjPkJnAQ,1360
-databricks/sdk/config.py,sha256=J5kKysLD4V6VfNkoD0tEPOW6-vfU0AjQZmoGw-DBuV8,23013
+databricks/sdk/config.py,sha256=rebzZAw0aMSxSwBeXKsF2VE9X_Y33Kjvcd1PO-5wgc4,23401
 databricks/sdk/core.py,sha256=6lsRl6BL3pLgqMMVFrOnQsx-RxxaJJL_Gt2jJfWUovs,3724
-databricks/sdk/credentials_provider.py,sha256=6ideQ4KtP2nB8a3BZxWlWLX4wNF2QoeoHgYJ356LAMo,39310
+databricks/sdk/credentials_provider.py,sha256=9_P3N52S87xPwI_yUSajnT49--kJWLhKCoHpn5Dwzps,41305
 databricks/sdk/data_plane.py,sha256=br5IPnOdE611IBubxP8xkUR9_qzbSRSYyVWSua6znWs,3109
 databricks/sdk/dbutils.py,sha256=PoDIwNAYGZhVZC7krox7tsudUDNVSk0gsFjFWlKJXVk,15753
 databricks/sdk/environments.py,sha256=9eVeb68cksqY2Lqwth2PJNmK0JEGdIjh-ebrrmUbqCc,3963
-databricks/sdk/oauth.py,sha256=8Cod3Fa4gPuVhd7zFBkVOVLB1fwiAKkyibU6_wnpqDE,28411
+databricks/sdk/oauth.py,sha256=wOcZVfi-Jd83XQDW1rbDIJbxFqJOjaeTSlUgQYD8VWQ,28406
+databricks/sdk/oidc.py,sha256=A9umMkfnL-Nwfw2GljGxqTtkz7PjMTzltGaeckfrvT4,5749
 databricks/sdk/oidc_token_supplier.py,sha256=QrO6J0QY4yFfcdQDL5h2OfxMxvBZJPtPmPeqLbPJ5Xw,1065
 databricks/sdk/py.typed,sha256=pSvaHpbY1UPNEXyVFUjlgBhjPFZMmVC_UNrPC7eMOHI,74
 databricks/sdk/retries.py,sha256=7k2kEexGqGKXHNAWHbPFSZSugU8UIU0qtyly_hix22Q,2581
 databricks/sdk/useragent.py,sha256=boEgzTv-Zmo6boipZKjSopNy0CXg4GShC1_lTKpJgqs,7361
-databricks/sdk/version.py,sha256=GL361gEztfNayz_pZjsGe3Es4PhpuAOve1lCqPCbSZs,23
+databricks/sdk/version.py,sha256=UQLHFqXC6UhDcoKB4MFAazCsPAeQEfBN68BEvb4D2Tk,23
 databricks/sdk/_widgets/__init__.py,sha256=VhI-VvLlr3rKUT1nbROslHJIbmZX_tPJ9rRhrdFsYUA,2811
 databricks/sdk/_widgets/default_widgets_utils.py,sha256=_hwCbptLbRzWEmknco0H1wQNAYcuy2pjFO9NiRbvFeo,1127
 databricks/sdk/_widgets/ipywidgets_utils.py,sha256=mg3rEPG9z76e0yVjGgcLybUvd_zSuN5ziGeKiZ-c8Ew,2927
@@ -63,9 +64,9 @@ databricks/sdk/service/sharing.py,sha256=i7d5ggmOnHJ2gVp4PfPmf14M-AUDRIU-XZb4Sqk
 databricks/sdk/service/sql.py,sha256=v6uspo4o0VoUW0T2eHhPDlJYeYah1vhR_WMxVwEhu2Y,442747
 databricks/sdk/service/vectorsearch.py,sha256=Y1yUJwoY7rg6HWKh7_90aJzIWNiWEG2WUw0i1n3KH7I,86701
 databricks/sdk/service/workspace.py,sha256=T0ZbnG1qcPjKysGO_tBzl5x1PyalydeYJRBZbooYNm0,130893
-databricks_sdk-0.54.0.dist-info/licenses/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
-databricks_sdk-0.54.0.dist-info/licenses/NOTICE,sha256=tkRcQYA1k68wDLcnOWbg2xJDsUOJw8G8DGBhb8dnI3w,1588
-databricks_sdk-0.54.0.dist-info/METADATA,sha256=ykrBHXQeYmMd6DMZrpcwXUrHhE_A2ipJTKiTLLiYtfM,39397
-databricks_sdk-0.54.0.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-databricks_sdk-0.54.0.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
-databricks_sdk-0.54.0.dist-info/RECORD,,
+databricks_sdk-0.55.0.dist-info/licenses/LICENSE,sha256=afBgTZo-JsYqj4VOjnejBetMuHKcFR30YobDdpVFkqY,11411
+databricks_sdk-0.55.0.dist-info/licenses/NOTICE,sha256=tkRcQYA1k68wDLcnOWbg2xJDsUOJw8G8DGBhb8dnI3w,1588
+databricks_sdk-0.55.0.dist-info/METADATA,sha256=ufiMw6V7EtmZD-5ZuCoanbezk94CfykKLbtNASbNFDI,39397
+databricks_sdk-0.55.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+databricks_sdk-0.55.0.dist-info/top_level.txt,sha256=7kRdatoSgU0EUurRQJ_3F1Nv4EOSHWAr6ng25tJOJKU,11
+databricks_sdk-0.55.0.dist-info/RECORD,,

{databricks_sdk-0.54.0.dist-info → databricks_sdk-0.55.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any