databricks-sdk 0.41.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

databricks/sdk/__init__.py

@@ -13,8 +13,9 @@ 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.apps import AppsAPI
-from databricks.sdk.service.billing import (BillableUsageAPI, BudgetsAPI,
-                                            LogDeliveryAPI, UsageDashboardsAPI)
+from databricks.sdk.service.billing import (BillableUsageAPI, BudgetPolicyAPI,
+                                            BudgetsAPI, LogDeliveryAPI,
+                                            UsageDashboardsAPI)
 from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
                                             AccountMetastoresAPI,
                                             AccountStorageCredentialsAPI,
@@ -41,7 +42,9 @@ from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI,
                                             InstanceProfilesAPI, LibrariesAPI,
                                             PolicyComplianceForClustersAPI,
                                             PolicyFamiliesAPI)
-from databricks.sdk.service.dashboards import GenieAPI, LakeviewAPI
+from databricks.sdk.service.dashboards import (GenieAPI, LakeviewAPI,
+                                               LakeviewEmbeddedAPI,
+                                               QueryExecutionAPI)
 from databricks.sdk.service.files import DbfsAPI, FilesAPI
 from databricks.sdk.service.iam import (AccessControlAPI,
                                         AccountAccessControlAPI,
@@ -80,7 +83,7 @@ from databricks.sdk.service.settings import (
     AibiDashboardEmbeddingApprovedDomainsAPI, AutomaticClusterUpdateAPI,
     ComplianceSecurityProfileAPI, CredentialsManagerAPI,
     CspEnablementAccountAPI, DefaultNamespaceAPI, DisableLegacyAccessAPI,
-    DisableLegacyDbfsAPI, DisableLegacyFeaturesAPI,
+    DisableLegacyDbfsAPI, DisableLegacyFeaturesAPI, EnableIpAccessListsAPI,
     EnhancedSecurityMonitoringAPI, EsmEnablementAccountAPI, IpAccessListsAPI,
     NetworkConnectivityAPI, NotificationDestinationsAPI, PersonalComputeAPI,
     RestrictWorkspaceAdminsAPI, SettingsAPI, TokenManagementAPI, TokensAPI,
@@ -95,7 +98,8 @@ from databricks.sdk.service.sql import (AlertsAPI, AlertsLegacyAPI,
                                         QueryHistoryAPI,
                                         QueryVisualizationsAPI,
                                         QueryVisualizationsLegacyAPI,
-                                        StatementExecutionAPI, WarehousesAPI)
+                                        RedashConfigAPI, StatementExecutionAPI,
+                                        WarehousesAPI)
 from databricks.sdk.service.vectorsearch import (VectorSearchEndpointsAPI,
                                                  VectorSearchIndexesAPI)
 from databricks.sdk.service.workspace import (GitCredentialsAPI, ReposAPI,
@@ -230,6 +234,7 @@ class WorkspaceClient:
         self._ip_access_lists = IpAccessListsAPI(self._api_client)
         self._jobs = JobsExt(self._api_client)
         self._lakeview = LakeviewAPI(self._api_client)
+        self._lakeview_embedded = LakeviewEmbeddedAPI(self._api_client)
         self._libraries = LibrariesAPI(self._api_client)
         self._metastores = MetastoresAPI(self._api_client)
         self._model_registry = ModelRegistryAPI(self._api_client)
@@ -254,11 +259,13 @@ class WorkspaceClient:
         self._quality_monitors = QualityMonitorsAPI(self._api_client)
         self._queries = QueriesAPI(self._api_client)
         self._queries_legacy = QueriesLegacyAPI(self._api_client)
+        self._query_execution = QueryExecutionAPI(self._api_client)
         self._query_history = QueryHistoryAPI(self._api_client)
         self._query_visualizations = QueryVisualizationsAPI(self._api_client)
         self._query_visualizations_legacy = QueryVisualizationsLegacyAPI(self._api_client)
         self._recipient_activation = RecipientActivationAPI(self._api_client)
         self._recipients = RecipientsAPI(self._api_client)
+        self._redash_config = RedashConfigAPI(self._api_client)
         self._registered_models = RegisteredModelsAPI(self._api_client)
         self._repos = ReposAPI(self._api_client)
         self._resource_quotas = ResourceQuotasAPI(self._api_client)
@@ -503,6 +510,11 @@ class WorkspaceClient:
         """These APIs provide specific management operations for Lakeview dashboards."""
         return self._lakeview
 
+    @property
+    def lakeview_embedded(self) -> LakeviewEmbeddedAPI:
+        """Token-based Lakeview APIs for embedding dashboards in external applications."""
+        return self._lakeview_embedded
+
     @property
     def libraries(self) -> LibrariesAPI:
         """The Libraries API allows you to install and uninstall libraries and get the status of libraries on a cluster."""
@@ -618,6 +630,11 @@ class WorkspaceClient:
         """These endpoints are used for CRUD operations on query definitions."""
         return self._queries_legacy
 
+    @property
+    def query_execution(self) -> QueryExecutionAPI:
+        """Query execution APIs for AI / BI Dashboards."""
+        return self._query_execution
+
     @property
     def query_history(self) -> QueryHistoryAPI:
         """A service responsible for storing and retrieving the list of queries run against SQL endpoints and serverless compute."""
@@ -643,6 +660,11 @@ class WorkspaceClient:
         """A recipient is an object you create using :method:recipients/create to represent an organization which you want to allow access shares."""
         return self._recipients
 
+    @property
+    def redash_config(self) -> RedashConfigAPI:
+        """Redash V2 service for workspace configurations (internal)."""
+        return self._redash_config
+
     @property
     def registered_models(self) -> RegisteredModelsAPI:
         """Databricks provides a hosted version of MLflow Model Registry in Unity Catalog."""
@@ -845,6 +867,7 @@ class AccountClient:
         self._api_client = client.ApiClient(self._config)
         self._access_control = AccountAccessControlAPI(self._api_client)
         self._billable_usage = BillableUsageAPI(self._api_client)
+        self._budget_policy = BudgetPolicyAPI(self._api_client)
         self._credentials = CredentialsAPI(self._api_client)
         self._custom_app_integration = CustomAppIntegrationAPI(self._api_client)
         self._encryption_keys = EncryptionKeysAPI(self._api_client)
@@ -890,6 +913,11 @@ class AccountClient:
         """This API allows you to download billable usage logs for the specified account and date range."""
         return self._billable_usage
 
+    @property
+    def budget_policy(self) -> BudgetPolicyAPI:
+        """A service serves REST API about Budget policies."""
+        return self._budget_policy
+
     @property
     def credentials(self) -> CredentialsAPI:
         """These APIs manage credential configurations for this workspace."""

databricks/sdk/_base_client.py

@@ -159,16 +159,29 @@ class _BaseClient:
         if isinstance(data, (str, bytes)):
             data = io.BytesIO(data.encode('utf-8') if isinstance(data, str) else data)
 
-        # Only retry if the request is not a stream or if the stream is seekable and
-        # we can rewind it. This is necessary to avoid bugs where the retry doesn't
-        # re-read already read data from the body.
-        if data is not None and not self._is_seekable_stream(data):
-            logger.debug(f"Retry disabled for non-seekable stream: type={type(data)}")
-            call = self._perform
-        else:
+        if not data:
+            # The request is not a stream.
             call = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
                            is_retryable=self._is_retryable,
                            clock=self._clock)(self._perform)
+        elif self._is_seekable_stream(data):
+            # Keep track of the initial position of the stream so that we can rewind to it
+            # if we need to retry the request.
+            initial_data_position = data.tell()
+
+            def rewind():
+                logger.debug(f"Rewinding input data to offset {initial_data_position} before retry")
+                data.seek(initial_data_position)
+
+            call = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
+                           is_retryable=self._is_retryable,
+                           clock=self._clock,
+                           before_retry=rewind)(self._perform)
+        else:
+            # Do not retry if the stream is not seekable. This is necessary to avoid bugs
+            # where the retry doesn't re-read already read data from the stream.
+            logger.debug(f"Retry disabled for non-seekable stream: type={type(data)}")
+            call = self._perform
 
         response = call(method,
                         url,
@@ -249,12 +262,6 @@ class _BaseClient:
                  files=None,
                  data=None,
                  auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None):
