databricks-sdk 0.59.0__py3-none-any.whl → 0.61.0__py3-none-any.whl

databricks/sdk/__init__.py

@@ -13,7 +13,7 @@ from databricks.sdk.mixins.files import DbfsExt, FilesExt
 from databricks.sdk.mixins.jobs import JobsExt
 from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
 from databricks.sdk.mixins.workspace import WorkspaceExt
-from databricks.sdk.service import aibuilder as pkg_aibuilder
+from databricks.sdk.service import agentbricks as pkg_agentbricks
 from databricks.sdk.service import apps as pkg_apps
 from databricks.sdk.service import billing as pkg_billing
 from databricks.sdk.service import catalog as pkg_catalog
@@ -36,7 +36,7 @@ from databricks.sdk.service import sharing as pkg_sharing
 from databricks.sdk.service import sql as pkg_sql
 from databricks.sdk.service import vectorsearch as pkg_vectorsearch
 from databricks.sdk.service import workspace as pkg_workspace
-from databricks.sdk.service.aibuilder import AiBuilderAPI
+from databricks.sdk.service.agentbricks import AgentBricksAPI
 from databricks.sdk.service.apps import AppsAPI
 from databricks.sdk.service.billing import (BillableUsageAPI, BudgetPolicyAPI,
                                             BudgetsAPI, LogDeliveryAPI,
@@ -59,7 +59,9 @@ from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
                                             TableConstraintsAPI, TablesAPI,
                                             TemporaryTableCredentialsAPI,
                                             VolumesAPI, WorkspaceBindingsAPI)
-from databricks.sdk.service.cleanrooms import (CleanRoomAssetsAPI,
+from databricks.sdk.service.cleanrooms import (CleanRoomAssetRevisionsAPI,
+                                               CleanRoomAssetsAPI,
+                                               CleanRoomAutoApprovalRulesAPI,
                                                CleanRoomsAPI,
                                                CleanRoomTaskRunsAPI)
 from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI,
@@ -240,14 +242,16 @@ class WorkspaceClient:
         serving_endpoints = ServingEndpointsExt(self._api_client)
         self._access_control = pkg_iam.AccessControlAPI(self._api_client)
         self._account_access_control_proxy = pkg_iam.AccountAccessControlProxyAPI(self._api_client)
-        self._ai_builder = pkg_aibuilder.AiBuilderAPI(self._api_client)
+        self._agent_bricks = pkg_agentbricks.AgentBricksAPI(self._api_client)
         self._alerts = pkg_sql.AlertsAPI(self._api_client)
         self._alerts_legacy = pkg_sql.AlertsLegacyAPI(self._api_client)
         self._alerts_v2 = pkg_sql.AlertsV2API(self._api_client)
         self._apps = pkg_apps.AppsAPI(self._api_client)
         self._artifact_allowlists = pkg_catalog.ArtifactAllowlistsAPI(self._api_client)
         self._catalogs = pkg_catalog.CatalogsAPI(self._api_client)
+        self._clean_room_asset_revisions = pkg_cleanrooms.CleanRoomAssetRevisionsAPI(self._api_client)
         self._clean_room_assets = pkg_cleanrooms.CleanRoomAssetsAPI(self._api_client)
+        self._clean_room_auto_approval_rules = pkg_cleanrooms.CleanRoomAutoApprovalRulesAPI(self._api_client)
         self._clean_room_task_runs = pkg_cleanrooms.CleanRoomTaskRunsAPI(self._api_client)
         self._clean_rooms = pkg_cleanrooms.CleanRoomsAPI(self._api_client)
         self._cluster_policies = pkg_compute.ClusterPoliciesAPI(self._api_client)
@@ -377,9 +381,9 @@ class WorkspaceClient:
         return self._account_access_control_proxy
 
     @property
-    def ai_builder(self) -> pkg_aibuilder.AiBuilderAPI:
+    def agent_bricks(self) -> pkg_agentbricks.AgentBricksAPI:
         """The Custom LLMs service manages state and powers the UI for the Custom LLM product."""
-        return self._ai_builder
+        return self._agent_bricks
 
     @property
     def alerts(self) -> pkg_sql.AlertsAPI:
@@ -411,11 +415,21 @@ class WorkspaceClient:
         """A catalog is the first layer of Unity Catalog’s three-level namespace."""
         return self._catalogs
 
+    @property
+    def clean_room_asset_revisions(self) -> pkg_cleanrooms.CleanRoomAssetRevisionsAPI:
+        """Clean Room Asset Revisions denote new versions of uploaded assets (e.g."""
+        return self._clean_room_asset_revisions
+
     @property
     def clean_room_assets(self) -> pkg_cleanrooms.CleanRoomAssetsAPI:
         """Clean room assets are data and code objects — Tables, volumes, and notebooks that are shared with the clean room."""
         return self._clean_room_assets
 
+    @property
+    def clean_room_auto_approval_rules(self) -> pkg_cleanrooms.CleanRoomAutoApprovalRulesAPI:
+        """Clean room auto-approval rules automatically create an approval on your behalf when an asset (e.g."""
+        return self._clean_room_auto_approval_rules
+
     @property
     def clean_room_task_runs(self) -> pkg_cleanrooms.CleanRoomTaskRunsAPI:
         """Clean room task runs are the executions of notebooks in a clean room."""

databricks/sdk/credentials_provider.py

@@ -331,7 +331,7 @@ def file_oidc(cfg) -> Optional[CredentialsProvider]:
 # that provides a Databricks token from an IdTokenSource.
 def _oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optional[CredentialsProvider]:
     try:
-        id_token = id_token_source.id_token()
+        id_token_source.id_token()  # validate the id_token_source
     except Exception as e:
         logger.debug(f"Failed to get OIDC token: {e}")
         return None
@@ -341,7 +341,7 @@ def _oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Opti
         token_endpoint=cfg.oidc_endpoints.token_endpoint,
         client_id=cfg.client_id,
         account_id=cfg.account_id,
-        id_token=id_token,
+        id_token_source=id_token_source,
         disable_async=cfg.disable_async_token_refresh,
     )
 

