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

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

databricks/sdk/service/billing.py

@@ -11,6 +11,8 @@ from ._internal import _enum, _from_dict, _repeated_dict
 
 _LOG = logging.getLogger('databricks.sdk')
 
+from databricks.sdk.service import compute
+
 # all definitions in this file are in alphabetical order
 
 
@@ -311,6 +313,44 @@ class BudgetConfigurationFilterWorkspaceIdClause:
                    values=d.get('values', None))
 
 
+@dataclass
+class BudgetPolicy:
+    """Contains the BudgetPolicy details."""
+
+    policy_id: str
+    """The Id of the policy. This field is generated by Databricks and globally unique."""
+
+    custom_tags: Optional[List[compute.CustomPolicyTag]] = None
+    """A list of tags defined by the customer. At most 20 entries are allowed per policy."""
+
+    policy_name: Optional[str] = None
+    """The name of the policy. - Must be unique among active policies. - Can contain only characters
+    from the ISO 8859-1 (latin1) set."""
+
+    def as_dict(self) -> dict:
+        """Serializes the BudgetPolicy into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.custom_tags: body['custom_tags'] = [v.as_dict() for v in self.custom_tags]
+        if self.policy_id is not None: body['policy_id'] = self.policy_id
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the BudgetPolicy into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.custom_tags: body['custom_tags'] = self.custom_tags
+        if self.policy_id is not None: body['policy_id'] = self.policy_id
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> BudgetPolicy:
+        """Deserializes the BudgetPolicy from a dictionary."""
+        return cls(custom_tags=_repeated_dict(d, 'custom_tags', compute.CustomPolicyTag),
+                   policy_id=d.get('policy_id', None),
+                   policy_name=d.get('policy_name', None))
+
+
 @dataclass
 class CreateBillingUsageDashboardRequest:
     dashboard_type: Optional[UsageDashboardType] = None
@@ -536,6 +576,45 @@ class CreateBudgetConfigurationResponse:
         return cls(budget=_from_dict(d, 'budget', BudgetConfiguration))
 
 
+@dataclass
+class CreateBudgetPolicyRequest:
+    """A request to create a BudgetPolicy."""
+
+    custom_tags: Optional[List[compute.CustomPolicyTag]] = None
+    """A list of tags defined by the customer. At most 40 entries are allowed per policy."""
+
+    policy_name: Optional[str] = None
+    """The name of the policy. - Must be unique among active policies. - Can contain only characters of
+    0-9, a-z, A-Z, -, =, ., :, /, @, _, +, whitespace."""
+
+    request_id: Optional[str] = None
+    """A unique identifier for this request. Restricted to 36 ASCII characters. A random UUID is
+    recommended. This request is only idempotent if a `request_id` is provided."""
+
+    def as_dict(self) -> dict:
+        """Serializes the CreateBudgetPolicyRequest into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.custom_tags: body['custom_tags'] = [v.as_dict() for v in self.custom_tags]
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        if self.request_id is not None: body['request_id'] = self.request_id
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the CreateBudgetPolicyRequest into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.custom_tags: body['custom_tags'] = self.custom_tags
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        if self.request_id is not None: body['request_id'] = self.request_id
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> CreateBudgetPolicyRequest:
+        """Deserializes the CreateBudgetPolicyRequest from a dictionary."""
+        return cls(custom_tags=_repeated_dict(d, 'custom_tags', compute.CustomPolicyTag),
+                   policy_name=d.get('policy_name', None),
+                   request_id=d.get('request_id', None))
+
+
 @dataclass
 class CreateLogDeliveryConfigurationParams:
     log_type: LogType
@@ -670,6 +749,25 @@ class DeleteBudgetConfigurationResponse:
         return cls()
 
 
+@dataclass
+class DeleteResponse:
+
+    def as_dict(self) -> dict:
+        """Serializes the DeleteResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the DeleteResponse into a shallow dictionary of its immediate attributes."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> DeleteResponse:
+        """Deserializes the DeleteResponse from a dictionary."""
+        return cls()
+
+
 class DeliveryStatus(Enum):
     """The status string for log delivery. Possible values are: * `CREATED`: There were no log delivery
     attempts since the config was created. * `SUCCEEDED`: The latest attempt of log delivery has
@@ -708,6 +806,44 @@ class DownloadResponse:
         return cls(contents=d.get('contents', None))
 
 