-        # Keep track of the initial position of the stream so that we can rewind it if
-        # we need to retry the request.
-        initial_data_position = 0
-        if self._is_seekable_stream(data):
-            initial_data_position = data.tell()
-
         response = self._session.request(method,
                                          url,
                                          params=self._fix_query_string(query),
@@ -266,16 +273,8 @@ class _BaseClient:
                                          stream=raw,
                                          timeout=self._http_timeout_seconds)
         self._record_request_log(response, raw=raw or data is not None or files is not None)
-
         error = self._error_parser.get_api_error(response)
         if error is not None:
-            # If the request body is a seekable stream, rewind it so that it is ready
-            # to be read again in case of a retry.
-            #
-            # TODO: This should be moved into a "before-retry" hook to avoid one
-            # unnecessary seek on the last failed retry before aborting.
-            if self._is_seekable_stream(data):
-                data.seek(initial_data_position)
             raise error from None
 
         return response

databricks/sdk/credentials_provider.py

@@ -676,12 +676,18 @@ class MetadataServiceTokenSource(Refreshable):
         self.host = cfg.host
 
     def refresh(self) -> Token:
-        resp = requests.get(self.url,
-                            timeout=self._metadata_service_timeout,
-                            headers={
-                                self.METADATA_SERVICE_VERSION_HEADER: self.METADATA_SERVICE_VERSION,
-                                self.METADATA_SERVICE_HOST_HEADER: self.host
-                            })
+        resp = requests.get(
+            self.url,
+            timeout=self._metadata_service_timeout,
+            headers={
+                self.METADATA_SERVICE_VERSION_HEADER: self.METADATA_SERVICE_VERSION,
+                self.METADATA_SERVICE_HOST_HEADER: self.host
+            },
+            proxies={
+                # Explicitly exclude localhost from being proxied. This is necessary
+                # for Metadata URLs which typically point to localhost.
+                "no_proxy": "localhost,127.0.0.1"
+            })
         json_resp: dict[str, Union[str, float]] = resp.json()
         access_token = json_resp.get("access_token", None)
         if access_token is None:

databricks/sdk/mixins/open_ai_client.py

@@ -1,8 +1,9 @@
 import json as js
 from typing import Dict, Optional
 
+from requests import Response
+
 from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
-                                            ExternalFunctionResponse,
                                             ServingEndpointsAPI)
 
 
@@ -63,7 +64,7 @@ class ServingEndpointsExt(ServingEndpointsAPI):
                      *,
                      headers: Optional[Dict[str, str]] = None,
                      json: Optional[Dict[str, str]] = None,
-                     params: Optional[Dict[str, str]] = None) -> ExternalFunctionResponse:
+                     params: Optional[Dict[str, str]] = None) -> Response:
         """Make external services call using the credentials stored in UC Connection.
         **NOTE:** Experimental: This API may change or be removed in a future release without warning.
         :param conn: str
@@ -79,13 +80,27 @@ class ServingEndpointsExt(ServingEndpointsAPI):
           JSON payload for the request.
         :param params: Dict[str,str] (optional)
           Query parameters for the request.
-        :returns: :class:`ExternalFunctionResponse`
+        :returns: :class:`Response`
         """
+        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)
+
+        # 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
+        else:
+            raise ValueError("Invalid response from the server.")
+
+        # Set the raw content
+        if isinstance(raw_content, bytes):
+            response._content = raw_content
+        else:
+            raise ValueError("Contents must be bytes.")
 
-        return super.http_request(connection_name=conn,
-                                  method=method,
-                                  path=path,
-                                  headers=js.dumps(headers),
-                                  json=js.dumps(json),
-                                  params=js.dumps(params),
-                                  )
+        return response

databricks/sdk/retries.py

@@ -13,7 +13,8 @@ def retried(*,
             on: Sequence[Type[BaseException]] = None,
             is_retryable: Callable[[BaseException], Optional[str]] = None,
             timeout=timedelta(minutes=20),
-            clock: Clock = None):
+            clock: Clock = None,
+            before_retry: Callable = None):
     has_allowlist = on is not None
     has_callback = is_retryable is not None
     if not (has_allowlist or has_callback) or (has_allowlist and has_callback):
@@ -54,6 +55,9 @@ def retried(*,
                         raise err
 
                     logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)')
+                    if before_retry:
+                        before_retry()
+
                     clock.sleep(sleep + random())
                     attempt += 1
             raise TimeoutError(f'Timed out after {timeout}') from last_err