databricks-sdk 0.44.1__py3-none-any.whl → 0.46.0__py3-none-any.whl

databricks/sdk/_base_client.py

@@ -17,7 +17,7 @@ from .errors import DatabricksError, _ErrorCustomizer, _Parser
 from .logger import RoundTrip
 from .retries import retried
 
-logger = logging.getLogger('databricks.sdk')
+logger = logging.getLogger("databricks.sdk")
 
 
 def _fix_host_if_needed(host: Optional[str]) -> Optional[str]:
@@ -25,35 +25,37 @@ def _fix_host_if_needed(host: Optional[str]) -> Optional[str]:
         return host
 
     # Add a default scheme if it's missing
-    if '://' not in host:
-        host = 'https://' + host
+    if "://" not in host:
+        host = "https://" + host
 
     o = urllib.parse.urlparse(host)
     # remove trailing slash
-    path = o.path.rstrip('/')
+    path = o.path.rstrip("/")
     # remove port if 443
     netloc = o.netloc
     if o.port == 443:
-        netloc = netloc.split(':')[0]
+        netloc = netloc.split(":")[0]
 
     return urllib.parse.urlunparse((o.scheme, netloc, path, o.params, o.query, o.fragment))
 
 
 class _BaseClient:
 
-    def __init__(self,
-                 debug_truncate_bytes: int = None,
-                 retry_timeout_seconds: int = None,
-                 user_agent_base: str = None,
-                 header_factory: Callable[[], dict] = None,
-                 max_connection_pools: int = None,
-                 max_connections_per_pool: int = None,
-                 pool_block: bool = True,
-                 http_timeout_seconds: float = None,
-                 extra_error_customizers: List[_ErrorCustomizer] = None,
-                 debug_headers: bool = False,
-                 clock: Clock = None,
-                 streaming_buffer_size: int = 1024 * 1024): # 1MB
+    def __init__(
+        self,
+        debug_truncate_bytes: Optional[int] = None,
+        retry_timeout_seconds: Optional[int] = None,
+        user_agent_base: Optional[str] = None,
+        header_factory: Optional[Callable[[], dict]] = None,
+        max_connection_pools: Optional[int] = None,
+        max_connections_per_pool: Optional[int] = None,
+        pool_block: Optional[bool] = True,
+        http_timeout_seconds: Optional[float] = None,
+        extra_error_customizers: Optional[List[_ErrorCustomizer]] = None,
+        debug_headers: Optional[bool] = False,
+        clock: Optional[Clock] = None,
+        streaming_buffer_size: int = 1024 * 1024,
+    ):  # 1MB
         """
         :param debug_truncate_bytes:
         :param retry_timeout_seconds:
@@ -87,9 +89,11 @@ class _BaseClient:
         # We don't use `max_retries` from HTTPAdapter to align with a more production-ready
         # retry strategy established in the Databricks SDK for Go. See _is_retryable and
         # @retried for more details.
-        http_adapter = requests.adapters.HTTPAdapter(pool_connections=max_connections_per_pool or 20,
-                                                     pool_maxsize=max_connection_pools or 20,
-                                                     pool_block=pool_block)
+        http_adapter = requests.adapters.HTTPAdapter(
+            pool_connections=max_connections_per_pool or 20,
+            pool_maxsize=max_connection_pools or 20,
+            pool_block=pool_block,
+        )
         self._session.mount("https://", http_adapter)
 
         # Default to 60 seconds
@@ -110,7 +114,7 @@ class _BaseClient:
         # See: https://github.com/databricks/databricks-sdk-py/issues/142
         if query is None:
             return None
-        with_fixed_bools = {k: v if type(v) != bool else ('true' if v else 'false') for k, v in query.items()}
+        with_fixed_bools = {k: v if type(v) != bool else ("true" if v else "false") for k, v in query.items()}
 
         # Query parameters may be nested, e.g.
         # {'filter_by': {'user_ids': [123, 456]}}
@@ -140,30 +144,34 @@ class _BaseClient:
             return False
         return data.seekable()
 
-    def do(self,
-           method: str,
-           url: str,
-           query: dict = None,
-           headers: dict = None,
-           body: dict = None,
-           raw: bool = False,
-           files=None,
-           data=None,
-           auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None,
-           response_headers: List[str] = None) -> Union[dict, list, BinaryIO]:
+    def do(
+        self,
+        method: str,
+        url: str,
+        query: Optional[dict] = None,
+        headers: Optional[dict] = None,
+        body: Optional[dict] = None,
+        raw: bool = False,
+        files=None,
+        data=None,
+        auth: Optional[Callable[[requests.PreparedRequest], requests.PreparedRequest]] = None,
+        response_headers: Optional[List[str]] = None,
+    ) -> Union[dict, list, BinaryIO]:
         if headers is None:
             headers = {}
-        headers['User-Agent'] = self._user_agent_base
+        headers["User-Agent"] = self._user_agent_base
 
         # Wrap strings and bytes in a seekable stream so that we can rewind them.
         if isinstance(data, (str, bytes)):
-            data = io.BytesIO(data.encode('utf-8') if isinstance(data, str) else data)
+            data = io.BytesIO(data.encode("utf-8") if isinstance(data, str) else data)
 
         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)
+            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.
@@ -173,25 +181,29 @@ class _BaseClient:
                 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)
+            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,
-                        query=query,
-                        headers=headers,
-                        body=body,
-                        raw=raw,
-                        files=files,
-                        data=data,
-                        auth=auth)
+        response = call(
+            method,
+            url,
+            query=query,
+            headers=headers,
+            body=body,
+            raw=raw,
+            files=files,
+            data=data,
+            auth=auth,
+        )
 
         resp = dict()
         for header in response_headers if response_headers else []:
@@ -220,6 +232,7 @@ class _BaseClient:
         # and Databricks SDK for Go retries
         # (see https://github.com/databricks/databricks-sdk-go/blob/main/apierr/errors.go)
         from urllib3.exceptions import ProxyError
+
         if isinstance(err, ProxyError):
             err = err.original_error
         if isinstance(err, requests.ConnectionError):
@@ -230,48 +243,55 @@ class _BaseClient:
             #
             # return a simple string for debug log readability, as `raise TimeoutError(...) from err`
             # will bubble up the original exception in case we reach max retries.
-            return f'cannot connect'
+            return f"cannot connect"
         if isinstance(err, requests.Timeout):
             # corresponds to `TLS handshake timeout` and `i/o timeout` in Go.
             #
             # return a simple string for debug log readability, as `raise TimeoutError(...) from err`
             # will bubble up the original exception in case we reach max retries.
-            return f'timeout'
+            return f"timeout"
         if isinstance(err, DatabricksError):
             message = str(err)
             transient_error_string_matches = [
                 "com.databricks.backend.manager.util.UnknownWorkerEnvironmentException",
-                "does not have any associated worker environments", "There is no worker environment with id",
-                "Unknown worker environment", "ClusterNotReadyException", "Unexpected error",
+                "does not have any associated worker environments",
+                "There is no worker environment with id",
+                "Unknown worker environment",
+                "ClusterNotReadyException",
+                "Unexpected error",
                 "Please try again later or try a faster operation.",
                 "RPC token bucket limit has been exceeded",
             ]
             for substring in transient_error_string_matches:
                 if substring not in message:
                     continue
-                return f'matched {substring}'
+                return f"matched {substring}"
         return None
 
-    def _perform(self,
-                 method: str,
-                 url: str,
-                 query: dict = None,
-                 headers: dict = None,
-                 body: dict = None,
-                 raw: bool = False,
-                 files=None,
-                 data=None,
-                 auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None):
-        response = self._session.request(method,
-                                         url,
-                                         params=self._fix_query_string(query),
-                                         json=body,
-                                         headers=headers,
-                                         files=files,
-                                         data=data,
-                                         auth=auth,
-                                         stream=raw,
-                                         timeout=self._http_timeout_seconds)
+    def _perform(
+        self,
+        method: str,
+        url: str,
+        query: Optional[dict] = None,
+        headers: Optional[dict] = None,
+        body: Optional[dict] = None,
+        raw: bool = False,
+        files=None,
+        data=None,
+        auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None,
+    ):
+        response = self._session.request(
+            method,
+            url,
+            params=self._fix_query_string(query),
+            json=body,
+            headers=headers,
+            files=files,
+            data=data,
+            auth=auth,
+            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:
@@ -305,14 +325,14 @@ class _StreamingResponse(BinaryIO):
     _closed: bool = False
 
     def fileno(self) -> int:
-        pass
+        return 0
 
-    def flush(self) -> int:
-        pass
+    def flush(self) -> int:  # type: ignore
+        return 0
 
     def __init__(self, response: _RawResponse, chunk_size: Union[int, None] = None):
         self._response = response
-        self._buffer = b''
+        self._buffer = b""
         self._content = None
         self._chunk_size = chunk_size
 
@@ -338,14 +358,14 @@ class _StreamingResponse(BinaryIO):
 
     def read(self, n: int = -1) -> bytes:
         """