databricks/sdk/mixins/files.py

@@ -28,7 +28,6 @@ from .._base_client import _BaseClient, _RawResponse, _StreamingResponse
 from .._property import _cached_property
 from ..config import Config
 from ..errors import AlreadyExists, NotFound
-from ..errors.customizer import _RetryAfterCustomizer
 from ..errors.mapper import _error_mapper
 from ..retries import retried
 from ..service import files
@@ -577,6 +576,27 @@ class _DbfsPath(_Path):
         return f"<_DbfsPath {self._path}>"
 
 
+class _RetryableException(Exception):
+    """Base class for retryable exceptions in DBFS operations."""
+
+    def __init__(self, message: str, http_status_code: int):
+        super().__init__()
+        self.message = message
+        self.http_status_code = http_status_code
+
+    def __str__(self) -> str:
+        return f"{self.message} (HTTP Status: {self.http_status_code})"
+
+    @staticmethod
+    def make_error(response: requests.Response) -> "_RetryableException":
+        """Map the response to a retryable exception."""
+
+        return _RetryableException(
+            message=response.text,
+            http_status_code=response.status_code,
+        )
+
+
 class DbfsExt(files.DbfsAPI):
     __doc__ = files.DbfsAPI.__doc__
 
@@ -885,7 +905,7 @@ class FilesExt(files.FilesAPI):
                         timeout=self._config.multipart_upload_single_chunk_upload_timeout_seconds,
                     )
 
-                upload_response = self._retry_idempotent_operation(perform, rewind)
+                upload_response = self._retry_cloud_idempotent_operation(perform, rewind)
 
                 if upload_response.status_code in (200, 201):
                     # Chunk upload successful
@@ -1097,7 +1117,7 @@ class FilesExt(files.FilesAPI):
                         )
 
                     try:
-                        return self._retry_idempotent_operation(perform)
+                        return self._retry_cloud_idempotent_operation(perform)
                     except RequestException:
                         _LOG.warning("Failed to retrieve upload status")
                         return None
@@ -1116,7 +1136,7 @@ class FilesExt(files.FilesAPI):
                     # a 503 or 500 response, then you need to resume the interrupted upload from where it left off.
 
                     # Let's follow that for all potentially retryable status codes.
-                    # Together with the catch block below we replicate the logic in _retry_idempotent_operation().
+                    # Together with the catch block below we replicate the logic in _retry_databricks_idempotent_operation().
                     if upload_response.status_code in self._RETRYABLE_STATUS_CODES:
                         if retry_count < self._config.multipart_upload_max_retries:
                             retry_count += 1
@@ -1243,7 +1263,7 @@ class FilesExt(files.FilesAPI):
                 timeout=self._config.multipart_upload_single_chunk_upload_timeout_seconds,
             )
 
-        abort_response = self._retry_idempotent_operation(perform)
+        abort_response = self._retry_cloud_idempotent_operation(perform)
 
         if abort_response.status_code not in (200, 201):
             raise ValueError(abort_response)
@@ -1265,7 +1285,7 @@ class FilesExt(files.FilesAPI):
                 timeout=self._config.multipart_upload_single_chunk_upload_timeout_seconds,
             )
 
-        abort_response = self._retry_idempotent_operation(perform)
+        abort_response = self._retry_cloud_idempotent_operation(perform)
 
         if abort_response.status_code not in (200, 201):
             raise ValueError(abort_response)
@@ -1283,23 +1303,31 @@ class FilesExt(files.FilesAPI):
         session.mount("http://", http_adapter)
         return session
 