+@dataclass
+class Filter:
+    """Structured representation of a filter to be applied to a list of policies. All specified filters
+    will be applied in conjunction."""
+
+    creator_user_id: Optional[int] = None
+    """The policy creator user id to be filtered on. If unspecified, all policies will be returned."""
+
+    creator_user_name: Optional[str] = None
+    """The policy creator user name to be filtered on. If unspecified, all policies will be returned."""
+
+    policy_name: Optional[str] = None
+    """The partial name of policies to be filtered on. If unspecified, all policies will be returned."""
+
+    def as_dict(self) -> dict:
+        """Serializes the Filter into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.creator_user_id is not None: body['creator_user_id'] = self.creator_user_id
+        if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the Filter into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.creator_user_id is not None: body['creator_user_id'] = self.creator_user_id
+        if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name
+        if self.policy_name is not None: body['policy_name'] = self.policy_name
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> Filter:
+        """Deserializes the Filter from a dictionary."""
+        return cls(creator_user_id=d.get('creator_user_id', None),
+                   creator_user_name=d.get('creator_user_name', None),
+                   policy_name=d.get('policy_name', None))
+
+
 @dataclass
 class GetBillingUsageDashboardResponse:
     dashboard_id: Optional[str] = None
@@ -787,6 +923,44 @@ class ListBudgetConfigurationsResponse:
                    next_page_token=d.get('next_page_token', None))
 
 
+@dataclass
+class ListBudgetPoliciesResponse:
+    """A list of policies."""
+
+    next_page_token: Optional[str] = None
+    """A token that can be sent as `page_token` to retrieve the next page. If this field is omitted,
+    there are no subsequent pages."""
+
+    policies: Optional[List[BudgetPolicy]] = None
+
+    previous_page_token: Optional[str] = None
+    """A token that can be sent as `page_token` to retrieve the previous page. In this field is
+    omitted, there are no previous pages."""
+
+    def as_dict(self) -> dict:
+        """Serializes the ListBudgetPoliciesResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.next_page_token is not None: body['next_page_token'] = self.next_page_token
+        if self.policies: body['policies'] = [v.as_dict() for v in self.policies]
+        if self.previous_page_token is not None: body['previous_page_token'] = self.previous_page_token
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the ListBudgetPoliciesResponse into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.next_page_token is not None: body['next_page_token'] = self.next_page_token
+        if self.policies: body['policies'] = self.policies
+        if self.previous_page_token is not None: body['previous_page_token'] = self.previous_page_token
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> ListBudgetPoliciesResponse:
+        """Deserializes the ListBudgetPoliciesResponse from a dictionary."""
+        return cls(next_page_token=d.get('next_page_token', None),
+                   policies=_repeated_dict(d, 'policies', BudgetPolicy),
+                   previous_page_token=d.get('previous_page_token', None))
+
+
 class LogDeliveryConfigStatus(Enum):
     """Status of log delivery configuration. Set to `ENABLED` (enabled) or `DISABLED` (disabled).
     Defaults to `ENABLED`. You can [enable or disable the
@@ -1046,6 +1220,39 @@ class PatchStatusResponse:
         return cls()
 
 
+@dataclass
+class SortSpec:
+    descending: Optional[bool] = None
+    """Whether to sort in descending order."""
+
+    field: Optional[SortSpecField] = None
+    """The filed to sort by"""
+
+    def as_dict(self) -> dict:
+        """Serializes the SortSpec into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.descending is not None: body['descending'] = self.descending
+        if self.field is not None: body['field'] = self.field.value
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the SortSpec into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.descending is not None: body['descending'] = self.descending
+        if self.field is not None: body['field'] = self.field
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> SortSpec:
+        """Deserializes the SortSpec from a dictionary."""
+        return cls(descending=d.get('descending', None), field=_enum(d, 'field', SortSpecField))
+
+
+class SortSpecField(Enum):
+
+    POLICY_NAME = 'POLICY_NAME'
+
+
 @dataclass
 class UpdateBudgetConfigurationBudget:
     account_id: Optional[str] = None
@@ -1315,6 +1522,147 @@ class BillableUsageAPI:
         return DownloadResponse.from_dict(res)
 
 
