Merge branch 'main' into main

Author: jijosg

.codegen/_openapi_sha

@@ -1 +1 @@
-033bcb9242b006001e2cf3956896711681de1a8c
+69902d1abe35bd9e78e0231927bf14d11b383a16

.release_metadata.json

@@ -1,3 +1,3 @@
 {
-    "timestamp": "2025-06-12 19:21:54+0000"
+    "timestamp": "2025-07-17 11:13:08+0000"
 }

CHANGELOG.md

@@ -1,25 +1,20 @@
 # NEXT CHANGELOG
 
-## Release v0.58.0
+## Release v0.60.0
 
 ### New Features and Improvements
 
+* Added headers to HttpRequestResponse in OpenAI client.
+
 ### Bug Fixes
 
+- Correctly issue in OIDC implementation that prevented the use of the feature (see #994).
+- Fix a reported issue where `FilesExt` fails to retry if it receives certain status code from server.
+
 ### Documentation
 
 ### Internal Changes
 
+- Refactor unit tests for `FilesExt` to improve its readability.
+
 ### API Changes
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.ClusterAttributes`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.ClusterDetails`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.ClusterSpec`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.CreateCluster`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.CreateInstancePool`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.EditCluster`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.EditInstancePool`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.GetInstancePool`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.InstancePoolAndStats`.
-* Added `remote_disk_throughput` and `total_initial_remote_disk_size` fields for `databricks.sdk.service.compute.UpdateClusterResource`.
-* Added `r` enum value for `databricks.sdk.service.compute.Language`.
-* Added `continuous` and `continuous_restart` enum values for `databricks.sdk.service.jobs.TriggerType`.

NEXT_CHANGELOG.md

@@ -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/__init__.py

@@ -28,7 +28,6 @@
 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 @@ def __repr__(self) -> str:
         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 @@ def perform():
                         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 @@ def perform():
                         )
 
                     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 @@ def perform():
                     # 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 @@ def perform():
                 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 @@ def perform():
                 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,31 +1303,39 @@ def _create_cloud_provider_session(self):
         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
 
         return retried(
             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/credentials_provider.py

@@ -4,6 +4,7 @@
 from requests import Response
 
 from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
+                                            HttpRequestResponse,
                                             ServingEndpointsAPI)
 
 
@@ -88,15 +89,30 @@ def http_request(
         """
         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 @@ def http_request(
         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