-    def _retry_idempotent_operation(
+    def _retry_cloud_idempotent_operation(
         self, operation: Callable[[], requests.Response], before_retry: Callable = None
     ) -> requests.Response:
-        """Perform given idempotent operation with necessary retries. Since operation is idempotent it's
-        safe to retry it for response codes where server state might have changed.
+        """Perform given idempotent operation with necessary retries for requests to non Databricks APIs.
+        For cloud APIs, we will retry on network errors and on server response codes.
+        Since operation is idempotent it's safe to retry it for response codes where server state might have changed.
         """
 
-        def delegate():
+        def delegate() -> requests.Response:
             response = operation()
             if response.status_code in self._RETRYABLE_STATUS_CODES:
-                attrs = {}
-                # this will assign "retry_after_secs" to the attrs, essentially making exception look retryable
-                _RetryAfterCustomizer().customize_error(response, attrs)
-                raise _error_mapper(response, attrs)
+                raise _RetryableException.make_error(response)
             else:
                 return response
 
+        def extended_is_retryable(e: BaseException) -> Optional[str]:
+            retry_reason_from_base = _BaseClient._is_retryable(e)
+            if retry_reason_from_base is not None:
+                return retry_reason_from_base
+
+            if isinstance(e, _RetryableException):
+                # this is a retriable exception, but not a network error
+                return f"retryable exception (status_code:{e.http_status_code})"
+            return None
+
         # following _BaseClient timeout
         retry_timeout_seconds = self._config.retry_timeout_seconds or 300
 
@@ -1307,7 +1335,7 @@ class FilesExt(files.FilesAPI):
             timeout=timedelta(seconds=retry_timeout_seconds),
             # also retry on network errors (connection error, connection timeout)
             # where we believe request didn't reach the server
-            is_retryable=_BaseClient._is_retryable,
+            is_retryable=extended_is_retryable,
             before_retry=before_retry,
         )(delegate)()
 

databricks/sdk/mixins/open_ai_client.py

@@ -4,6 +4,7 @@ from typing import Dict, Optional
 from requests import Response
 
 from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
+                                            HttpRequestResponse,
                                             ServingEndpointsAPI)
 
 
@@ -88,15 +89,30 @@ class ServingEndpointsExt(ServingEndpointsAPI):
         """
         response = Response()
         response.status_code = 200
-        server_response = super().http_request(
-            connection_name=conn,
-            method=method,
-            path=path,
-            headers=js.dumps(headers) if headers is not None else None,
-            json=js.dumps(json) if json is not None else None,
-            params=js.dumps(params) if params is not None else None,
+
+        # We currently don't call super.http_request because we need to pass in response_headers
+        # This is a temporary fix to get the headers we need for the MCP session id
+        # TODO: Remove this once we have a better way to get back the response headers
+        headers_to_capture = ["mcp-session-id"]
+        res = self._api.do(
+            "POST",
+            "/api/2.0/external-function",
+            body={
+                "connection_name": conn,
+                "method": method.value,
+                "path": path,
+                "headers": js.dumps(headers) if headers is not None else None,
+                "json": js.dumps(json) if json is not None else None,
+                "params": js.dumps(params) if params is not None else None,
+            },
+            headers={"Accept": "text/plain", "Content-Type": "application/json"},
+            raw=True,
+            response_headers=headers_to_capture,
         )
 
+        # Create HttpRequestResponse from the raw response
+        server_response = HttpRequestResponse.from_dict(res)
+
         # Read the content from the HttpRequestResponse object
         if hasattr(server_response, "contents") and hasattr(server_response.contents, "read"):
             raw_content = server_response.contents.read()  # Read the bytes
@@ -109,4 +125,9 @@ class ServingEndpointsExt(ServingEndpointsAPI):
         else:
             raise ValueError("Contents must be bytes.")
 
+        # Copy headers from raw response to Response
+        for header_name in headers_to_capture:
+            if header_name in res:
+                response.headers[header_name] = res[header_name]
+
         return response

databricks/sdk/oidc.py

@@ -188,14 +188,18 @@ class DatabricksOidcTokenSource(oauth.TokenSource):
             logger.debug("Client ID provided, authenticating with Workload Identity Federation")
 
         id_token = self._id_token_source.id_token()
+        return self._exchange_id_token(id_token)
 
+    # This function is used to create the OAuth client.
+    # It exists to make it easier to test.
+    def _exchange_id_token(self, id_token: IdToken) -> oauth.Token:
         client = oauth.ClientCredentials(
             client_id=self._client_id,
-            client_secret="",  # we have no (rotatable) secrets in OIDC flow
+            client_secret="",  # there is no (rotatable) secrets in the OIDC flow
             token_url=self._token_endpoint,
             endpoint_params={
                 "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
-                "subject_token": id_token,
+                "subject_token": id_token.jwt,
                 "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
             },
             scopes=["all-apis"],

databricks/sdk/service/{aibuilder.py → agentbricks.py}

@@ -23,9 +23,6 @@ class CustomLlm:
     instructions: str
     """Instructions for the custom LLM to follow"""
 
-    optimization_state: State
-    """If optimization is kicked off, tracks the state of the custom LLM"""
-
     agent_artifact_path: Optional[str] = None
 
     creation_time: Optional[str] = None
@@ -45,6 +42,9 @@ class CustomLlm:
 
     id: Optional[str] = None
 
+    optimization_state: Optional[State] = None
+    """If optimization is kicked off, tracks the state of the custom LLM"""
+
     def as_dict(self) -> dict:
         """Serializes the CustomLlm into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -190,7 +190,7 @@ class Table:
         )
 
 
-class AiBuilderAPI:
+class AgentBricksAPI:
     """The Custom LLMs service manages state and powers the UI for the Custom LLM product."""
 
     def __init__(self, api_client):
@@ -270,7 +270,7 @@ class AiBuilderAPI:
             "Accept": "application/json",
         }
 
-        self._api.do("DELETE", f"/api/2.0/custom-lms/{id}", headers=headers)
+        self._api.do("DELETE", f"/api/2.0/custom-llms/{id}", headers=headers)
 
     def get_custom_llm(self, id: str) -> CustomLlm:
         """Get a Custom LLM.