-        Read up to n bytes from the response stream. If n is negative, read 
-        until the end of the stream. 
+        Read up to n bytes from the response stream. If n is negative, read
+        until the end of the stream.
         """
 
         self._open()
         read_everything = n < 0
         remaining_bytes = n
-        res = b''
+        res = b""
         while remaining_bytes > 0 or read_everything:
             if len(self._buffer) == 0:
                 try:
@@ -383,10 +403,10 @@ class _StreamingResponse(BinaryIO):
     def writable(self) -> bool:
         return False
 
-    def write(self, s: Union[bytes, bytearray]) -> int:
+    def write(self, s: Union[bytes, bytearray]) -> int:  # type: ignore
         raise NotImplementedError()
 
-    def writelines(self, lines: Iterable[bytes]) -> None:
+    def writelines(self, lines: Iterable[bytes]) -> None:  # type: ignore
         raise NotImplementedError()
 
     def __next__(self) -> bytes:
@@ -395,8 +415,12 @@ class _StreamingResponse(BinaryIO):
     def __iter__(self) -> Iterator[bytes]:
         return self._content
 
-    def __exit__(self, t: Union[Type[BaseException], None], value: Union[BaseException, None],
-                 traceback: Union[TracebackType, None]) -> None:
+    def __exit__(
+        self,
+        t: Union[Type[BaseException], None],
+        value: Union[BaseException, None],
+        traceback: Union[TracebackType, None],
+    ) -> None:
         self._content = None
-        self._buffer = b''
+        self._buffer = b""
         self.close()

databricks/sdk/_property.py

@@ -16,8 +16,9 @@ class _cached_property:
         if self.attrname is None:
             self.attrname = name
         elif name != self.attrname:
-            raise TypeError("Cannot assign the same cached_property to two different names "
-                            f"({self.attrname!r} and {name!r}).")
+            raise TypeError(
+                "Cannot assign the same cached_property to two different names " f"({self.attrname!r} and {name!r})."
+            )
 
     def __get__(self, instance, owner=None):
         if instance is None:
@@ -26,9 +27,11 @@ class _cached_property:
             raise TypeError("Cannot use cached_property instance without calling __set_name__ on it.")
         try:
             cache = instance.__dict__
-        except AttributeError: # not all objects have __dict__ (e.g. class defines slots)
-            msg = (f"No '__dict__' attribute on {type(instance).__name__!r} "
-                   f"instance to cache {self.attrname!r} property.")
+        except AttributeError:  # not all objects have __dict__ (e.g. class defines slots)
+            msg = (
+                f"No '__dict__' attribute on {type(instance).__name__!r} "
+                f"instance to cache {self.attrname!r} property."
+            )
             raise TypeError(msg) from None
         val = cache.get(self.attrname, _NOT_FOUND)
         if val is _NOT_FOUND:
@@ -36,7 +39,9 @@ class _cached_property:
             try:
                 cache[self.attrname] = val
             except TypeError:
-                msg = (f"The '__dict__' attribute on {type(instance).__name__!r} instance "
-                       f"does not support item assignment for caching {self.attrname!r} property.")
+                msg = (
+                    f"The '__dict__' attribute on {type(instance).__name__!r} instance "
+                    f"does not support item assignment for caching {self.attrname!r} property."
+                )
                 raise TypeError(msg) from None
         return val

databricks/sdk/_widgets/__init__.py

@@ -43,7 +43,17 @@ try:
     # Detect if we are in an interactive notebook by iterating over the mro of the current ipython instance,
     # to find ZMQInteractiveShell (jupyter). When used from REPL or file, this check will fail, since the
     # mro only contains TerminalInteractiveShell.
-    if len(list(filter(lambda i: i.__name__ == 'ZMQInteractiveShell', get_ipython().__class__.__mro__))) == 0:
+    if (
+        len(
+            list(
+                filter(
+                    lambda i: i.__name__ == "ZMQInteractiveShell",
+                    get_ipython().__class__.__mro__,
+                )
+            )
+        )
+        == 0
+    ):
         logging.debug("Not in an interactive notebook. Skipping ipywidgets implementation for dbutils.")
         raise EnvironmentError("Not in an interactive notebook.")
 
@@ -61,7 +71,8 @@ try:
         warnings.warn(
             "\nTo use databricks widgets interactively in your notebook, please install databricks sdk using:\n"
             "\tpip install 'databricks-sdk[notebook]'\n"
-            "Falling back to default_value_only implementation for databricks widgets.")
+            "Falling back to default_value_only implementation for databricks widgets."
+        )
         logging.debug(f"{e.msg}. Skipping ipywidgets implementation for dbutils.")
         raise e
 

databricks/sdk/_widgets/default_widgets_utils.py

@@ -11,25 +11,31 @@ class DefaultValueOnlyWidgetUtils(WidgetUtils):
     def text(self, name: str, defaultValue: str, label: typing.Optional[str] = None):
         self._widgets[name] = defaultValue
 
-    def dropdown(self,
-                 name: str,
-                 defaultValue: str,
-                 choices: typing.List[str],
-                 label: typing.Optional[str] = None):
+    def dropdown(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
         self._widgets[name] = defaultValue
 
-    def combobox(self,
-                 name: str,
-                 defaultValue: str,
-                 choices: typing.List[str],
-                 label: typing.Optional[str] = None):
+    def combobox(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
         self._widgets[name] = defaultValue
 
-    def multiselect(self,
-                    name: str,
-                    defaultValue: str,
-                    choices: typing.List[str],
-                    label: typing.Optional[str] = None):
+    def multiselect(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
         self._widgets[name] = defaultValue
 
     def _get(self, name: str) -> str:

databricks/sdk/_widgets/ipywidgets_utils.py

@@ -28,9 +28,9 @@ class DbUtilsWidget:
         if type(value) == str or value is None:
             return value
         if type(value) == list or type(value) == tuple:
-            return ','.join(value)
+            return ",".join(value)
 
-        raise ValueError("The returned value has invalid type (" + type(value) + ").")
+        raise ValueError(f"The returned value has invalid type ({type(value)}).")
 
 
 class IPyWidgetUtil(WidgetUtils):
@@ -38,7 +38,12 @@ class IPyWidgetUtil(WidgetUtils):
     def __init__(self) -> None:
         self._widgets: typing.Dict[str, DbUtilsWidget] = {}
 
-    def _register(self, name: str, widget: ValueWidget, label: typing.Optional[str] = None):
+    def _register(
+        self,
+        name: str,
+        widget: ValueWidget,
+        label: typing.Optional[str] = None,
+    ):
         label = label if label is not None else name
         w = DbUtilsWidget(label, widget)
 
@@ -51,29 +56,47 @@ class IPyWidgetUtil(WidgetUtils):
     def text(self, name: str, defaultValue: str, label: typing.Optional[str] = None):
         self._register(name, widget_string.Text(defaultValue), label)
 
-    def dropdown(self,
-                 name: str,
-                 defaultValue: str,
-                 choices: typing.List[str],
-                 label: typing.Optional[str] = None):
-        self._register(name, widget_selection.Dropdown(value=defaultValue, options=choices), label)
-
-    def combobox(self,
-                 name: str,
-                 defaultValue: str,
-                 choices: typing.List[str],
-                 label: typing.Optional[str] = None):
-        self._register(name, widget_string.Combobox(value=defaultValue, options=choices), label)
-
-    def multiselect(self,
-                    name: str,
-                    defaultValue: str,
-                    choices: typing.List[str],
-                    label: typing.Optional[str] = None):
+    def dropdown(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
         self._register(
             name,
-            widget_selection.SelectMultiple(value=(defaultValue, ),
-                                            options=[("__EMPTY__", ""), *list(zip(choices, choices))]), label)
+            widget_selection.Dropdown(value=defaultValue, options=choices),
+            label,
+        )
+
+    def combobox(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
+        self._register(
+            name,
+            widget_string.Combobox(value=defaultValue, options=choices),
+            label,
+        )
+
+    def multiselect(
+        self,
+        name: str,
+        defaultValue: str,
+        choices: typing.List[str],
+        label: typing.Optional[str] = None,
+    ):
+        self._register(
+            name,
+            widget_selection.SelectMultiple(
+                value=(defaultValue,),
+                options=[("__EMPTY__", ""), *list(zip(choices, choices))],
+            ),
+            label,
+        )
 
     def _get(self, name: str) -> str:
         return self._widgets[name].value

databricks/sdk/azure.py

@@ -4,14 +4,14 @@ from .oauth import TokenSource
 from .service.provisioning import Workspace
 
 
-def add_workspace_id_header(cfg: 'Config', headers: Dict[str, str]):
+def add_workspace_id_header(cfg: "Config", headers: Dict[str, str]):
     if cfg.azure_workspace_resource_id:
         headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id
 
 
-def add_sp_management_token(token_source: 'TokenSource', headers: Dict[str, str]):
+def add_sp_management_token(token_source: "TokenSource", headers: Dict[str, str]):
     mgmt_token = token_source.token()
-    headers['X-Databricks-Azure-SP-Management-Token'] = mgmt_token.access_token
+    headers["X-Databricks-Azure-SP-Management-Token"] = mgmt_token.access_token
 
 
 def get_azure_resource_id(workspace: Workspace):
@@ -22,6 +22,8 @@ def get_azure_resource_id(workspace: Workspace):
     """
     if workspace.azure_workspace_info is None:
         return None