+class BudgetPolicyAPI:
+    """A service serves REST API about Budget policies"""
+
+    def __init__(self, api_client):
+        self._api = api_client
+
+    def create(self,
+               *,
+               custom_tags: Optional[List[compute.CustomPolicyTag]] = None,
+               policy_name: Optional[str] = None,
+               request_id: Optional[str] = None) -> BudgetPolicy:
+        """Create a budget policy.
+        
+        Creates a new policy.
+        
+        :param custom_tags: List[:class:`CustomPolicyTag`] (optional)
+          A list of tags defined by the customer. At most 40 entries are allowed per policy.
+        :param policy_name: str (optional)
+          The name of the policy. - Must be unique among active policies. - Can contain only characters of
+          0-9, a-z, A-Z, -, =, ., :, /, @, _, +, whitespace.
+        :param request_id: str (optional)
+          A unique identifier for this request. Restricted to 36 ASCII characters. A random UUID is
+          recommended. This request is only idempotent if a `request_id` is provided.
+        
+        :returns: :class:`BudgetPolicy`
+        """
+        body = {}
+        if custom_tags is not None: body['custom_tags'] = [v.as_dict() for v in custom_tags]
+        if policy_name is not None: body['policy_name'] = policy_name
+        if request_id is not None: body['request_id'] = request_id
+        headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
+        res = self._api.do('POST',
+                           f'/api/2.1/accounts/{self._api.account_id}/budget-policies',
+                           body=body,
+                           headers=headers)
+        return BudgetPolicy.from_dict(res)
+
+    def delete(self, policy_id: str):
+        """Delete a budget policy.
+        
+        Deletes a policy
+        
+        :param policy_id: str
+          The Id of the policy.
+        
+        
+        """
+
+        headers = {'Accept': 'application/json', }
+
+        self._api.do('DELETE',
+                     f'/api/2.1/accounts/{self._api.account_id}/budget-policies/{policy_id}',
+                     headers=headers)
+
+    def get(self, policy_id: str) -> BudgetPolicy:
+        """Get a budget policy.
+        
+        Retrieves a policy by it's ID.
+        
+        :param policy_id: str
+          The Id of the policy.
+        
+        :returns: :class:`BudgetPolicy`
+        """
+
+        headers = {'Accept': 'application/json', }
+
+        res = self._api.do('GET',
+                           f'/api/2.1/accounts/{self._api.account_id}/budget-policies/{policy_id}',
+                           headers=headers)
+        return BudgetPolicy.from_dict(res)
+
+    def list(self,
+             *,
+             filter_by: Optional[Filter] = None,
+             page_size: Optional[int] = None,
+             page_token: Optional[str] = None,
+             sort_spec: Optional[SortSpec] = None) -> Iterator[BudgetPolicy]:
+        """List policies.
+        
+        Lists all policies. Policies are returned in the alphabetically ascending order of their names.
+        
+        :param filter_by: :class:`Filter` (optional)
+          A filter to apply to the list of policies.
+        :param page_size: int (optional)
+          The maximum number of budget policies to return. If unspecified, at most 100 budget policies will be
+          returned. The maximum value is 1000; values above 1000 will be coerced to 1000.
+        :param page_token: str (optional)
+          A page token, received from a previous `ListServerlessPolicies` call. Provide this to retrieve the
+          subsequent page. If unspecified, the first page will be returned.
+          
+          When paginating, all other parameters provided to `ListServerlessPoliciesRequest` must match the
+          call that provided the page token.
+        :param sort_spec: :class:`SortSpec` (optional)
+          The sort specification.
+        
+        :returns: Iterator over :class:`BudgetPolicy`
+        """
+
+        query = {}
+        if filter_by is not None: query['filter_by'] = filter_by.as_dict()
+        if page_size is not None: query['page_size'] = page_size
+        if page_token is not None: query['page_token'] = page_token
+        if sort_spec is not None: query['sort_spec'] = sort_spec.as_dict()
+        headers = {'Accept': 'application/json', }
+
+        while True:
+            json = self._api.do('GET',
+                                f'/api/2.1/accounts/{self._api.account_id}/budget-policies',
+                                query=query,
+                                headers=headers)
+            if 'policies' in json:
+                for v in json['policies']:
+                    yield BudgetPolicy.from_dict(v)
+            if 'next_page_token' not in json or not json['next_page_token']:
+                return
+            query['page_token'] = json['next_page_token']
+
+    def update(self, policy_id: str, *, policy: Optional[BudgetPolicy] = None) -> BudgetPolicy:
+        """Update a budget policy.
+        
+        Updates a policy
+        
+        :param policy_id: str
+          The Id of the policy. This field is generated by Databricks and globally unique.
+        :param policy: :class:`BudgetPolicy` (optional)
+          Contains the BudgetPolicy details.
+        
+        :returns: :class:`BudgetPolicy`
+        """
+        body = policy.as_dict()
+        headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
+        res = self._api.do('PATCH',
+                           f'/api/2.1/accounts/{self._api.account_id}/budget-policies/{policy_id}',
+                           body=body,
+                           headers=headers)
+        return BudgetPolicy.from_dict(res)
+
+
 class BudgetsAPI:
     """These APIs manage budget configurations for this account. Budgets enable you to monitor usage across your
     account. You can set up budgets to either track account-wide spending, or apply filters to track the