-    return (f'/subscriptions/{workspace.azure_workspace_info.subscription_id}'
-            f'/resourceGroups/{workspace.azure_workspace_info.resource_group}'
-            f'/providers/Microsoft.Databricks/workspaces/{workspace.workspace_name}')
+    return (
+        f"/subscriptions/{workspace.azure_workspace_info.subscription_id}"
+        f"/resourceGroups/{workspace.azure_workspace_info.resource_group}"
+        f"/providers/Microsoft.Databricks/workspaces/{workspace.workspace_name}"
+    )

databricks/sdk/casing.py

@@ -8,22 +8,22 @@ class _Name(object):
         for ch in raw_name:
             if ch.isupper():
                 if segment:
-                    self._segments.append(''.join(segment))
+                    self._segments.append("".join(segment))
                 segment = [ch.lower()]
             elif ch.islower():
                 segment.append(ch)
             else:
                 if segment:
-                    self._segments.append(''.join(segment))
+                    self._segments.append("".join(segment))
                 segment = []
         if segment:
-            self._segments.append(''.join(segment))
+            self._segments.append("".join(segment))
 
     def to_snake_case(self) -> str:
-        return '_'.join(self._segments)
+        return "_".join(self._segments)
 
     def to_header_case(self) -> str:
-        return '-'.join([s.capitalize() for s in self._segments])
+        return "-".join([s.capitalize() for s in self._segments])
 
 
 class Casing(object):