star

venv/lib/python3.9/site-packages/databricks/sdk/_base_client.py

@@ -0,0 +1,429 @@
+import io
+import logging
+import urllib.parse
+from abc import ABC, abstractmethod
+from datetime import timedelta
+from types import TracebackType
+from typing import (Any, BinaryIO, Callable, Dict, Iterable, Iterator, List,
+                    Optional, Type, Union)
+
+import requests
+import requests.adapters
+
+from . import useragent
+from .casing import Casing
+from .clock import Clock, RealClock
+from .errors import DatabricksError, _ErrorCustomizer, _Parser
+from .logger import RoundTrip
+from .retries import retried
+
+logger = logging.getLogger("databricks.sdk")
+
+
+def _fix_host_if_needed(host: Optional[str]) -> Optional[str]:
+    if not host:
+        return host
+
+    # Add a default scheme if it's missing
+    if "://" not in host:
+        host = "https://" + host
+
+    o = urllib.parse.urlparse(host)
+    # remove trailing slash
+    path = o.path.rstrip("/")
+    # remove port if 443
+    netloc = o.netloc
+    if o.port == 443:
+        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: 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:
+        :param user_agent_base:
+        :param header_factory: A function that returns a dictionary of headers to include in the request.
+        :param max_connection_pools: Number of urllib3 connection pools to cache before discarding the least
+            recently used pool. Python requests default value is 10.
+        :param max_connections_per_pool: The maximum number of connections to save in the pool. Improves performance
+            in multithreaded situations. For now, we're setting it to the same value as connection_pool_size.
+        :param pool_block: If pool_block is False, then more connections will are created, but not saved after the
+            first use. Blocks when no free connections are available. urllib3 ensures that no more than
+            pool_maxsize connections are used at a time. Prevents platform from flooding. By default, requests library
+            doesn't block.
+        :param http_timeout_seconds:
+        :param extra_error_customizers:
+        :param debug_headers: Whether to include debug headers in the request log.
+        :param clock: Clock object to use for time-related operations.
+        :param streaming_buffer_size: The size of the buffer to use for streaming responses.
+        """
+
+        self._debug_truncate_bytes = debug_truncate_bytes or 96
+        self._debug_headers = debug_headers
+        self._retry_timeout_seconds = retry_timeout_seconds or 300
+        self._user_agent_base = user_agent_base or useragent.to_string()
+        self._header_factory = header_factory
+        self._clock = clock or RealClock()
+        self._session = requests.Session()
+        self._session.auth = self._authenticate
+        self._streaming_buffer_size = streaming_buffer_size
+
+        # 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,
+        )
+        self._session.mount("https://", http_adapter)
+
+        # Default to 60 seconds
+        self._http_timeout_seconds = http_timeout_seconds or 60
+
+        self._error_parser = _Parser(
+            extra_error_customizers=extra_error_customizers,
+            debug_headers=debug_headers,
+        )
+
+    def _authenticate(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
+        if self._header_factory:
+            headers = self._header_factory()
+            for k, v in headers.items():
+                r.headers[k] = v
+        return r
+
+    @staticmethod
+    def _fix_query_string(query: Optional[dict] = None) -> Optional[dict]:
+        # Convert True -> "true" for Databricks APIs to understand booleans.
+        # 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()}
+
+        # Query parameters may be nested, e.g.
+        # {'filter_by': {'user_ids': [123, 456]}}
+        # The HTTP-compatible representation of this is
+        # filter_by.user_ids=123&filter_by.user_ids=456
+        # To achieve this, we convert the above dictionary to
+        # {'filter_by.user_ids': [123, 456]}
+        # See the following for more information:
+        # https://cloud.google.com/endpoints/docs/grpc-service-config/reference/rpc/google.api#google.api.HttpRule
+        def flatten_dict(d: Dict[str, Any]) -> Dict[str, Any]:
+            for k1, v1 in d.items():
+                if isinstance(v1, dict):
+                    v1 = dict(flatten_dict(v1))
+                    for k2, v2 in v1.items():
+                        yield f"{k1}.{k2}", v2
+                else:
+                    yield k1, v1
+
+        flattened = dict(flatten_dict(with_fixed_bools))
+        return flattened
+
+    @staticmethod
+    def _is_seekable_stream(data) -> bool:
+        if data is None:
+            return False
+        if not isinstance(data, io.IOBase):
+            return False
+        return data.seekable()
+
+    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
+
+        # 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)
+
+        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,
+            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 []:
+            resp[header] = response.headers.get(Casing.to_header_case(header))
+        if raw:
+            streaming_response = _StreamingResponse(response)
+            streaming_response.set_chunk_size(self._streaming_buffer_size)
+            resp["contents"] = streaming_response
+            return resp
+        if not len(response.content):
+            return resp
+
+        json_response = response.json()
+        if json_response is None:
+            return resp
+
+        if isinstance(json_response, list):
+            return json_response
+
+        return {**resp, **json_response}
+
+    @staticmethod
+    def _is_retryable(err: BaseException) -> Optional[str]:
+        # this method is Databricks-specific port of urllib3 retries
+        # (see https://github.com/urllib3/urllib3/blob/main/src/urllib3/util/retry.py)
+        # 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):
+            # corresponds to `connection reset by peer` and `connection refused` errors from Go,
+            # which are generally related to the temporary glitches in the networking stack,
+            # also caused by endpoint protection software, like ZScaler, to drop connections while
+            # not yet authenticated.
+            #
+            # 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"
+        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"
+        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",
+                "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 None
+
+    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:
+            raise error from None
+
+        return response
+
+    def _record_request_log(self, response: requests.Response, raw: bool = False) -> None:
+        if not logger.isEnabledFor(logging.DEBUG):
+            return
+        logger.debug(RoundTrip(response, self._debug_headers, self._debug_truncate_bytes, raw).generate())
+
+
+class _RawResponse(ABC):
+
+    @abstractmethod
+    # follows Response signature: https://github.com/psf/requests/blob/main/src/requests/models.py#L799
+    def iter_content(self, chunk_size: int = 1, decode_unicode: bool = False):
+        pass
+
+    @abstractmethod
+    def close(self):
+        pass
+
+
+class _StreamingResponse(BinaryIO):
+    _response: _RawResponse
+    _buffer: bytes
+    _content: Union[Iterator[bytes], None]
+    _chunk_size: Union[int, None]
+    _closed: bool = False
+
+    def fileno(self) -> int:
+        return 0
+
+    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._content = None
+        self._chunk_size = chunk_size
+
+    def _open(self) -> None:
+        if self._closed:
+            raise ValueError("I/O operation on closed file")
+        if not self._content:
+            self._content = self._response.iter_content(chunk_size=self._chunk_size, decode_unicode=False)
+
+    def __enter__(self) -> BinaryIO:
+        self._open()
+        return self
+
+    def set_chunk_size(self, chunk_size: Union[int, None]) -> None:
+        self._chunk_size = chunk_size
+
+    def close(self) -> None:
+        self._response.close()
+        self._closed = True
+
+    def isatty(self) -> bool:
+        return False
+
+    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.
+        """
+
+        self._open()
+        read_everything = n < 0
+        remaining_bytes = n
+        res = b""
+        while remaining_bytes > 0 or read_everything:
+            if len(self._buffer) == 0:
+                try:
+                    self._buffer = next(self._content)
+                except StopIteration:
+                    break
+            bytes_available = len(self._buffer)
+            to_read = bytes_available if read_everything else min(remaining_bytes, bytes_available)
+            res += self._buffer[:to_read]
+            self._buffer = self._buffer[to_read:]
+            remaining_bytes -= to_read
+        return res
+
+    def readable(self) -> bool:
+        return self._content is not None
+
+    def readline(self, __limit: int = ...) -> bytes:
+        raise NotImplementedError()
+
+    def readlines(self, __hint: int = ...) -> List[bytes]:
+        raise NotImplementedError()
+
+    def seek(self, __offset: int, __whence: int = ...) -> int:
+        raise NotImplementedError()
+
+    def seekable(self) -> bool:
+        return False
+
+    def tell(self) -> int:
+        raise NotImplementedError()
+
+    def truncate(self, __size: Union[int, None] = ...) -> int:
+        raise NotImplementedError()
+
+    def writable(self) -> bool:
+        return False
+
+    def write(self, s: Union[bytes, bytearray]) -> int:  # type: ignore
+        raise NotImplementedError()
+
+    def writelines(self, lines: Iterable[bytes]) -> None:  # type: ignore
+        raise NotImplementedError()
+
+    def __next__(self) -> bytes:
+        return self.read(1)
+
+    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:
+        self._content = None
+        self._buffer = b""
+        self.close()

venv/lib/python3.9/site-packages/databricks/sdk/_property.py

@@ -0,0 +1,47 @@
+# Copied from functools.py
+# Remove when Python 3.8 is the minimum supported version.
+
+_NOT_FOUND = object()
+
+
+class _cached_property:
+
+    def __init__(self, func):
+        self.func = func
+        self.attrname = None
+        self.__doc__ = func.__doc__
+        self.__module__ = func.__module__
+
+    def __set_name__(self, owner, name):
+        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})."
+            )
+
+    def __get__(self, instance, owner=None):
+        if instance is None:
+            return self
+        if self.attrname is None:
+            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."
+            )
+            raise TypeError(msg) from None
+        val = cache.get(self.attrname, _NOT_FOUND)
+        if val is _NOT_FOUND:
+            val = self.func(instance)
+            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."
+                )
+                raise TypeError(msg) from None
+        return val

venv/lib/python3.9/site-packages/databricks/sdk/_widgets/__init__.py

@@ -0,0 +1,83 @@
+import logging
+import typing
+import warnings
+from abc import ABC, abstractmethod
+
+
+class WidgetUtils(ABC):
+
+    def get(self, name: str):
+        return self._get(name)
+
+    @abstractmethod
+    def _get(self, name: str) -> str:
+        pass
+
+    def getArgument(self, name: str, defaultValue: typing.Optional[str] = None):
+        try:
+            return self.get(name)
+        except Exception:
+            return defaultValue
+
+    def remove(self, name: str):
+        self._remove(name)
+
+    @abstractmethod
+    def _remove(self, name: str):
+        pass
+
+    def removeAll(self):
+        self._remove_all()
+
+    @abstractmethod
+    def _remove_all(self):
+        pass
+
+
+try:
+    # We only use ipywidgets if we are in a notebook interactive shell otherwise we raise error,
+    # to fallback to using default_widgets. Also, users WILL have IPython in their notebooks (jupyter),
+    # because we DO NOT SUPPORT any other notebook backends, and hence fallback to default_widgets.
+    from IPython.core.getipython import get_ipython
+
+    # 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
+    ):
+        logging.debug("Not in an interactive notebook. Skipping ipywidgets implementation for dbutils.")
+        raise EnvironmentError("Not in an interactive notebook.")
+
+    # For import errors in IPyWidgetUtil, we provide a warning message, prompting users to install the
+    # correct installation group of the sdk.
+    try:
+        from .ipywidgets_utils import IPyWidgetUtil
+
+        widget_impl = IPyWidgetUtil
+        logging.debug("Using ipywidgets implementation for dbutils.")
+
+    except ImportError as e:
+        # Since we are certain that we are in an interactive notebook, we can make assumptions about
+        # formatting and make the warning nicer for the user.
+        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."
+        )
+        logging.debug(f"{e.msg}. Skipping ipywidgets implementation for dbutils.")
+        raise e
+
+except:
+    from .default_widgets_utils import DefaultValueOnlyWidgetUtils
+
+    widget_impl = DefaultValueOnlyWidgetUtils
+    logging.debug("Using default_value_only implementation for dbutils.")

venv/lib/python3.9/site-packages/databricks/sdk/_widgets/default_widgets_utils.py

@@ -0,0 +1,48 @@
+import typing
+
+from . import WidgetUtils
+
+
+class DefaultValueOnlyWidgetUtils(WidgetUtils):
+
+    def __init__(self) -> None:
+        self._widgets: typing.Dict[str, str] = {}
+
+    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,
+    ):
+        self._widgets[name] = defaultValue
+
+    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,
+    ):
+        self._widgets[name] = defaultValue
+
+    def _get(self, name: str) -> str:
+        return self._widgets[name]
+
+    def _remove(self, name: str):
+        del self._widgets[name]
+
+    def _remove_all(self):
+        self._widgets = {}

venv/lib/python3.9/site-packages/databricks/sdk/_widgets/ipywidgets_utils.py

@@ -0,0 +1,110 @@
+import typing
+
+from IPython.core.display_functions import display
+from ipywidgets.widgets import (ValueWidget, Widget, widget_box,
+                                widget_selection, widget_string)
+
+from .default_widgets_utils import WidgetUtils
+
+
+class DbUtilsWidget:
+
+    def __init__(self, label: str, value_widget: ValueWidget) -> None:
+        self.label_widget = widget_string.Label(label)
+        self.value_widget = value_widget
+        self.box = widget_box.Box([self.label_widget, self.value_widget])
+
+    def display(self):
+        display(self.box)
+
+    def close(self):
+        self.label_widget.close()
+        self.value_widget.close()
+        self.box.close()
+
+    @property
+    def value(self):
+        value = self.value_widget.value
+        if type(value) == str or value is None:
+            return value
+        if type(value) == list or type(value) == tuple:
+            return ",".join(value)
+
+        raise ValueError(f"The returned value has invalid type ({type(value)}).")
+
+
+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,
+    ):
+        label = label if label is not None else name
+        w = DbUtilsWidget(label, widget)
+
+        if name in self._widgets:
+            self.remove(name)
+
+        self._widgets[name] = w
+        w.display()
+
+    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,
+    ):
+        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
+
+    def _remove(self, name: str):
+        self._widgets[name].close()
+        del self._widgets[name]
+
+    def _remove_all(self):
+        Widget.close_all()
+        self._widgets = {}

venv/lib/python3.9/site-packages/databricks/sdk/azure.py

@@ -0,0 +1,29 @@
+from typing import Dict
+
+from .oauth import TokenSource
+from .service.provisioning import Workspace
+
+
+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]):
+    mgmt_token = token_source.token()
+    headers["X-Databricks-Azure-SP-Management-Token"] = mgmt_token.access_token
+
+
+def get_azure_resource_id(workspace: Workspace):
+    """
+    Returns the Azure Resource ID for the given workspace, if it is an Azure workspace.
+    :param workspace:
+    :return:
+    """
+    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}"
+    )

venv/lib/python3.9/site-packages/databricks/sdk/casing.py

@@ -0,0 +1,38 @@
+class _Name(object):
+    """Parses a name in camelCase, PascalCase, snake_case, or kebab-case into its segments."""
+
+    def __init__(self, raw_name: str):
+        #
+        self._segments = []
+        segment = []
+        for ch in raw_name:
+            if ch.isupper():
+                if segment:
+                    self._segments.append("".join(segment))
+                segment = [ch.lower()]
+            elif ch.islower():
+                segment.append(ch)
+            else:
+                if segment:
+                    self._segments.append("".join(segment))
+                segment = []
+        if segment:
+            self._segments.append("".join(segment))
+
+    def to_snake_case(self) -> str:
+        return "_".join(self._segments)
+
+    def to_header_case(self) -> str:
+        return "-".join([s.capitalize() for s in self._segments])
+
+
+class Casing(object):
+
+    @staticmethod
+    def to_header_case(name: str) -> str:
+        """
+        Convert a name from camelCase, PascalCase, snake_case, or kebab-case to header-case.
+        :param name:
+        :return:
+        """
+        return _Name(name).to_header_case()

venv/lib/python3.9/site-packages/databricks/sdk/client_types.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+
+class HostType(Enum):
+    """Enum representing the type of Databricks host."""
+
+    ACCOUNTS = "accounts"
+    WORKSPACE = "workspace"
+    UNIFIED = "unified"
+
+
+class ClientType(Enum):
+    """Enum representing the type of client configuration."""
+
+    ACCOUNT = "account"
+    WORKSPACE = "workspace"

venv/lib/python3.9/site-packages/databricks/sdk/clock.py

@@ -0,0 +1,49 @@
+import abc
+import time
+
+
+class Clock(metaclass=abc.ABCMeta):
+
+    @abc.abstractmethod
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+
+    @abc.abstractmethod
+    def sleep(self, seconds: float) -> None:
+        """
+        Delay execution for a given number of seconds.  The argument may be
+        a floating point number for subsecond precision.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+
+
+class RealClock(Clock):
+    """
+    A real clock that uses the ``time`` module to get the current time and sleep.
+    """
+
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+        return time.time()
+
+    def sleep(self, seconds: float) -> None:
+        """
+        Delay execution for a given number of seconds.  The argument may be
+        a floating point number for subsecond precision.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+        time.sleep(seconds)

venv/lib/python3.9/site-packages/databricks/sdk/common/lro.py

@@ -0,0 +1,17 @@
+from datetime import timedelta
+from typing import Optional
+
+
+class LroOptions:
+    """LroOptions is the options for the Long Running Operations.
+    DO NOT USE THIS OPTION. This option is still under development
+    and can be updated in the future without notice.
+    """
+
+    def __init__(self, *, timeout: Optional[timedelta] = None):
+        """
+        Args:
+            timeout: The timeout for the Long Running Operations.
+                if not set, then operation will wait forever.
+        """
+        self.timeout = timeout

venv/lib/python3.9/site-packages/databricks/sdk/common/types/fieldmask.py

@@ -0,0 +1,39 @@
+class FieldMask(object):
+    """Class for FieldMask message type."""
+
+    # This is based on the base implementation from protobuf.
+    # https://pigweed.googlesource.com/third_party/github/protocolbuffers/protobuf/+/HEAD/python/google/protobuf/internal/field_mask.py
+    # The original implementation only works with proto generated classes.
+    # Since our classes are not generated from proto files, we need to implement it manually.
+
+    def __init__(self, field_mask=None):
+        """Initializes the FieldMask."""
+        if field_mask:
+            self.paths = field_mask
+
+    def ToJsonString(self) -> str:
+        """Converts FieldMask to string."""
+        return ",".join(self.paths)
+
+    def FromJsonString(self, value: str) -> None:
+        """Converts string to FieldMask."""
+        if not isinstance(value, str):
+            raise ValueError("FieldMask JSON value not a string: {!r}".format(value))
+        if value:
+            self.paths = value.split(",")
+        else:
+            self.paths = []
+
+    def __eq__(self, other) -> bool:
+        """Check equality based on paths."""
+        if not isinstance(other, FieldMask):
+            return False
+        return self.paths == other.paths
+
+    def __hash__(self) -> int:
+        """Hash based on paths tuple."""
+        return hash(tuple(self.paths))
+
+    def __repr__(self) -> str:
+        """String representation for debugging."""
+        return f"FieldMask(paths={self.paths})"

venv/lib/python3.9/site-packages/databricks/sdk/config.py

@@ -0,0 +1,808 @@
+import configparser
+import copy
+import datetime
+import logging
+import os
+import pathlib
+import re
+import sys
+import urllib.parse
+from typing import Dict, Iterable, List, Optional
+
+import requests
+
+from . import useragent
+from ._base_client import _fix_host_if_needed
+from .client_types import ClientType, HostType
+from .clock import Clock, RealClock
+from .credentials_provider import (CredentialsStrategy, DefaultCredentials,
+                                   OAuthCredentialsProvider)
+from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
+                           DatabricksEnvironment, get_environment_for_hostname)
+from .oauth import (OidcEndpoints, Token, get_account_endpoints,
+                    get_azure_entra_id_workspace_endpoints,
+                    get_endpoints_from_url, get_host_metadata,
+                    get_unified_endpoints, get_workspace_endpoints)
+
+logger = logging.getLogger("databricks.sdk")
+
+
+class ConfigAttribute:
+    """Configuration attribute metadata and descriptor protocols."""
+
+    # name and transform are discovered from Config.__new__
+    name: str = None
+    transform: type = str
+    _custom_transform = None
+
+    def __init__(self, env: str = None, auth: str = None, sensitive: bool = False, transform=None):
+        self.env = env
+        self.auth = auth
+        self.sensitive = sensitive
+        self._custom_transform = transform
+
+    def __get__(self, cfg: "Config", owner):
+        if not cfg:
+            return None
+        return cfg._inner.get(self.name, None)
+
+    def __set__(self, cfg: "Config", value: any):
+        cfg._inner[self.name] = self.transform(value)
+
+    def __repr__(self) -> str:
+        return f"<ConfigAttribute '{self.name}' {self.transform.__name__}>"
+
+
+def _parse_scopes(value):
+    """Parse scopes into a deduplicated, sorted list."""
+    if value is None:
+        return None
+    if isinstance(value, list):
+        result = sorted(set(s for s in value if s))
+        return result if result else None
+    if isinstance(value, str):
+        parsed: list = sorted(set(s for s in re.split(r"[, ]+", value) if s))
+        return parsed if parsed else None
+    return None
+
+
+def with_product(product: str, product_version: str):
+    """[INTERNAL API] Change the product name and version used in the User-Agent header."""
+    useragent.with_product(product, product_version)
+
+
+def with_user_agent_extra(key: str, value: str):
+    """[INTERNAL API] Add extra metadata to the User-Agent header when developing a library."""
+    useragent.with_extra(key, value)
+
+
+class Config:
+    host: str = ConfigAttribute(env="DATABRICKS_HOST")
+    account_id: str = ConfigAttribute(env="DATABRICKS_ACCOUNT_ID")
+    workspace_id: str = ConfigAttribute(env="DATABRICKS_WORKSPACE_ID")
+
+    # Experimental flag to indicate if the host is a unified host (supports both workspace and account APIs)
+    experimental_is_unified_host: bool = ConfigAttribute(env="DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST")
+
+    # [Experimental] OpenID Connect discovery URL. When set, OIDC endpoints are fetched directly
+    # from this URL instead of the default host-type-based well-known endpoint logic.
+    discovery_url: str = ConfigAttribute(env="DATABRICKS_DISCOVERY_URL")
+
+    # PAT token.
+    token: str = ConfigAttribute(env="DATABRICKS_TOKEN", auth="pat", sensitive=True)
+
+    # Audience for OIDC ID token source accepting an audience as a parameter.
+    # For example, the GitHub action ID token source.
+    token_audience: str = ConfigAttribute(env="DATABRICKS_TOKEN_AUDIENCE", auth="github-oidc")
+
+    # Environment variable for OIDC token.
+    oidc_token_env: str = ConfigAttribute(env="DATABRICKS_OIDC_TOKEN_ENV", auth="env-oidc")
+    oidc_token_filepath: str = ConfigAttribute(env="DATABRICKS_OIDC_TOKEN_FILE", auth="file-oidc")
+
+    username: str = ConfigAttribute(env="DATABRICKS_USERNAME", auth="basic")
+    password: str = ConfigAttribute(env="DATABRICKS_PASSWORD", auth="basic", sensitive=True)
+
+    client_id: str = ConfigAttribute(env="DATABRICKS_CLIENT_ID", auth="oauth")
+    client_secret: str = ConfigAttribute(env="DATABRICKS_CLIENT_SECRET", auth="oauth", sensitive=True)
+    profile: str = ConfigAttribute(env="DATABRICKS_CONFIG_PROFILE")
+    config_file: str = ConfigAttribute(env="DATABRICKS_CONFIG_FILE")
+    google_service_account: str = ConfigAttribute(env="DATABRICKS_GOOGLE_SERVICE_ACCOUNT", auth="google")
+    google_credentials: str = ConfigAttribute(env="GOOGLE_CREDENTIALS", auth="google", sensitive=True)
+    azure_workspace_resource_id: str = ConfigAttribute(env="DATABRICKS_AZURE_RESOURCE_ID", auth="azure")
+    azure_use_msi: bool = ConfigAttribute(env="ARM_USE_MSI", auth="azure")
+    azure_client_secret: str = ConfigAttribute(env="ARM_CLIENT_SECRET", auth="azure", sensitive=True)
+    azure_client_id: str = ConfigAttribute(env="ARM_CLIENT_ID", auth="azure")
+    azure_tenant_id: str = ConfigAttribute(env="ARM_TENANT_ID", auth="azure")
+    azure_environment: str = ConfigAttribute(env="ARM_ENVIRONMENT")
+    databricks_cli_path: str = ConfigAttribute(env="DATABRICKS_CLI_PATH")
+    auth_type: str = ConfigAttribute(env="DATABRICKS_AUTH_TYPE")
+    cluster_id: str = ConfigAttribute(env="DATABRICKS_CLUSTER_ID")
+    warehouse_id: str = ConfigAttribute(env="DATABRICKS_WAREHOUSE_ID")
+    serverless_compute_id: str = ConfigAttribute(env="DATABRICKS_SERVERLESS_COMPUTE_ID")
+    skip_verify: bool = ConfigAttribute()
+    http_timeout_seconds: float = ConfigAttribute()
+    debug_truncate_bytes: int = ConfigAttribute(env="DATABRICKS_DEBUG_TRUNCATE_BYTES")
+    debug_headers: bool = ConfigAttribute(env="DATABRICKS_DEBUG_HEADERS")
+    rate_limit: int = ConfigAttribute(env="DATABRICKS_RATE_LIMIT")
+    retry_timeout_seconds: int = ConfigAttribute()
+    metadata_service_url = ConfigAttribute(
+        env="DATABRICKS_METADATA_SERVICE_URL",
+        auth="metadata-service",
+        sensitive=True,
+    )
+    max_connection_pools: int = ConfigAttribute()
+    max_connections_per_pool: int = ConfigAttribute()
+    databricks_environment: Optional[DatabricksEnvironment] = None
+
+    disable_async_token_refresh: bool = ConfigAttribute(env="DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
+
+    disable_experimental_files_api_client: bool = ConfigAttribute(
+        env="DATABRICKS_DISABLE_EXPERIMENTAL_FILES_API_CLIENT"
+    )
+
+    scopes: list = ConfigAttribute(transform=_parse_scopes)
+    authorization_details: str = ConfigAttribute()
+
+    # disable_oauth_refresh_token controls whether a refresh token should be requested
+    # during the U2M authentication flow (default to false).
+    disable_oauth_refresh_token: bool = ConfigAttribute(env="DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN")
+
+    files_ext_client_download_streaming_chunk_size: int = 2 * 1024 * 1024  # 2 MiB
+
+    # When downloading a file, the maximum number of attempts to retry downloading the whole file. Default is no limit.
+    files_ext_client_download_max_total_recovers: Optional[int] = None
+
+    # When downloading a file, the maximum number of attempts to retry downloading from the same offset without progressing.
+    # This is to avoid infinite retrying when the download is not making any progress. Default is 1.
+    files_ext_client_download_max_total_recovers_without_progressing = 1
+
+    # File multipart upload/download parameters
+    # ----------------------
+
+    # Minimal input stream size (bytes) to use multipart / resumable uploads.
+    # For small files it's more efficient to make one single-shot upload request.
+    # When uploading a file, SDK will initially buffer this many bytes from input stream.
+    # This parameter can be less or bigger than multipart_upload_chunk_size.
+    files_ext_multipart_upload_min_stream_size: int = 50 * 1024 * 1024
+
+    # Maximum number of presigned URLs that can be requested at a time.
+    #
+    # The more URLs we request at once, the higher chance is that some of the URLs will expire
+    # before we get to use it. We discover the presigned URL is expired *after* sending the
+    # input stream partition to the server. So to retry the upload of this partition we must rewind
+    # the stream back. In case of a non-seekable stream we cannot rewind, so we'll abort
+    # the upload. To reduce the chance of this, we're requesting presigned URLs one by one
+    # and using them immediately.
+    files_ext_multipart_upload_batch_url_count: int = 1
+
+    # Size of the chunk to use for multipart uploads & downloads.
+    #
+    # The smaller chunk is, the less chance for network errors (or URL get expired),
+    # but the more requests we'll make.
+    # For AWS, minimum is 5Mb: https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
+    # For GCP, minimum is 256 KiB (and also recommended multiple is 256 KiB)
+    # boto uses 8Mb: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.TransferConfig
+    files_ext_multipart_upload_default_part_size: int = 10 * 1024 * 1024  # 10 MiB
+
+    # List of multipart upload part sizes that can be automatically selected
+    files_ext_multipart_upload_part_size_options: List[int] = [
+        10 * 1024 * 1024,  # 10 MiB
+        20 * 1024 * 1024,  # 20 MiB
+        50 * 1024 * 1024,  # 50 MiB
+        100 * 1024 * 1024,  # 100 MiB
+        200 * 1024 * 1024,  # 200 MiB
+        500 * 1024 * 1024,  # 500 MiB
+        1 * 1024 * 1024 * 1024,  # 1 GiB
+        2 * 1024 * 1024 * 1024,  # 2 GiB
+        4 * 1024 * 1024 * 1024,  # 4 GiB
+    ]
+
+    # Maximum size of a single part in multipart upload.
+    # For AWS, maximum is 5 GiB: https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
+    # For Azure, maximum is 4 GiB: https://learn.microsoft.com/en-us/rest/api/storageservices/put-block
+    # For CloudFlare R2, maximum is 5 GiB: https://developers.cloudflare.com/r2/objects/multipart-objects/
+    files_ext_multipart_upload_max_part_size: int = 4 * 1024 * 1024 * 1024  # 4 GiB
+
+    # Default parallel multipart upload concurrency. Set to 10 because of the experiment results show that it
+    # gives good performance result.
+    files_ext_multipart_upload_default_parallelism: int = 10
+
+    # The expiration duration for presigned URLs used in multipart uploads and downloads.
+    # The client will request new presigned URLs if the previous one is expired. The duration should be long enough
+    # to complete the upload or download of a single part.
+    files_ext_multipart_upload_url_expiration_duration: datetime.timedelta = datetime.timedelta(hours=1)
+    files_ext_presigned_download_url_expiration_duration: datetime.timedelta = datetime.timedelta(hours=1)
+
+    # When downloading a file in parallel, how many worker threads to use.
+    files_ext_parallel_download_default_parallelism: int = 10
+
+    # When downloading a file, if the file size is smaller than this threshold,
+    # We'll use a single-threaded download even if the parallel download is enabled.
+    files_ext_parallel_download_min_file_size: int = 50 * 1024 * 1024  # 50 MiB
+
+    # Default chunk size to use when downloading a file in parallel. Not effective for single threaded download.
+    files_ext_parallel_download_default_part_size: int = 10 * 1024 * 1024  # 10 MiB
+
+    # This is not a "wall time" cutoff for the whole upload request,
+    # but a maximum time between consecutive data reception events (even 1 byte) from the server
+    files_ext_network_transfer_inactivity_timeout_seconds: float = 60
+
+    # Cap on the number of custom retries during incremental uploads:
+    # 1) multipart: upload part URL is expired, so new upload URLs must be requested to continue upload
+    # 2) resumable: chunk upload produced a retryable response (or exception), so upload status must be
+    # retrieved to continue the upload.
+    # In these two cases standard SDK retries (which are capped by the `retry_timeout_seconds` option) are not used.
+    # Note that retry counter is reset when upload is successfully resumed.
+    files_ext_multipart_upload_max_retries = 3
+
+    # Cap on the number of custom retries during parallel downloads.
+    files_ext_parallel_download_max_retries = 3
+
+    # Maximum number of retry attempts for FilesExt cloud API operations.
+    # This works in conjunction with retry_timeout_seconds - whichever limit
+    # is hit first will stop the retry loop.
+    experimental_files_ext_cloud_api_max_retries: int = 3
+
+    def __init__(
+        self,
+        *,
+        # Deprecated. Use credentials_strategy instead.
+        credentials_provider: Optional[CredentialsStrategy] = None,
+        credentials_strategy: Optional[CredentialsStrategy] = None,
+        product=None,
+        product_version=None,
+        clock: Optional[Clock] = None,
+        custom_headers: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ):
+        """Initialize a Config object.
+
+        Args:
+            credentials_provider: (Deprecated) Use credentials_strategy instead.
+            credentials_strategy: Custom credentials strategy for authentication.
+            product: Product name for User-Agent header.
+            product_version: Product version for User-Agent header.
+            clock: Clock instance for time-related operations.
+            custom_headers: Optional dictionary of custom HTTP headers to include in all API requests.
+                These headers will be automatically added to every request made by the client.
+                Request-specific headers passed to individual API calls will override these custom headers
+                if there is a conflict. Example: {"X-Request-ID": "123", "X-Custom-Header": "value"}
+            **kwargs: Additional configuration parameters.
+        """
+        self._header_factory = None
+        self._inner = {}
+        self._user_agent_other_info = []
+        self._custom_headers = custom_headers or {}
+        if credentials_strategy and credentials_provider:
+            raise ValueError("When providing `credentials_strategy` field, `credential_provider` cannot be specified.")
+        if credentials_provider:
+            logger.warning("parameter 'credentials_provider' is deprecated. Use 'credentials_strategy' instead.")
+        self._credentials_strategy = next(
+            s
+            for s in [
+                credentials_strategy,
+                credentials_provider,
+                DefaultCredentials(),
+            ]
+            if s is not None
+        )
+        if "databricks_environment" in kwargs:
+            self.databricks_environment = kwargs["databricks_environment"]
+            del kwargs["databricks_environment"]
+        self._clock = clock if clock is not None else RealClock()
+        try:
+            self._set_inner_config(kwargs)
+            self._load_from_env()
+            self._known_file_config_loader()
+            self._fix_host_if_needed()
+            self._validate()
+            self.init_auth()
+            self._init_product(product, product_version)
+        except ValueError as e:
+            message = self.wrap_debug_info(str(e))
+            raise ValueError(message) from e
+
+    def oauth_token(self) -> Token:
+        """Returns the OAuth token from the current credential provider.
+
+        This method only works when using OAuth-based authentication methods.
+        If the current credential provider is an OAuthCredentialsProvider, it reuses
+        the existing provider. Otherwise, it raises a ValueError indicating that
+        OAuth tokens are not available for the current authentication method.
+        """
+        if isinstance(self._header_factory, OAuthCredentialsProvider):
+            return self._header_factory.oauth_token()
+        raise ValueError(
+            f"OAuth tokens are not available for {self.auth_type} authentication. "
+            f"Use an OAuth-based authentication method to access OAuth tokens."
+        )
+
+    def wrap_debug_info(self, message: str) -> str:
+        debug_string = self.debug_string()
+        if debug_string:
+            message = f"{message.rstrip('.')}. {debug_string}"
+        return message
+
+    @staticmethod
+    def parse_dsn(dsn: str) -> "Config":
+        uri = urllib.parse.urlparse(dsn)
+        if uri.scheme != "databricks":
+            raise ValueError(f"Expected databricks:// scheme, got {uri.scheme}://")
+        kwargs = {"host": f"https://{uri.hostname}"}
+        if uri.username:
+            kwargs["username"] = uri.username
+        if uri.password:
+            kwargs["password"] = uri.password
+        query = dict(urllib.parse.parse_qsl(uri.query))
+        for attr in Config.attributes():
+            if attr.name not in query:
+                continue
+            kwargs[attr.name] = query[attr.name]
+        return Config(**kwargs)
+
+    def authenticate(self) -> Dict[str, str]:
+        """Returns a list of fresh authentication headers"""
+        return self._header_factory()
+
+    def as_dict(self) -> dict:
+        return self._inner
+
+    def _get_azure_environment_name(self) -> str:
+        if not self.azure_environment:
+            return "PUBLIC"
+        env = self.azure_environment.upper()
+        # Compatibility with older versions of the SDK that allowed users to specify AzurePublicCloud or AzureChinaCloud
+        if env.startswith("AZURE"):
+            env = env[len("AZURE") :]
+        if env.endswith("CLOUD"):
+            env = env[: -len("CLOUD")]
+        return env
+
+    @property
+    def environment(self) -> DatabricksEnvironment:
+        """Returns the environment based on configuration."""
+        if self.databricks_environment:
+            return self.databricks_environment
+        if not self.host and self.azure_workspace_resource_id:
+            azure_env = self._get_azure_environment_name()
+            for environment in ALL_ENVS:
+                if environment.cloud != Cloud.AZURE:
+                    continue
+                if environment.azure_environment.name != azure_env:
+                    continue
+                if environment.dns_zone.startswith(".dev") or environment.dns_zone.startswith(".staging"):
+                    continue
+                return environment
+        return get_environment_for_hostname(self.host)
+
+    @property
+    def is_azure(self) -> bool:
+        if self.azure_workspace_resource_id:
+            return True
+        return self.environment.cloud == Cloud.AZURE
+
+    @property
+    def is_gcp(self) -> bool:
+        return self.environment.cloud == Cloud.GCP
+
+    @property
+    def is_aws(self) -> bool:
+        return self.environment.cloud == Cloud.AWS
+
+    @property
+    def host_type(self) -> HostType:
+        """Determine the type of host based on the configuration.
+
+        Returns the HostType which can be ACCOUNTS, WORKSPACE, or UNIFIED.
+        """
+        # Check if explicitly marked as unified host
+        if self.experimental_is_unified_host:
+            return HostType.UNIFIED
+
+        if not self.host:
+            return HostType.WORKSPACE
+
+        # Check for accounts host pattern
+        if self.host.startswith("https://accounts.") or self.host.startswith("https://accounts-dod."):
+            return HostType.ACCOUNTS
+
+        return HostType.WORKSPACE
+
+    @property
+    def client_type(self) -> ClientType:
+        """Determine the type of client configuration.
+
+        This is separate from host_type. For example, a unified host can support both
+        workspace and account client types.
+
+        Returns ClientType.ACCOUNT or ClientType.WORKSPACE based on the configuration.
+
+        For unified hosts, account_id must be set. If workspace_id is also set,
+        returns WORKSPACE, otherwise returns ACCOUNT.
+        """
+        host_type = self.host_type
+
+        if host_type == HostType.ACCOUNTS:
+            return ClientType.ACCOUNT
+
+        if host_type == HostType.WORKSPACE:
+            return ClientType.WORKSPACE
+
+        if host_type == HostType.UNIFIED:
+            if not self.account_id:
+                raise ValueError("Unified host requires account_id to be set")
+            if self.workspace_id:
+                return ClientType.WORKSPACE
+            return ClientType.ACCOUNT
+
+        # Default to workspace for backward compatibility
+        return ClientType.WORKSPACE
+
+    @property
+    def is_account_client(self) -> bool:
+        """[Deprecated] Use host_type or client_type instead.
+
+        Determines if this is an account client based on the host URL.
+        """
+        if self.experimental_is_unified_host:
+            raise ValueError(
+                "is_account_client cannot be used with unified hosts; use host_type or client_type instead"
+            )
+        if not self.host:
+            return False
+        return self.host.startswith("https://accounts.") or self.host.startswith("https://accounts-dod.")
+
+    @property
+    def arm_environment(self) -> AzureEnvironment:
+        return self.environment.azure_environment
+
+    @property
+    def effective_azure_login_app_id(self):
+        return self.environment.azure_application_id
+
+    @property
+    def hostname(self) -> str:
+        url = urllib.parse.urlparse(self.host)
+        return url.netloc
+
+    @property
+    def is_any_auth_configured(self) -> bool:
+        for attr in Config.attributes():
+            if not attr.auth:
+                continue
+            value = self._inner.get(attr.name, None)
+            if value:
+                return True
+        return False
+
+    @property
+    def user_agent(self):
+        """Returns User-Agent header used by this SDK"""
+
+        # global user agent includes SDK version, product name & version, platform info,
+        # and global extra info. Config can have specific extra info associated with it,
+        # such as an override product, auth type, and other user-defined information.
+        return useragent.to_string(
+            self._product_info,
+            [("auth", self.auth_type)] + self._user_agent_other_info,
+        )
+
+    @property
+    def _upstream_user_agent(self) -> str:
+        return " ".join(f"{k}/{v}" for k, v in useragent._get_upstream_user_agent_info())
+
+    def with_user_agent_extra(self, key: str, value: str) -> "Config":
+        self._user_agent_other_info.append((key, value))
+        return self
+
+    @property
+    def databricks_oidc_endpoints(self) -> Optional[OidcEndpoints]:
+        """Get OIDC endpoints for Databricks OAuth.
+
+        If discovery_url is set, OIDC endpoints are fetched directly from it. Otherwise
+        falls back to the host-type-based well-known endpoint logic.
+
+        Note: This method does NOT return Azure Entra ID endpoints. For Azure authentication,
+        use get_azure_entra_id_workspace_endpoints() directly.
+
+        Returns:
+            OidcEndpoints for Databricks OAuth, or None if host is not configured.
+        """
+        self._fix_host_if_needed()
+        if not self.host:
+            return None
+
+        if self.discovery_url:
+            return get_endpoints_from_url(self.discovery_url)
+
+        # Handle unified hosts
+        if self.host_type == HostType.UNIFIED:
+            if not self.account_id:
+                raise ValueError("Unified host requires account_id to be set for OAuth endpoints")
+            return get_unified_endpoints(self.host, self.account_id)
+
+        # Handle traditional account hosts
+        if self.host_type == HostType.ACCOUNTS and self.account_id:
+            return get_account_endpoints(self.host, self.account_id)
+
+        # Default to workspace endpoints
+        return get_workspace_endpoints(self.host)
+
+    @property
+    def oidc_endpoints(self) -> Optional[OidcEndpoints]:
+        """[DEPRECATED] Get OIDC endpoints with automatic Azure detection (deprecated).
+
+        This method incorrectly returns Azure OIDC endpoints when azure_client_id
+        is set, even for Databricks OAuth flows that don't use Azure authentication. This caused
+        bugs where Databricks M2M OAuth would fail when ARM_CLIENT_ID was set for other purposes.
+
+        Use instead:
+        - databricks_oidc_endpoints: For Databricks OAuth (oauth-m2m, external-browser, etc.)
+        - get_azure_entra_id_workspace_endpoints(): For Azure Entra ID authentication
+
+        Returns:
+            OidcEndpoints (Azure or Databricks depending on config), or None if host is not configured.
+        """
+        self._fix_host_if_needed()
+        if not self.host:
+            return None
+        if self.is_azure and self.azure_client_id:
+            return get_azure_entra_id_workspace_endpoints(self.host)
+        return self.databricks_oidc_endpoints
+
+    def debug_string(self) -> str:
+        """Returns log-friendly representation of configured attributes"""
+        buf = []
+        attrs_used = []
+        envs_used = []
+        for attr in Config.attributes():
+            if attr.env and os.environ.get(attr.env):
+                envs_used.append(attr.env)
+            value = getattr(self, attr.name)
+            if not value:
+                continue
+            safe = "***" if attr.sensitive else f"{value}"
+            attrs_used.append(f"{attr.name}={safe}")
+        if attrs_used:
+            buf.append(f"Config: {', '.join(attrs_used)}")
+        if envs_used:
+            buf.append(f"Env: {', '.join(envs_used)}")
+        return ". ".join(buf)
+
+    def to_dict(self) -> Dict[str, any]:
+        return self._inner
+
+    @property
+    def sql_http_path(self) -> Optional[str]:
+        """(Experimental) Return HTTP path for SQL Drivers.
+
+        If `cluster_id` or `warehouse_id` are configured, return a valid HTTP Path argument
+        used in construction of JDBC/ODBC DSN string.
+
+        See https://docs.databricks.com/integrations/jdbc-odbc-bi.html
+        """
+        if (not self.cluster_id) and (not self.warehouse_id):
+            return None
+        if self.cluster_id and self.warehouse_id:
+            raise ValueError("cannot have both cluster_id and warehouse_id")
+        headers = self.authenticate()
+        headers["User-Agent"] = f"{self.user_agent} sdk-feature/sql-http-path"
+        if self.cluster_id:
+            response = requests.get(f"{self.host}/api/2.0/preview/scim/v2/Me", headers=headers)
+            # get workspace ID from the response header
+            workspace_id = response.headers.get("x-databricks-org-id")
+            return f"sql/protocolv1/o/{workspace_id}/{self.cluster_id}"
+        if self.warehouse_id:
+            return f"/sql/1.0/warehouses/{self.warehouse_id}"
+
+    @property
+    def clock(self) -> Clock:
+        return self._clock
+
+    @classmethod
+    def attributes(cls) -> Iterable[ConfigAttribute]:
+        """Returns a list of Databricks SDK configuration metadata"""
+        if hasattr(cls, "_attributes"):
+            return cls._attributes
+        if sys.version_info[1] >= 10:
+            import inspect
+
+            anno = inspect.get_annotations(cls)
+        else:
+            # Python 3.7 compatibility: getting type hints require extra hop, as described in
+            # "Accessing The Annotations Dict Of An Object In Python 3.9 And Older" section of
+            # https://docs.python.org/3/howto/annotations.html
+            anno = cls.__dict__["__annotations__"]
+        attrs = []
+        for name, v in cls.__dict__.items():
+            if type(v) != ConfigAttribute:
+                continue
+            v.name = name
+            v.transform = v._custom_transform if v._custom_transform else anno.get(name, str)
+            attrs.append(v)
+        cls._attributes = attrs
+        return cls._attributes
+
+    def _resolve_host_metadata(self) -> None:
+        """[Experimental] Populate missing config fields from the host's
+        /.well-known/databricks-config discovery endpoint.
+
+        Fills in account_id, workspace_id, and discovery_url (derived from oidc_endpoint,
+        with any {account_id} placeholder substituted) if not already set.
+        """
+        if not self.host:
+            return
+        meta = get_host_metadata(self.host)
+        if not self.account_id and meta.account_id:
+            logger.debug(f"Resolved account_id from host metadata: {meta.account_id}")
+            self.account_id = meta.account_id
+        if not self.account_id:
+            raise ValueError("account_id is not configured and could not be resolved from host metadata")
+        if not self.workspace_id and meta.workspace_id:
+            logger.debug(f"Resolved workspace_id from host metadata: {meta.workspace_id}")
+            self.workspace_id = meta.workspace_id
+        if not self.discovery_url:
+            if meta.oidc_endpoint:
+                logger.debug(f"Resolved discovery_url from host metadata: {meta.oidc_endpoint}")
+                self.discovery_url = meta.oidc_endpoint.replace("{account_id}", self.account_id)
+            else:
+                raise ValueError("discovery_url is not configured and could not be resolved from host metadata")
+
+    def _fix_host_if_needed(self):
+        updated_host = _fix_host_if_needed(self.host)
+        if updated_host:
+            self.host = updated_host
+
+    def load_azure_tenant_id(self):
+        """[Internal] Load the Azure tenant ID from the Azure Databricks login page.
+
+        If the tenant ID is already set, this method does nothing."""
+        if self.azure_tenant_id is not None or self.host is None:
+            return
+        login_url = f"{self.host}/aad/auth"
+        logger.debug(f"Loading tenant ID from {login_url}")
+        resp = requests.get(login_url, allow_redirects=False)
+        if resp.status_code // 100 != 3:
+            logger.debug(f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}")
+            return
+        entra_id_endpoint = resp.headers.get("Location")
+        if entra_id_endpoint is None:
+            logger.debug(f"No Location header in response from {login_url}")
+            return
+        # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
+        # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).
+        url = urllib.parse.urlparse(entra_id_endpoint)
+        path_segments = url.path.split("/")
+        if len(path_segments) < 2:
+            logger.debug(f"Invalid path in Location header: {url.path}")
+            return
+        self.azure_tenant_id = path_segments[1]
+        logger.debug(f"Loaded tenant ID: {self.azure_tenant_id}")
+
+    def _set_inner_config(self, keyword_args: Dict[str, any]):
+        for attr in self.attributes():
+            if attr.name not in keyword_args:
+                continue
+            if keyword_args.get(attr.name, None) is None:
+                continue
+            self.__setattr__(attr.name, keyword_args[attr.name])
+
+    def _load_from_env(self):
+        found = False
+        for attr in self.attributes():
+            if not attr.env:
+                continue
+            if attr.name in self._inner:
+                continue
+            value = os.environ.get(attr.env)
+            if not value:
+                continue
+            self.__setattr__(attr.name, value)
+            found = True
+        if found:
+            logger.debug("Loaded from environment")
+
+    def _known_file_config_loader(self):
+        if not self.profile and (self.is_any_auth_configured or self.host or self.azure_workspace_resource_id):
+            # skip loading configuration file if there's any auth configured
+            # directly as part of the Config() constructor.
+            return
+        config_file = self.config_file
+        if not config_file:
+            config_file = "~/.databrickscfg"
+        config_path = pathlib.Path(config_file).expanduser()
+        if not config_path.exists():
+            logger.debug("%s does not exist", config_path)
+            return
+        ini_file = configparser.ConfigParser()
+        ini_file.read(config_path)
+        profile = self.profile
+        has_explicit_profile = self.profile is not None
+        # In Go SDK, we skip merging the profile with DEFAULT section, though Python's ConfigParser.items()
+        # is returning profile key-value pairs _including those from DEFAULT_. This is not what we expect
+        # from Unified Auth test suite at the moment. Hence, the private variable access.
+        # See: https://docs.python.org/3/library/configparser.html#mapping-protocol-access
+        if not has_explicit_profile and not ini_file.defaults():
+            logger.debug(f"{config_path} has no DEFAULT profile configured")
+            return
+        if not has_explicit_profile:
+            profile = "DEFAULT"
+        profiles = ini_file._sections
+        if ini_file.defaults():
+            profiles["DEFAULT"] = ini_file.defaults()
+        if profile not in profiles:
+            raise ValueError(f"resolve: {config_path} has no {profile} profile configured")
+        raw_config = profiles[profile]
+        logger.info(f"loading {profile} profile from {config_file}: {', '.join(raw_config.keys())}")
+        for k, v in raw_config.items():
+            if k in self._inner:
+                # don't overwrite a value previously set
+                continue
+            self.__setattr__(k, v)
+
+    def _validate(self):
+        auths_used = set()
+        for attr in Config.attributes():
+            if attr.name not in self._inner:
+                continue
+            if not attr.auth:
+                continue
+            auths_used.add(attr.auth)
+        if len(auths_used) <= 1:
+            return
+        if self.auth_type:
+            # client has auth preference set
+            return
+        names = " and ".join(sorted(auths_used))
+        raise ValueError(f"validate: more than one authorization method configured: {names}")
+
+    def init_auth(self):
+        try:
+            self._header_factory = self._credentials_strategy(self)
+            self.auth_type = self._credentials_strategy.auth_type()
+            if not self._header_factory:
+                raise ValueError("not configured")
+        except ValueError as e:
+            raise ValueError(f"{self._credentials_strategy.auth_type()} auth: {e}") from e
+
+    def _init_product(self, product, product_version):
+        if product is not None or product_version is not None:
+            default_product, default_version = useragent.product()
+            self._product_info = (
+                product or default_product,
+                product_version or default_version,
+            )
+        else:
+            self._product_info = None
+
+    def get_scopes(self) -> list:
+        """Get OAuth scopes with proper defaulting.
+
+        Returns ["all-apis"] if no scopes configured.
+        This is the single source of truth for scope defaulting across all OAuth methods.
+        """
+        return self.scopes if self.scopes else ["all-apis"]
+
+    def get_scopes_as_string(self) -> str:
+        """Get OAuth scopes as a space-separated string.
+
+        Returns "all-apis" if no scopes configured.
+        """
+        return " ".join(self.get_scopes())
+
+    def __repr__(self):
+        return f"<{self.debug_string()}>"
+
+    def copy(self):
+        """Creates a copy of the config object.
+        All the copies share most of their internal state (ie, shared reference to fields such as credential_provider).
+        Copies have their own instances of the following fields
+            - `_user_agent_other_info`
+        """
+        cpy: Config = copy.copy(self)
+        cpy._user_agent_other_info = copy.deepcopy(self._user_agent_other_info)
+        return cpy
+
+    def deep_copy(self):
+        """Creates a deep copy of the config object."""
+        return copy.deepcopy(self)

venv/lib/python3.9/site-packages/databricks/sdk/core.py

@@ -0,0 +1,115 @@
+import re
+from typing import BinaryIO
+from urllib.parse import urlencode
+
+from ._base_client import _BaseClient
+from .config import *
+# To preserve backwards compatibility (as these definitions were previously in this module)
+from .credentials_provider import *
+from .errors import DatabricksError, _ErrorCustomizer
+from .oauth import retrieve_token
+
+__all__ = ["Config", "DatabricksError"]
+
+logger = logging.getLogger("databricks.sdk")
+
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
+
+class ApiClient:
+
+    def __init__(self, cfg: Config):
+        self._cfg = cfg
+
+        self._api_client = _BaseClient(
+            debug_truncate_bytes=cfg.debug_truncate_bytes,
+            retry_timeout_seconds=cfg.retry_timeout_seconds,
+            user_agent_base=cfg.user_agent,
+            header_factory=cfg.authenticate,
+            max_connection_pools=cfg.max_connection_pools,
+            max_connections_per_pool=cfg.max_connections_per_pool,
+            pool_block=True,
+            http_timeout_seconds=cfg.http_timeout_seconds,
+            extra_error_customizers=[_AddDebugErrorCustomizer(cfg)],
+            clock=cfg.clock,
+        )
+
+    @property
+    def account_id(self) -> str:
+        return self._cfg.account_id
+
+    @property
+    def is_account_client(self) -> bool:
+        return self._cfg.is_account_client
+
+    def get_oauth_token(self, auth_details: str) -> Token:
+        if not self._cfg.auth_type:
+            self._cfg.authenticate()
+        original_token = self._cfg.oauth_token()
+        headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
+        params = urlencode(
+            {
+                "grant_type": JWT_BEARER_GRANT_TYPE,
+                "authorization_details": auth_details,
+                "assertion": original_token.access_token,
+            }
+        )
+        return retrieve_token(
+            client_id=self._cfg.client_id,
+            client_secret=self._cfg.client_secret,
+            token_url=self._cfg.host + OIDC_TOKEN_PATH,
+            params=params,
+            headers=headers,
+        )
+
+    def do(
+        self,
+        method: str,
+        path: Optional[str] = None,
+        url: Optional[str] = None,
+        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 url is None:
+            # Remove extra `/` from path for Files API
+            # Once we've fixed the OpenAPI spec, we can remove this
+            path = re.sub("^/api/2.0/fs/files//", "/api/2.0/fs/files/", path)
+            url = f"{self._cfg.host}{path}"
+
+        # Merge custom headers with request-specific headers
+        # Request-specific headers take precedence
+        merged_headers = {**self._cfg._custom_headers, **(headers or {})}
+
+        return self._api_client.do(
+            method=method,
+            url=url,
+            query=query,
+            headers=merged_headers,
+            body=body,
+            raw=raw,
+            files=files,
+            data=data,
+            auth=auth,
+            response_headers=response_headers,
+        )
+
+
+class _AddDebugErrorCustomizer(_ErrorCustomizer):
+    """An error customizer that adds debug information about the configuration to unauthenticated and
+    unauthorized errors."""
+
+    def __init__(self, cfg: Config):
+        self._cfg = cfg
+
+    def customize_error(self, response: requests.Response, kwargs: dict):
+        if response.status_code in (401, 403):
+            message = kwargs.get("message", "request failed")
+            kwargs["message"] = self._cfg.wrap_debug_info(message)

venv/lib/python3.9/site-packages/databricks/sdk/data_plane.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from typing import Callable, Optional
+from urllib import parse
+
+from databricks.sdk import oauth
+from databricks.sdk.oauth import Token
+
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
+
+class DataPlaneTokenSource:
+    """
+    EXPERIMENTAL Manages token sources for multiple DataPlane endpoints.
+    """
+
+    # TODO: Enable async once its stable. @oauth_credentials_provider must also have async enabled.
+    def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], disable_async: Optional[bool] = True):
+        self._cpts = cpts
+        self._token_exchange_host = token_exchange_host
+        self._token_sources = {}
+        self._disable_async = disable_async
+        self._lock = threading.Lock()
+
+    def token(self, endpoint, auth_details):
+        key = f"{endpoint}:{auth_details}"
+
+        # First, try to read without acquiring the lock to avoid contention.
+        # Reads are atomic, so this is safe.
+        token_source = self._token_sources.get(key)
+        if token_source:
+            return token_source.token()
+
+        # If token_source is not found, acquire the lock and check again.
+        with self._lock:
+            # Another thread might have created it while we were waiting for the lock.
+            token_source = self._token_sources.get(key)
+            if not token_source:
+                token_source = DataPlaneEndpointTokenSource(
+                    self._token_exchange_host, self._cpts, auth_details, self._disable_async
+                )
+                self._token_sources[key] = token_source
+
+        return token_source.token()
+
+
+class DataPlaneEndpointTokenSource(oauth.Refreshable):
+    """
+    EXPERIMENTAL A token source for a specific DataPlane endpoint.
+    """
+
+    def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], auth_details: str, disable_async: bool):
+        super().__init__(disable_async=disable_async)
+        self._auth_details = auth_details
+        self._cpts = cpts
+        self._token_exchange_host = token_exchange_host
+
+    def refresh(self) -> Token:
+        control_plane_token = self._cpts()
+        headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
+        params = parse.urlencode(
+            {
+                "grant_type": JWT_BEARER_GRANT_TYPE,
+                "authorization_details": self._auth_details,
+                "assertion": control_plane_token.access_token,
+            }
+        )
+        return oauth.retrieve_token(
+            client_id="",
+            client_secret="",
+            token_url=self._token_exchange_host + OIDC_TOKEN_PATH,
+            params=params,
+            headers=headers,
+        )
+
+
+@dataclass
+class DataPlaneDetails:
+    """
+    Contains details required to query a DataPlane endpoint.
+    """
+
+    endpoint_url: str
+    """URL used to query the endpoint through the DataPlane."""
+    token: Token
+    """Token to query the DataPlane endpoint."""

venv/lib/python3.9/site-packages/databricks/sdk/dbutils.py

@@ -0,0 +1,475 @@
+import base64
+import json
+import logging
+import os
+import threading
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from .core import ApiClient, Config, DatabricksError
+from .mixins import compute as compute_ext
+from .mixins import files as dbfs_ext
+from .service import compute, workspace
+
+_LOG = logging.getLogger("databricks.sdk")
+
+
+class FileInfo(namedtuple("FileInfo", ["path", "name", "size", "modificationTime"])):
+    pass
+
+
+class MountInfo(namedtuple("MountInfo", ["mountPoint", "source", "encryptionType"])):
+    pass
+
+
+class SecretScope(namedtuple("SecretScope", ["name"])):
+
+    def getName(self):
+        return self.name
+
+
+class SecretMetadata(namedtuple("SecretMetadata", ["key"])):
+    pass
+
+
+class _FsUtil:
+    """Manipulates the Databricks filesystem (DBFS)"""
+
+    def __init__(
+        self,
+        dbfs_ext: dbfs_ext.DbfsExt,
+        proxy_factory: Callable[[str], "_ProxyUtil"],
+    ):
+        self._dbfs = dbfs_ext
+        self._proxy_factory = proxy_factory
+
+    def cp(self, from_: str, to: str, recurse: bool = False) -> bool:
+        """Copies a file or directory, possibly across FileSystems"""
+        self._dbfs.copy(from_, to, recursive=recurse)
+        return True
+
+    def head(self, file: str, maxBytes: int = 65536) -> str:
+        """Returns up to the first 'maxBytes' bytes of the given file as a String encoded in UTF-8"""
+        with self._dbfs.download(file) as f:
+            return f.read(maxBytes).decode("utf8")
+
+    def ls(self, dir: str) -> List[FileInfo]:
+        """Lists the contents of a directory"""
+        return [
+            FileInfo(
+                f.path,
+                os.path.basename(f.path),
+                f.file_size,
+                f.modification_time,
+            )
+            for f in self._dbfs.list(dir)
+        ]
+
+    def mkdirs(self, dir: str) -> bool:
+        """Creates the given directory if it does not exist, also creating any necessary parent directories"""
+        self._dbfs.mkdirs(dir)
+        return True
+
+    def mv(self, from_: str, to: str, recurse: bool = False) -> bool:
+        """Moves a file or directory, possibly across FileSystems"""
+        self._dbfs.move_(from_, to, recursive=recurse, overwrite=True)
+        return True
+
+    def put(self, file: str, contents: str, overwrite: bool = False) -> bool:
+        """Writes the given String out to a file, encoded in UTF-8"""
+        with self._dbfs.open(file, write=True, overwrite=overwrite) as f:
+            f.write(contents.encode("utf8"))
+        return True
+
+    def rm(self, dir: str, recurse: bool = False) -> bool:
+        """Removes a file or directory"""
+        self._dbfs.delete(dir, recursive=recurse)
+        return True
+
+    def mount(
+        self,
+        source: str,
+        mount_point: str,
+        encryption_type: str = None,
+        owner: str = None,
+        extra_configs: Dict[str, str] = None,
+    ) -> bool:
+        """Mounts the given source directory into DBFS at the given mount point"""
+        fs = self._proxy_factory("fs")
+        kwargs = {}
+        if encryption_type:
+            kwargs["encryption_type"] = encryption_type
+        if owner:
+            kwargs["owner"] = owner
+        if extra_configs:
+            kwargs["extra_configs"] = extra_configs
+        return fs.mount(source, mount_point, **kwargs)
+
+    def unmount(self, mount_point: str) -> bool:
+        """Deletes a DBFS mount point"""
+        fs = self._proxy_factory("fs")
+        return fs.unmount(mount_point)
+
+    def updateMount(
+        self,
+        source: str,
+        mount_point: str,
+        encryption_type: str = None,
+        owner: str = None,
+        extra_configs: Dict[str, str] = None,
+    ) -> bool:
+        """Similar to mount(), but updates an existing mount point (if present) instead of creating a new one"""
+        fs = self._proxy_factory("fs")
+        kwargs = {}
+        if encryption_type:
+            kwargs["encryption_type"] = encryption_type
+        if owner:
+            kwargs["owner"] = owner
+        if extra_configs:
+            kwargs["extra_configs"] = extra_configs
+        return fs.updateMount(source, mount_point, **kwargs)
+
+    def mounts(self) -> List[MountInfo]:
+        """Displays information about what is mounted within DBFS"""
+        result = []
+        fs = self._proxy_factory("fs")
+        for info in fs.mounts():
+            result.append(MountInfo(info[0], info[1], info[2]))
+        return result
+
+    def refreshMounts(self) -> bool:
+        """Forces all machines in this cluster to refresh their mount cache,
+        ensuring they receive the most recent information"""
+        fs = self._proxy_factory("fs")
+        return fs.refreshMounts()
+
+
+class _SecretsUtil:
+    """Remote equivalent of secrets util"""
+
+    def __init__(self, secrets_api: workspace.SecretsAPI):
+        self._api = secrets_api  # nolint
+
+    def getBytes(self, scope: str, key: str) -> bytes:
+        """Gets the bytes representation of a secret value for the specified scope and key."""
+        query = {"scope": scope, "key": key}
+        raw = self._api._api.do("GET", "/api/2.0/secrets/get", query=query)
+        return base64.b64decode(raw["value"])
+
+    def get(self, scope: str, key: str) -> str:
+        """Gets the string representation of a secret value for the specified secrets scope and key."""
+        val = self.getBytes(scope, key)
+        string_value = val.decode()
+        return string_value
+
+    def list(self, scope) -> List[SecretMetadata]:
+        """Lists the metadata for secrets within the specified scope."""
+
+        # transform from SDK dataclass to dbutils-compatible namedtuple
+        return [SecretMetadata(v.key) for v in self._api.list_secrets(scope)]
+
+    def listScopes(self) -> List[SecretScope]:
+        """Lists the available scopes."""
+
+        # transform from SDK dataclass to dbutils-compatible namedtuple
+        return [SecretScope(v.name) for v in self._api.list_scopes()]
+
+
+class _JobsUtil:
+    """Remote equivalent of jobs util"""
+
+    class _TaskValuesUtil:
+        """Remote equivalent of task values util"""
+
+        def get(
+            self,
+            taskKey: str,
+            key: str,
+            default: any = None,
+            debugValue: any = None,
+        ) -> None:
+            """
+            Returns `debugValue` if present, throws an error otherwise as this implementation is always run outside of a job run
+            """
+            if debugValue is None:
+                raise TypeError(
+                    "Must pass debugValue when calling get outside of a job context. debugValue cannot be None."
+                )
+            return debugValue
+
+        def set(self, key: str, value: any) -> None:
+            """
+            Sets a task value on the current task run
+            """
+
+    def __init__(self) -> None:
+        self.taskValues = self._TaskValuesUtil()
+
+
+class RemoteDbUtils:
+
+    def __init__(self, config: "Config" = None):
+        # Create a shallow copy of the config to allow the use of a custom
+        # user-agent while avoiding modifying the original config.
+        self._config = Config() if not config else config.copy()
+        self._config.with_user_agent_extra("dbutils", "remote")
+
+        self._client = ApiClient(self._config)
+        self._clusters = compute_ext.ClustersExt(self._client)
+        self._commands = compute.CommandExecutionAPI(self._client)
+        self._lock = threading.Lock()
+        self._ctx = None
+
+        self.fs = _FsUtil(dbfs_ext.DbfsExt(self._client), self.__getattr__)
+        self.secrets = _SecretsUtil(workspace.SecretsAPI(self._client))
+        self.jobs = _JobsUtil()
+        self._widgets = None
+
+    # When we import widget_impl, the init file checks whether user has the
+    # correct dependencies required for running on notebook or not (ipywidgets etc).
+    # We only want these checks (and the subsequent errors and warnings), to
+    # happen when the user actually uses widgets.
+    @property
+    def widgets(self):
+        if self._widgets is None:
+            from ._widgets import widget_impl
+
+            self._widgets = widget_impl()
+
+        return self._widgets
+
+    @property
+    def _cluster_id(self) -> str:
+        cluster_id = self._config.cluster_id
+        if not cluster_id:
+            message = "cluster_id is required in the configuration"
+            raise ValueError(self._config.wrap_debug_info(message))
+        return cluster_id
+
+    def _running_command_context(self) -> compute.ContextStatusResponse:
+        if self._ctx:
+            return self._ctx
+        with self._lock:
+            if self._ctx:
+                return self._ctx
+            self._clusters.ensure_cluster_is_running(self._cluster_id)
+            self._ctx = self._commands.create(cluster_id=self._cluster_id, language=compute.Language.PYTHON).result()
+        return self._ctx
+
+    def __getattr__(self, util) -> "_ProxyUtil":
+        return _ProxyUtil(
+            command_execution=self._commands,
+            context_factory=self._running_command_context,
+            cluster_id=self._cluster_id,
+            name=util,
+        )
+
+
+@dataclass
+class OverrideResult:
+    result: Any
+
+
+def get_local_notebook_path():
+    value = os.getenv("DATABRICKS_SOURCE_FILE")
+    if value is None:
+        raise ValueError(
+            "Getting the current notebook path is only supported when running a notebook using the `Databricks Connect: Run as File` or `Databricks Connect: Debug as File` commands in the Databricks extension for VS Code. To bypass this error, set environment variable `DATABRICKS_SOURCE_FILE` to the desired notebook path."
+        )
+
+    return value
+
+
+def not_supported_method_err_msg(methodName):
+    return f"Method '{methodName}' is not supported in the SDK version of DBUtils"
+
+
+class _OverrideProxyUtil:
+
+    @classmethod
+    def new(cls, path: str):
+        if path in cls.not_supported_override_paths:
+            raise ValueError(cls.not_supported_override_paths[path])
+
+        if len(cls.__get_matching_overrides(path)) > 0:
+            return _OverrideProxyUtil(path)
+        return None
+
+    def __init__(self, name: str):
+        self._name = name
+
+    # These are the paths that we want to override and not send to remote dbutils. NOTE, for each of these paths, no prefixes
+    # are sent to remote either. This could lead to unintentional breakage.
+    # Our current proxy implementation (which sends everything to remote dbutils) uses `{util}.{method}(*args, **kwargs)` ONLY.
+    # This means, it is completely safe to override paths starting with `{util}.{attribute}.<other_parts>`, since none of the prefixes
+    # are being proxied to remote dbutils currently.
+    proxy_override_paths = {
+        "notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()": get_local_notebook_path,
+    }
+
+    # These paths work the same as 'proxy_override_paths' but instead of using a local implementation we raise an exception.
+    not_supported_override_paths = {
+        # The object returned by 'credentials.getServiceCredentialProvider()' can't be serialized to JSON.
+        # Without this override, the command would fail with an error 'TypeError: Object of type Session is not JSON serializable'.
+        # We override it to show a better error message
+        "credentials.getServiceCredentialsProvider": not_supported_method_err_msg(
+            "credentials.getServiceCredentialsProvider"
+        ),
+    }
+
+    @classmethod
+    def __get_matching_overrides(cls, path: str):
+        return [x for x in cls.proxy_override_paths.keys() if x.startswith(path)]
+
+    def __run_override(self, path: str) -> Optional[OverrideResult]:
+        overrides = self.__get_matching_overrides(path)
+        if len(overrides) == 1 and overrides[0] == path:
+            return OverrideResult(self.proxy_override_paths[overrides[0]]())
+
+        if len(overrides) > 0:
+            return OverrideResult(_OverrideProxyUtil(name=path))
+
+        return None
+
+    def __call__(self, *args, **kwds) -> Any:
+        if len(args) != 0 or len(kwds) != 0:
+            raise TypeError(
+                f"Arguments are not supported for overridden method {self._name}. Invoke as: {self._name}()"
+            )
+
+        callable_path = f"{self._name}()"
+        result = self.__run_override(callable_path)
+        if result:
+            return result.result
+
+        raise TypeError(f"{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> Any:
+        result = self.__run_override(f"{self._name}.{method}")
+        if result:
+            return result.result
+
+        raise AttributeError(f"module {self._name} has no attribute {method}")
+
+
+class _ProxyUtil:
+    """Enables temporary workaround to call remote in-REPL dbutils without having to re-implement them"""
+
+    def __init__(
+        self,
+        *,
+        command_execution: compute.CommandExecutionAPI,
+        context_factory: Callable[[], compute.ContextStatusResponse],
+        cluster_id: str,
+        name: str,
+    ):
+        self._commands = command_execution
+        self._cluster_id = cluster_id
+        self._context_factory = context_factory
+        self._name = name
+
+    def __call__(self):
+        raise NotImplementedError(f"dbutils.{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> "_ProxyCall | _ProxyUtil | _OverrideProxyUtil":
+        override = _OverrideProxyUtil.new(f"{self._name}.{method}")
+        if override:
+            return override
+
+        return _ProxyCall(
+            command_execution=self._commands,
+            cluster_id=self._cluster_id,
+            context_factory=self._context_factory,
+            util=self._name,
+            method=method,
+        )
+
+
+import html
+import re
+
+
+class _ProxyCall:
+
+    def __init__(
+        self,
+        *,
+        command_execution: compute.CommandExecutionAPI,
+        context_factory: Callable[[], compute.ContextStatusResponse],
+        cluster_id: str,
+        util: str,
+        method: str,
+    ):
+        self._commands = command_execution
+        self._cluster_id = cluster_id
+        self._context_factory = context_factory
+        self._util = util
+        self._method = method
+
+    _out_re = re.compile(r"Out\[[\d\s]+]:\s")
+    _tag_re = re.compile(r"<[^>]*>")
+    _exception_re = re.compile(r".*Exception:\s+(.*)")
+    _execution_error_re = re.compile(r"ExecutionError: ([\s\S]*)\n(StatusCode=[0-9]*)\n(StatusDescription=.*)\n")
+    _error_message_re = re.compile(r"ErrorMessage=(.+)\n")
+    _ascii_escape_re = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]")
+
+    def _is_failed(self, results: compute.Results) -> bool:
+        return results.result_type == compute.ResultType.ERROR
+
+    def _text(self, results: compute.Results) -> str:
+        if results.result_type != compute.ResultType.TEXT:
+            return ""
+        return self._out_re.sub("", str(results.data))
+
+    def _raise_if_failed(self, results: compute.Results):
+        if not self._is_failed(results):
+            return
+        raise DatabricksError(self._error_from_results(results))
+
+    def _error_from_results(self, results: compute.Results):
+        if not self._is_failed(results):
+            return
+        if results.cause:
+            _LOG.debug(f'{self._ascii_escape_re.sub("", results.cause)}')
+
+        summary = self._tag_re.sub("", results.summary)
+        summary = html.unescape(summary)
+
+        exception_matches = self._exception_re.findall(summary)
+        if len(exception_matches) == 1:
+            summary = exception_matches[0].replace("; nested exception is:", "")
+            summary = summary.rstrip(" ")
+            return summary
+
+        execution_error_matches = self._execution_error_re.findall(results.cause)
+        if len(execution_error_matches) == 1:
+            return "\n".join(execution_error_matches[0])
+
+        error_message_matches = self._error_message_re.findall(results.cause)
+        if len(error_message_matches) == 1:
+            return error_message_matches[0]
+
+        return summary
+
+    def __call__(self, *args, **kwargs):
+        raw = json.dumps((args, kwargs))
+        code = f"""
+        import json
+        (args, kwargs) = json.loads('{raw}')
+        result = dbutils.{self._util}.{self._method}(*args, **kwargs)
+        dbutils.notebook.exit(json.dumps(result))
+        """
+        ctx = self._context_factory()
+        result = self._commands.execute(
+            cluster_id=self._cluster_id,
+            language=compute.Language.PYTHON,
+            context_id=ctx.id,
+            command=code,
+        ).result()
+        if result.status == compute.CommandStatus.FINISHED:
+            self._raise_if_failed(result.results)
+            raw = result.results.data
+            return json.loads(raw)
+        else:
+            raise Exception(result.results.summary)

venv/lib/python3.9/site-packages/databricks/sdk/environments.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+
+@dataclass
+class AzureEnvironment:
+    name: str
+    service_management_endpoint: str
+    resource_manager_endpoint: str
+    active_directory_endpoint: str
+
+
+ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+
+ENVIRONMENTS = dict(
+    PUBLIC=AzureEnvironment(
+        name="PUBLIC",
+        service_management_endpoint="https://management.core.windows.net/",
+        resource_manager_endpoint="https://management.azure.com/",
+        active_directory_endpoint="https://login.microsoftonline.com/",
+    ),
+    USGOVERNMENT=AzureEnvironment(
+        name="USGOVERNMENT",
+        service_management_endpoint="https://management.core.usgovcloudapi.net/",
+        resource_manager_endpoint="https://management.usgovcloudapi.net/",
+        active_directory_endpoint="https://login.microsoftonline.us/",
+    ),
+    CHINA=AzureEnvironment(
+        name="CHINA",
+        service_management_endpoint="https://management.core.chinacloudapi.cn/",
+        resource_manager_endpoint="https://management.chinacloudapi.cn/",
+        active_directory_endpoint="https://login.chinacloudapi.cn/",
+    ),
+)
+
+
+class Cloud(Enum):
+    AWS = "AWS"
+    AZURE = "AZURE"
+    GCP = "GCP"
+
+
+@dataclass
+class DatabricksEnvironment:
+    cloud: Cloud
+    dns_zone: str
+    azure_application_id: Optional[str] = None
+    azure_environment: Optional[AzureEnvironment] = None
+
+    def deployment_url(self, name: str) -> str:
+        return f"https://{name}{self.dns_zone}"
+
+    @property
+    def azure_service_management_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.service_management_endpoint
+
+    @property
+    def azure_resource_manager_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.resource_manager_endpoint
+
+    @property
+    def azure_active_directory_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.active_directory_endpoint
+
+
+DEFAULT_ENVIRONMENT = DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.com")
+
+ALL_ENVS = [
+    DatabricksEnvironment(Cloud.AWS, ".dev.databricks.com"),
+    DatabricksEnvironment(Cloud.AWS, ".staging.cloud.databricks.com"),
+    DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.us"),
+    DEFAULT_ENVIRONMENT,
+    DatabricksEnvironment(
+        Cloud.AZURE,
+        ".dev.azuredatabricks.net",
+        azure_application_id="62a912ac-b58e-4c1d-89ea-b2dbfc7358fc",
+        azure_environment=ENVIRONMENTS["PUBLIC"],
+    ),
+    DatabricksEnvironment(
+        Cloud.AZURE,
+        ".staging.azuredatabricks.net",
+        azure_application_id="4a67d088-db5c-48f1-9ff2-0aace800ae68",
+        azure_environment=ENVIRONMENTS["PUBLIC"],
+    ),
+    DatabricksEnvironment(
+        Cloud.AZURE,
+        ".azuredatabricks.net",
+        azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+        azure_environment=ENVIRONMENTS["PUBLIC"],
+    ),
+    DatabricksEnvironment(
+        Cloud.AZURE,
+        ".databricks.azure.us",
+        azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+        azure_environment=ENVIRONMENTS["USGOVERNMENT"],
+    ),
+    DatabricksEnvironment(
+        Cloud.AZURE,
+        ".databricks.azure.cn",
+        azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+        azure_environment=ENVIRONMENTS["CHINA"],
+    ),
+    DatabricksEnvironment(Cloud.GCP, ".dev.gcp.databricks.com"),
+    DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
+    DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com"),
+]
+
+
+def get_environment_for_hostname(hostname: Optional[str]) -> DatabricksEnvironment:
+    if not hostname:
+        return DEFAULT_ENVIRONMENT
+    for env in ALL_ENVS:
+        if hostname.endswith(env.dns_zone):
+            return env
+    return DEFAULT_ENVIRONMENT

venv/lib/python3.9/site-packages/databricks/sdk/errors/__init__.py

@@ -0,0 +1,6 @@
+from .base import DatabricksError, ErrorDetail
+from .customizer import _ErrorCustomizer
+from .parser import _Parser
+from .platform import *
+from .private_link import PrivateLinkValidationError
+from .sdk import *

venv/lib/python3.9/site-packages/databricks/sdk/errors/base.py

@@ -0,0 +1,136 @@
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from . import details as errdetails
+
+
+# Deprecated.
+class ErrorDetail:
+    def __init__(
+        self,
+        type: Optional[str] = None,
+        reason: Optional[str] = None,
+        domain: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        **kwargs,
+    ):
+        self.type = type
+        self.reason = reason
+        self.domain = domain
+        self.metadata = metadata
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "ErrorDetail":
+        # Key "@type" is not a valid keyword argument name in Python. Rename
+        # it to "type" to avoid conflicts.
+        safe_args = {}
+        for k, v in d.items():
+            safe_args[k if k != "@type" else "type"] = v
+
+        return cls(**safe_args)
+
+
+class DatabricksError(IOError):
+    """Generic error from Databricks REST API"""
+
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        *,
+        error_code: Optional[str] = None,
+        detail: Optional[str] = None,
+        status: Optional[str] = None,
+        scimType: Optional[str] = None,
+        error: Optional[str] = None,
+        retry_after_secs: Optional[int] = None,
+        details: Optional[List[Dict[str, Any]]] = None,
+        **kwargs,
+    ):
+        """
+
+        :param message:
+        :param error_code:
+        :param detail: [Deprecated]
+        :param status: [Deprecated]
+        :param scimType: [Deprecated]
+        :param error: [Deprecated]
+        :param retry_after_secs: [Deprecated]
+        :param details:
+        :param kwargs:
+        """
+
+        if detail:
+            # Handle SCIM error message details
+            # @see https://tools.ietf.org/html/rfc7644#section-3.7.3
+            if detail == "null":
+                message = "SCIM API Internal Error"
+            else:
+                message = detail
+            # add more context from SCIM responses
+            message = f"{scimType} {message}".strip(" ")
+            error_code = f"SCIM_{status}"
+
+        super().__init__(message if message else error)
+        self.error_code = error_code
+        self.retry_after_secs = retry_after_secs
+        self._error_details = errdetails.parse_error_details(details or [])
+        self.kwargs = kwargs
+
+        # Deprecated.
+        self.details = []
+        if details:
+            for d in details:
+                if not isinstance(d, dict):
+                    continue
+                self.details.append(ErrorDetail.from_dict(d))
+
+    def get_error_info(self) -> List[ErrorDetail]:
+        if self.details is None:
+            return []
+        return [detail for detail in self.details if detail.type == errdetails._ERROR_INFO_TYPE]
+
+    def get_error_details(self) -> errdetails.ErrorDetails:
+        return self._error_details
+
+
+@dataclass
+class _ErrorOverride:
+    # The name of the override. Used for logging purposes.
+    debug_name: str
+
+    # A regex that must match the path of the request for this override to be applied.
+    path_regex: re.Pattern
+
+    # The HTTP method of the request for the override to apply
+    verb: str
+
+    # The custom error class to use for this override.
+    custom_error: type
+
+    # A regular expression that must match the error code for this override to be applied. If None,
+    # this field is ignored.
+    status_code_matcher: Optional[re.Pattern] = None
+
+    # A regular expression that must match the error code for this override to be applied. If None,
+    # this field is ignored.
+    error_code_matcher: Optional[re.Pattern] = None
+
+    # A regular expression that must match the message for this override to be applied. If None,
+    # this field is ignored.
+    message_matcher: Optional[re.Pattern] = None
+
+    def matches(self, response: requests.Response, raw_error: dict):
+        if response.request.method != self.verb:
+            return False
+        if not self.path_regex.match(response.request.path_url):
+            return False
+        if self.status_code_matcher and not self.status_code_matcher.match(str(response.status_code)):
+            return False
+        if self.error_code_matcher and not self.error_code_matcher.match(raw_error.get("error_code", "")):
+            return False
+        if self.message_matcher and not self.message_matcher.match(raw_error.get("message", "")):
+            return False
+        return True

venv/lib/python3.9/site-packages/databricks/sdk/errors/customizer.py

@@ -0,0 +1,50 @@
+import abc
+import logging
+
+import requests
+
+
+class _ErrorCustomizer(abc.ABC):
+    """A customizer for errors from the Databricks REST API."""
+
+    @abc.abstractmethod
+    def customize_error(self, response: requests.Response, kwargs: dict):
+        """Customize the error constructor parameters."""
+
+
+class _RetryAfterCustomizer(_ErrorCustomizer):
+    """An error customizer that sets the retry_after_secs parameter based on the Retry-After header."""
+
+    _DEFAULT_RETRY_AFTER_SECONDS = 1
+    """The default number of seconds to wait before retrying a request if the Retry-After header is missing or is not
+    a valid integer."""
+
+    @classmethod
+    def _parse_retry_after(cls, response: requests.Response) -> int:
+        retry_after = response.headers.get("Retry-After")
+        if retry_after is None:
+            logging.debug(
+                f"No Retry-After header received in response with status code 429 or 503. Defaulting to {cls._DEFAULT_RETRY_AFTER_SECONDS}"
+            )
+            # 429 requests should include a `Retry-After` header, but if it's missing,
+            # we default to 1 second.
+            return cls._DEFAULT_RETRY_AFTER_SECONDS
+        # If the request is throttled, try parse the `Retry-After` header and sleep
+        # for the specified number of seconds. Note that this header can contain either
+        # an integer or a RFC1123 datetime string.
+        # See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
+        #
+        # For simplicity, we only try to parse it as an integer, as this is what Databricks
+        # platform returns. Otherwise, we fall back and don't sleep.
+        try:
+            return int(retry_after)
+        except ValueError:
+            logging.debug(
+                f"Invalid Retry-After header received: {retry_after}. Defaulting to {cls._DEFAULT_RETRY_AFTER_SECONDS}"
+            )
+            # defaulting to 1 sleep second to make self._is_retryable() simpler
+            return cls._DEFAULT_RETRY_AFTER_SECONDS
+
+    def customize_error(self, response: requests.Response, kwargs: dict):
+        if response.status_code in (429, 503):
+            kwargs["retry_after_secs"] = self._parse_retry_after(response)

venv/lib/python3.9/site-packages/databricks/sdk/errors/deserializer.py

@@ -0,0 +1,119 @@
+import abc
+import json
+import logging
+import re
+from typing import Optional
+
+import requests
+
+
+class _ErrorDeserializer(abc.ABC):
+    """A parser for errors from the Databricks REST API."""
+
+    @abc.abstractmethod
+    def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        """Parses an error from the Databricks REST API. If the error cannot be parsed, returns None."""
+
+
+class _EmptyDeserializer(_ErrorDeserializer):
+    """A parser that handles empty responses."""
+
+    def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        if len(response_body) == 0:
+            return {"message": response.reason}
+        return None
+
+
+class _StandardErrorDeserializer(_ErrorDeserializer):
+    """
+    Parses errors from the Databricks REST API using the standard error format.
+    """
+
+    def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        try:
+            payload_str = response_body.decode("utf-8")
+            resp = json.loads(payload_str)
+        except UnicodeDecodeError as e:
+            logging.debug(
+                "_StandardErrorParser: unable to decode response using utf-8",
+                exc_info=e,
+            )
+            return None
+        except json.JSONDecodeError as e:
+            logging.debug(
+                "_StandardErrorParser: unable to deserialize response as json",
+                exc_info=e,
+            )
+            return None
+        if not isinstance(resp, dict):
+            logging.debug("_StandardErrorParser: response is valid JSON but not a dictionary")
+            return None
+
+        error_args = {
+            "message": resp.get("message", "request failed"),
+            "error_code": resp.get("error_code"),
+            "details": resp.get("details"),
+        }
+
+        # Handle API 1.2-style errors
+        if "error" in resp:
+            error_args["message"] = resp["error"]
+
+        # Handle SCIM Errors
+        detail = resp.get("detail")
+        status = resp.get("status")
+        scim_type = resp.get("scimType")
+        if detail:
+            # Handle SCIM error message details
+            # @see https://tools.ietf.org/html/rfc7644#section-3.7.3
+            if detail == "null":
+                detail = "SCIM API Internal Error"
+            error_args["message"] = f"{scim_type} {detail}".strip(" ")
+            error_args["error_code"] = f"SCIM_{status}"
+        return error_args
+
+
+class _StringErrorDeserializer(_ErrorDeserializer):
+    """
+    Parses errors from the Databricks REST API in the format "ERROR_CODE: MESSAGE".
+    """
+
+    __STRING_ERROR_REGEX = re.compile(r"([A-Z_]+): (.*)")
+
+    def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        payload_str = response_body.decode("utf-8")
+        match = self.__STRING_ERROR_REGEX.match(payload_str)
+        if not match:
+            logging.debug("_StringErrorParser: unable to parse response as string")
+            return None
+        error_code, message = match.groups()
+        return {
+            "error_code": error_code,
+            "message": message,
+            "status": response.status_code,
+        }
+
+
+class _HtmlErrorDeserializer(_ErrorDeserializer):
+    """
+    Parses errors from the Databricks REST API in HTML format.
+    """
+
+    __HTML_ERROR_REGEXES = [
+        re.compile(r"<pre>(.*)</pre>"),
+        re.compile(r"<title>(.*)</title>"),
+    ]
+
+    def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        payload_str = response_body.decode("utf-8")
+        for regex in self.__HTML_ERROR_REGEXES:
+            match = regex.search(payload_str)
+            if match:
+                message = match.group(1) if match.group(1) else response.reason
+                return {
+                    "status": response.status_code,
+                    "message": message,
+                    "error_code": response.reason.upper().replace(" ", "_"),
+                }
+        logging.debug("_HtmlErrorParser: no <pre> tag found in error response")
+        return None

venv/lib/python3.9/site-packages/databricks/sdk/errors/details.py

@@ -0,0 +1,417 @@
+import re
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class ErrorInfo:
+    """Describes the cause of the error with structured details."""
+
+    reason: str
+    domain: str
+    metadata: Dict[str, str]
+
+
+@dataclass
+class RequestInfo:
+    """
+    Contains metadata about the request that clients can attach when
+    filing a bug or providing other forms of feedback.
+    """
+
+    request_id: str
+    serving_data: str
+
+
+@dataclass
+class RetryInfo:
+    """
+    Describes when the clients can retry a failed request. Clients could
+    ignore the recommendation here or retry when this information is missing
+    from error responses.
+
+    It's always recommended that clients should use exponential backoff
+    when retrying.
+
+    Clients should wait until `retry_delay` amount of time has passed since
+    receiving the error response before retrying.  If retrying requests also
+    fail, clients should use an exponential backoff scheme to gradually
+    increase the delay between retries based on `retry_delay`, until either
+    a maximum number of retries have been reached or a maximum retry delay
+    cap has been reached.
+    """
+
+    retry_delay_seconds: float
+
+
+@dataclass
+class DebugInfo:
+    """Describes additional debugging info."""
+
+    stack_entries: List[str]
+    detail: str
+
+
+@dataclass
+class QuotaFailureViolation:
+    """Describes a single quota violation."""
+
+    subject: str
+    description: str
+
+
+@dataclass
+class QuotaFailure:
+    """
+    Describes how a quota check failed.
+
+    For example if a daily limit was exceeded for the calling project, a
+    service could respond with a QuotaFailure detail containing the project
+    id and the description of the quota limit that was exceeded.  If the
+    calling project hasn't enabled the service in the developer console,
+    then a service could respond with the project id and set
+    `service_disabled` to true.
+
+    Also see RetryInfo and Help types for other details about handling a
+    quota failure.
+    """
+
+    violations: List[QuotaFailureViolation]
+
+
+@dataclass
+class PreconditionFailureViolation:
+    """Describes a single precondition violation."""
+
+    type: str
+    subject: str
+    description: str
+
+
+@dataclass
+class PreconditionFailure:
+    """Describes what preconditions have failed."""
+
+    violations: List[PreconditionFailureViolation]
+
+
+@dataclass
+class BadRequestFieldViolation:
+    """Describes a single field violation in a bad request."""
+
+    field: str
+    description: str
+
+
+@dataclass
+class BadRequest:
+    """
+    Describes violations in a client request. This error type
+    focuses on the syntactic aspects of the request.
+    """
+
+    field_violations: List[BadRequestFieldViolation]
+
+
+@dataclass
+class ResourceInfo:
+    """Describes the resource that is being accessed."""
+
+    resource_type: str
+    resource_name: str
+    owner: str
+    description: str
+
+
+@dataclass
+class HelpLink:
+    """Describes a single help link."""
+
+    description: str
+    url: str
+
+
+@dataclass
+class Help:
+    """
+    Provides links to documentation or for performing an out of
+    band action.
+
+    For example, if a quota check failed with an error indicating
+    the calling project hasn't enabled the accessed service, this
+    can contain a URL pointing directly to the right place in the
+    developer console to flip the bit.
+    """
+
+    links: List[HelpLink]
+
+
+@dataclass
+class ErrorDetails:
+    """
+    ErrorDetails contains the error details of an API error. It
+    is the union of known error details types and unknown details.
+    """
+
+    error_info: Optional[ErrorInfo] = None
+    request_info: Optional[RequestInfo] = None
+    retry_info: Optional[RetryInfo] = None
+    debug_info: Optional[DebugInfo] = None
+    quota_failure: Optional[QuotaFailure] = None
+    precondition_failure: Optional[PreconditionFailure] = None
+    bad_request: Optional[BadRequest] = None
+    resource_info: Optional[ResourceInfo] = None
+    help: Optional[Help] = None
+    unknown_details: List[Any] = field(default_factory=list)
+
+
+# Supported error details proto types.
+_ERROR_INFO_TYPE = "type.googleapis.com/google.rpc.ErrorInfo"
+_REQUEST_INFO_TYPE = "type.googleapis.com/google.rpc.RequestInfo"
+_RETRY_INFO_TYPE = "type.googleapis.com/google.rpc.RetryInfo"
+_DEBUG_INFO_TYPE = "type.googleapis.com/google.rpc.DebugInfo"
+_QUOTA_FAILURE_TYPE = "type.googleapis.com/google.rpc.QuotaFailure"
+_PRECONDITION_FAILURE_TYPE = "type.googleapis.com/google.rpc.PreconditionFailure"
+_BAD_REQUEST_TYPE = "type.googleapis.com/google.rpc.BadRequest"
+_RESOURCE_INFO_TYPE = "type.googleapis.com/google.rpc.ResourceInfo"
+_HELP_TYPE = "type.googleapis.com/google.rpc.Help"
+
+
+def parse_error_details(details: List[Any]) -> ErrorDetails:
+    ed = ErrorDetails()
+
+    if not details:
+        return ed
+
+    for d in details:
+        pd = _parse_json_error_details(d)
+
+        if isinstance(pd, ErrorInfo):
+            ed.error_info = pd
+        elif isinstance(pd, RequestInfo):
+            ed.request_info = pd
+        elif isinstance(pd, RetryInfo):
+            ed.retry_info = pd
+        elif isinstance(pd, DebugInfo):
+            ed.debug_info = pd
+        elif isinstance(pd, QuotaFailure):
+            ed.quota_failure = pd
+        elif isinstance(pd, PreconditionFailure):
+            ed.precondition_failure = pd
+        elif isinstance(pd, BadRequest):
+            ed.bad_request = pd
+        elif isinstance(pd, ResourceInfo):
+            ed.resource_info = pd
+        elif isinstance(pd, Help):
+            ed.help = pd
+        else:
+            ed.unknown_details.append(pd)
+
+    return ed
+
+
+def _parse_json_error_details(value: Any) -> Any:
+    """
+    Attempts to parse an error details type from the given JSON value. If the
+    value is not a known error details type, it returns the input as is.
+
+    :param value: The JSON value to parse.
+    :return: The parsed error details type or the input value if it is not
+        a known error details type.
+    """
+
+    if not isinstance(value, dict):
+        return value  # not a JSON object
+
+    t = value.get("@type")
+    if not isinstance(t, str):
+        return value  # JSON object with no @type field
+
+    try:
+        if t == _ERROR_INFO_TYPE:
+            return _parse_error_info(value)
+        elif t == _REQUEST_INFO_TYPE:
+            return _parse_req_info(value)
+        elif t == _RETRY_INFO_TYPE:
+            return _parse_retry_info(value)
+        elif t == _DEBUG_INFO_TYPE:
+            return _parse_debug_info(value)
+        elif t == _QUOTA_FAILURE_TYPE:
+            return _parse_quota_failure(value)
+        elif t == _PRECONDITION_FAILURE_TYPE:
+            return _parse_precondition_failure(value)
+        elif t == _BAD_REQUEST_TYPE:
+            return _parse_bad_request(value)
+        elif t == _RESOURCE_INFO_TYPE:
+            return _parse_resource_info(value)
+        elif t == _HELP_TYPE:
+            return _parse_help(value)
+        else:  # unknown type
+            return value
+    except (TypeError, ValueError):
+        return value  # not a valid known type
+    except Exception:
+        return value
+
+
+def _parse_error_info(d: Dict[str, Any]) -> ErrorInfo:
+    return ErrorInfo(
+        domain=_parse_string(d.get("domain", "")),
+        reason=_parse_string(d.get("reason", "")),
+        metadata=_parse_dict(d.get("metadata", {})),
+    )
+
+
+def _parse_req_info(d: Dict[str, Any]) -> RequestInfo:
+    return RequestInfo(
+        request_id=_parse_string(d.get("request_id", "")),
+        serving_data=_parse_string(d.get("serving_data", "")),
+    )
+
+
+def _parse_retry_info(d: Dict[str, Any]) -> RetryInfo:
+    delay = 0.0
+    if "retry_delay" in d:
+        delay = _parse_seconds(d["retry_delay"])
+
+    return RetryInfo(
+        retry_delay_seconds=delay,
+    )
+
+
+def _parse_debug_info(d: Dict[str, Any]) -> DebugInfo:
+    di = DebugInfo(
+        stack_entries=[],
+        detail=_parse_string(d.get("detail", "")),
+    )
+
+    if "stack_entries" not in d:
+        return di
+
+    if not isinstance(d["stack_entries"], list):
+        raise ValueError(f"Expected list, got {d['stack_entries']!r}")
+    for entry in d["stack_entries"]:
+        di.stack_entries.append(_parse_string(entry))
+
+    return di
+
+
+def _parse_quota_failure_violation(d: Dict[str, Any]) -> QuotaFailureViolation:
+    return QuotaFailureViolation(
+        subject=_parse_string(d.get("subject", "")),
+        description=_parse_string(d.get("description", "")),
+    )
+
+
+def _parse_quota_failure(d: Dict[str, Any]) -> QuotaFailure:
+    violations = []
+    if "violations" in d:
+        if not isinstance(d["violations"], list):
+            raise ValueError(f"Expected list, got {d['violations']!r}")
+        for violation in d["violations"]:
+            if not isinstance(violation, dict):
+                raise ValueError(f"Expected dict, got {violation!r}")
+            violations.append(_parse_quota_failure_violation(violation))
+    return QuotaFailure(violations=violations)
+
+
+def _parse_precondition_failure_violation(d: Dict[str, Any]) -> PreconditionFailureViolation:
+    return PreconditionFailureViolation(
+        type=_parse_string(d.get("type", "")),
+        subject=_parse_string(d.get("subject", "")),
+        description=_parse_string(d.get("description", "")),
+    )
+
+
+def _parse_precondition_failure(d: Dict[str, Any]) -> PreconditionFailure:
+    violations = []
+    if "violations" in d:
+        if not isinstance(d["violations"], list):
+            raise ValueError(f"Expected list, got {d['violations']!r}")
+        for v in d["violations"]:
+            if not isinstance(v, dict):
+                raise ValueError(f"Expected dict, got {v!r}")
+            violations.append(_parse_precondition_failure_violation(v))
+    return PreconditionFailure(violations=violations)
+
+
+def _parse_bad_request_field_violation(d: Dict[str, Any]) -> BadRequestFieldViolation:
+    return BadRequestFieldViolation(
+        field=_parse_string(d.get("field", "")),
+        description=_parse_string(d.get("description", "")),
+    )
+
+
+def _parse_bad_request(d: Dict[str, Any]) -> BadRequest:
+    field_violations = []
+    if "field_violations" in d:
+        if not isinstance(d["field_violations"], list):
+            raise ValueError(f"Expected list, got {d['field_violations']!r}")
+        for violation in d["field_violations"]:
+            if not isinstance(violation, dict):
+                raise ValueError(f"Expected dict, got {violation!r}")
+            field_violations.append(_parse_bad_request_field_violation(violation))
+    return BadRequest(field_violations=field_violations)
+
+
+def _parse_resource_info(d: Dict[str, Any]) -> ResourceInfo:
+    return ResourceInfo(
+        resource_type=_parse_string(d.get("resource_type", "")),
+        resource_name=_parse_string(d.get("resource_name", "")),
+        owner=_parse_string(d.get("owner", "")),
+        description=_parse_string(d.get("description", "")),
+    )
+
+
+def _parse_help_link(d: Dict[str, Any]) -> HelpLink:
+    return HelpLink(
+        description=_parse_string(d.get("description", "")),
+        url=_parse_string(d.get("url", "")),
+    )
+
+
+def _parse_help(d: Dict[str, Any]) -> Help:
+    links = []
+    if "links" in d:
+        if not isinstance(d["links"], list):
+            raise ValueError(f"Expected list, got {d['links']!r}")
+        for link in d["links"]:
+            if not isinstance(link, dict):
+                raise ValueError(f"Expected dict, got {link!r}")
+            links.append(_parse_help_link(link))
+    return Help(links=links)
+
+
+def _parse_string(a: Any) -> str:
+    if isinstance(a, str):
+        return a
+    raise ValueError(f"Expected string, got {a!r}")
+
+
+def _parse_dict(a: Any) -> Dict[str, str]:
+    if not isinstance(a, dict):
+        raise ValueError(f"Expected Dict[str, str], got {a!r}")
+    for key, value in a.items():
+        if not isinstance(key, str) or not isinstance(value, str):
+            raise ValueError(f"Expected Dict[str, str], got {a!r}")
+    return a
+
+
+def _parse_seconds(a: Any) -> float:
+    """
+    Parse a duration string into a float representing the number of seconds.
+
+    The duration type is encoded as a string rather than an where the string
+    ends in the suffix "s" (indicating seconds) and is preceded by a decimal
+    number of seconds. For example, "3.000000001s", represents a duration of
+    3 seconds and 1 nanosecond.
+    """
+
+    if not isinstance(a, str):
+        raise ValueError(f"Expected string, got {a!r}")
+
+    match = re.match(r"^(\d+(\.\d+)?)s$", a)
+    if match:
+        return float(match.group(1))
+
+    raise ValueError(f"Expected duration string, got {a!r}")

venv/lib/python3.9/site-packages/databricks/sdk/errors/mapper.py

@@ -0,0 +1,27 @@
+import requests
+
+from databricks.sdk.errors import platform
+from databricks.sdk.errors.base import DatabricksError
+
+from .overrides import _ALL_OVERRIDES
+
+
+def _error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
+    for override in _ALL_OVERRIDES:
+        if override.matches(response, raw):
+            return override.custom_error(**raw)
+    status_code = response.status_code
+    error_code = raw.get("error_code", None)
+    if error_code in platform.ERROR_CODE_MAPPING:
+        # more specific error codes override more generic HTTP status codes
+        return platform.ERROR_CODE_MAPPING[error_code](**raw)
+
+    if status_code in platform.STATUS_CODE_MAPPING:
+        # more generic HTTP status codes matched after more specific error codes,
+        # where there's a default exception class per HTTP status code, and we do
+        # rely on Databricks platform exception mapper to do the right thing.
+        return platform.STATUS_CODE_MAPPING[status_code](**raw)
+
+    # backwards-compatible error creation for cases like using older versions of
+    # the SDK on way never releases of the platform.
+    return DatabricksError(**raw)

venv/lib/python3.9/site-packages/databricks/sdk/errors/overrides.py

@@ -0,0 +1,36 @@
+# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
+
+import re
+
+from .base import _ErrorOverride
+from .platform import ResourceDoesNotExist
+
+_ALL_OVERRIDES = [
+    _ErrorOverride(
+        debug_name="Clusters InvalidParameterValue=>ResourceDoesNotExist",
+        path_regex=re.compile(r"^/api/2\.\d/clusters/get"),
+        verb="GET",
+        status_code_matcher=re.compile(r"^400$"),
+        error_code_matcher=re.compile(r"INVALID_PARAMETER_VALUE"),
+        message_matcher=re.compile(r"Cluster .* does not exist"),
+        custom_error=ResourceDoesNotExist,
+    ),
+    _ErrorOverride(
+        debug_name="Jobs InvalidParameterValue=>ResourceDoesNotExist",
+        path_regex=re.compile(r"^/api/2\.\d/jobs/get"),
+        verb="GET",
+        status_code_matcher=re.compile(r"^400$"),
+        error_code_matcher=re.compile(r"INVALID_PARAMETER_VALUE"),
+        message_matcher=re.compile(r"Job .* does not exist"),
+        custom_error=ResourceDoesNotExist,
+    ),
+    _ErrorOverride(
+        debug_name="Job Runs InvalidParameterValue=>ResourceDoesNotExist",
+        path_regex=re.compile(r"^/api/2\.\d/jobs/runs/get"),
+        verb="GET",
+        status_code_matcher=re.compile(r"^400$"),
+        error_code_matcher=re.compile(r"INVALID_PARAMETER_VALUE"),
+        message_matcher=re.compile(r"(Run .* does not exist|Run: .* in job: .* doesn\'t exist)"),
+        custom_error=ResourceDoesNotExist,
+    ),
+]

venv/lib/python3.9/site-packages/databricks/sdk/errors/parser.py

@@ -0,0 +1,100 @@
+import logging
+from typing import List, Optional
+
+import requests
+
+from ..logger import RoundTrip
+from .base import DatabricksError
+from .customizer import _ErrorCustomizer, _RetryAfterCustomizer
+from .deserializer import (_EmptyDeserializer, _ErrorDeserializer,
+                           _HtmlErrorDeserializer, _StandardErrorDeserializer,
+                           _StringErrorDeserializer)
+from .mapper import _error_mapper
+from .private_link import (_get_private_link_validation_error,
+                           _is_private_link_redirect)
+
+# A list of _ErrorDeserializers that are tried in order to parse an API error from a response body. Most errors should
+# be parsable by the _StandardErrorDeserializer, but additional parsers can be added here for specific error formats.
+# The order of the parsers is not important, as the set of errors that can be parsed by each parser should be disjoint.
+_error_deserializers = [
+    _EmptyDeserializer(),
+    _StandardErrorDeserializer(),
+    _StringErrorDeserializer(),
+    _HtmlErrorDeserializer(),
+]
+
+# A list of _ErrorCustomizers that are applied to the error arguments after they are parsed. Customizers can modify the
+# error arguments in any way, including adding or removing fields. Customizers are applied in order, so later
+# customizers can override the changes made by earlier customizers.
+_error_customizers = [
+    _RetryAfterCustomizer(),
+]
+
+
+def _unknown_error(response: requests.Response, debug_headers: bool = False) -> str:
+    """A standard error message that can be shown when an API response cannot be parsed.
+
+    This error message includes a link to the issue tracker for the SDK for users to report the issue to us.
+
+    :param response: The response object from the API request.
+    :param debug_headers: Whether to include headers in the request log. Defaults to False to defensively handle cases where request headers might contain sensitive data (e.g. tokens).
+    """
+    request_log = RoundTrip(response, debug_headers=debug_headers, debug_truncate_bytes=10 * 1024).generate()
+    return (
+        "This is likely a bug in the Databricks SDK for Python or the underlying "
+        "API. Please report this issue with the following debugging information to the SDK issue tracker at "
+        f"https://github.com/databricks/databricks-sdk-go/issues. Request log:```{request_log}```"
+    )
+
+
+class _Parser:
+    """
+    A parser for errors from the Databricks REST API. It attempts to deserialize an error using a sequence of
+    deserializers, and then customizes the deserialized error using a sequence of customizers. If the error cannot be
+    deserialized, it returns a generic error with debugging information and instructions to report the issue to the SDK
+    issue tracker.
+    """
+
+    def __init__(
+        self,
+        extra_error_parsers: List[_ErrorDeserializer] = [],
+        extra_error_customizers: List[_ErrorCustomizer] = [],
+        debug_headers: bool = False,
+    ):
+        self._error_parsers = _error_deserializers + (extra_error_parsers if extra_error_parsers is not None else [])
+        self._error_customizers = _error_customizers + (
+            extra_error_customizers if extra_error_customizers is not None else []
+        )
+        self._debug_headers = debug_headers
+
+    def get_api_error(self, response: requests.Response) -> Optional[DatabricksError]:
+        """
+        Handles responses from the REST API and returns a DatabricksError if the response indicates an error.
+        :param response: The response from the REST API.
+        :return: A DatabricksError if the response indicates an error, otherwise None.
+        """
+        if not response.ok:
+            content = response.content
+            for parser in self._error_parsers:
+                try:
+                    error_args = parser.deserialize_error(response, content)
+                    if error_args:
+                        for customizer in self._error_customizers:
+                            customizer.customize_error(response, error_args)
+                        return _error_mapper(response, error_args)
+                except Exception as e:
+                    logging.debug(
+                        f"Error parsing response with {parser}, continuing",
+                        exc_info=e,
+                    )
+            return _error_mapper(
+                response,
+                {"message": "unable to parse response. " + _unknown_error(response, self._debug_headers)},
+            )
+
+        # Private link failures happen via a redirect to the login page. From a requests-perspective, the request
+        # is successful, but the response is not what we expect. We need to handle this case separately.
+        if _is_private_link_redirect(response):
+            return _get_private_link_validation_error(response.url)
+
+        return None

venv/lib/python3.9/site-packages/databricks/sdk/errors/platform.py

@@ -0,0 +1,116 @@
+# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
+
+from .base import DatabricksError
+
+
+class BadRequest(DatabricksError):
+    """the request is invalid"""
+
+
+class Unauthenticated(DatabricksError):
+    """the request does not have valid authentication (AuthN) credentials for the operation"""
+
+
+class PermissionDenied(DatabricksError):
+    """the caller does not have permission to execute the specified operation"""
+
+
+class NotFound(DatabricksError):
+    """the operation was performed on a resource that does not exist"""
+
+
+class ResourceConflict(DatabricksError):
+    """maps to all HTTP 409 (Conflict) responses"""
+
+
+class TooManyRequests(DatabricksError):
+    """maps to HTTP code: 429 Too Many Requests"""
+
+
+class Cancelled(DatabricksError):
+    """the operation was explicitly canceled by the caller"""
+
+
+class InternalError(DatabricksError):
+    """some invariants expected by the underlying system have been broken"""
+
+
+class NotImplemented(DatabricksError):
+    """the operation is not implemented or is not supported/enabled in this service"""
+
+
+class TemporarilyUnavailable(DatabricksError):
+    """the service is currently unavailable"""
+
+
+class DeadlineExceeded(DatabricksError):
+    """the deadline expired before the operation could complete"""
+
+
+class InvalidState(BadRequest):
+    """unexpected state"""
+
+
+class InvalidParameterValue(BadRequest):
+    """supplied value for a parameter was invalid"""
+
+
+class ResourceDoesNotExist(NotFound):
+    """operation was performed on a resource that does not exist"""
+
+
+class Aborted(ResourceConflict):
+    """the operation was aborted, typically due to a concurrency issue such as a sequencer check
+    failure"""
+
+
+class AlreadyExists(ResourceConflict):
+    """operation was rejected due a conflict with an existing resource"""
+
+
+class ResourceAlreadyExists(ResourceConflict):
+    """operation was rejected due a conflict with an existing resource"""
+
+
+class ResourceExhausted(TooManyRequests):
+    """operation is rejected due to per-user rate limiting"""
+
+
+class RequestLimitExceeded(TooManyRequests):
+    """cluster request was rejected because it would exceed a resource limit"""
+
+
+class Unknown(InternalError):
+    """this error is used as a fallback if the platform-side mapping is missing some reason"""
+
+
+class DataLoss(InternalError):
+    """unrecoverable data loss or corruption"""
+
+
+STATUS_CODE_MAPPING = {
+    400: BadRequest,
+    401: Unauthenticated,
+    403: PermissionDenied,
+    404: NotFound,
+    409: ResourceConflict,
+    429: TooManyRequests,
+    499: Cancelled,
+    500: InternalError,
+    501: NotImplemented,
+    503: TemporarilyUnavailable,
+    504: DeadlineExceeded,
+}
+
+ERROR_CODE_MAPPING = {
+    "INVALID_STATE": InvalidState,
+    "INVALID_PARAMETER_VALUE": InvalidParameterValue,
+    "RESOURCE_DOES_NOT_EXIST": ResourceDoesNotExist,
+    "ABORTED": Aborted,
+    "ALREADY_EXISTS": AlreadyExists,
+    "RESOURCE_ALREADY_EXISTS": ResourceAlreadyExists,
+    "RESOURCE_EXHAUSTED": ResourceExhausted,
+    "REQUEST_LIMIT_EXCEEDED": RequestLimitExceeded,
+    "UNKNOWN": Unknown,
+    "DATA_LOSS": DataLoss,
+}

venv/lib/python3.9/site-packages/databricks/sdk/errors/private_link.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from urllib import parse
+
+import requests
+
+from ..environments import Cloud, get_environment_for_hostname
+from .platform import PermissionDenied
+
+
+@dataclass
+class _PrivateLinkInfo:
+    serviceName: str
+    endpointName: str
+    referencePage: str
+
+    def error_message(self):
+        return (
+            f"The requested workspace has {self.serviceName} enabled and is not accessible from the current network. "
+            f"Ensure that {self.serviceName} is properly configured and that your device has access to the "
+            f"{self.endpointName}. For more information, see {self.referencePage}."
+        )
+
+
+_private_link_info_map = {
+    Cloud.AWS: _PrivateLinkInfo(
+        serviceName="AWS PrivateLink",
+        endpointName="AWS VPC endpoint",
+        referencePage="https://docs.databricks.com/en/security/network/classic/privatelink.html",
+    ),
+    Cloud.AZURE: _PrivateLinkInfo(
+        serviceName="Azure Private Link",
+        endpointName="Azure Private Link endpoint",
+        referencePage="https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting",
+    ),
+    Cloud.GCP: _PrivateLinkInfo(
+        serviceName="Private Service Connect",
+        endpointName="GCP VPC endpoint",
+        referencePage="https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html",
+    ),
+}
+
+
+class PrivateLinkValidationError(PermissionDenied):
+    """Raised when a user tries to access a Private Link-enabled workspace, but the user's network does not have access
+    to the workspace."""
+
+
+def _is_private_link_redirect(resp: requests.Response) -> bool:
+    parsed = parse.urlparse(resp.url)
+    return parsed.path == "/login.html" and "error=private-link-validation-error" in parsed.query
+
+
+def _get_private_link_validation_error(url: str) -> PrivateLinkValidationError:
+    parsed = parse.urlparse(url)
+    env = get_environment_for_hostname(parsed.hostname)
+    return PrivateLinkValidationError(
+        message=_private_link_info_map[env.cloud].error_message(),
+        error_code="PRIVATE_LINK_VALIDATION_ERROR",
+        status_code=403,
+    )

venv/lib/python3.9/site-packages/databricks/sdk/errors/sdk.py

@@ -0,0 +1,6 @@
+class OperationFailed(RuntimeError):
+    pass
+
+
+class OperationTimeout(RuntimeError, TimeoutError):
+    pass

venv/lib/python3.9/site-packages/databricks/sdk/logger/__init__.py

@@ -0,0 +1 @@
+from .round_trip_logger import RoundTrip

venv/lib/python3.9/site-packages/databricks/sdk/logger/round_trip_logger.py

@@ -0,0 +1,127 @@
+import json
+import urllib.parse
+from typing import Any, Dict, List
+
+import requests
+
+
+class RoundTrip:
+    """
+    A utility class for converting HTTP requests and responses to strings.
+
+    :param response: The response object to stringify.
+    :param debug_headers: Whether to include headers in the generated string.
+    :param debug_truncate_bytes: The maximum number of bytes to include in the generated string.
+    :param raw: Whether the response is a stream or not. If True, the response will not be logged directly.
+    """
+
+    def __init__(
+        self,
+        response: requests.Response,
+        debug_headers: bool,
+        debug_truncate_bytes: int,
+        raw=False,
+    ):
+        self._debug_headers = debug_headers
+        self._debug_truncate_bytes = max(debug_truncate_bytes, 96)
+        self._raw = raw
+        self._response = response
+
+    def generate(self) -> str:
+        """
+        Generate a string representation of the request and response. The string will include the request method, URL,
+        headers, and body, as well as the response status code, reason, headers, and body. Outgoing information
+        will be prefixed with `>`, and incoming information will be prefixed with `<`.
+        :return: A string representation of the request.
+        """
+        request = self._response.request
+        url = urllib.parse.urlparse(request.url)
+        query = ""
+        if url.query:
+            query = f"?{urllib.parse.unquote(url.query)}"
+        sb = [f"{request.method} {urllib.parse.unquote(url.path)}{query}"]
+        if self._debug_headers:
+            for k, v in request.headers.items():
+                sb.append(f"> * {k}: {self._only_n_bytes(v, self._debug_truncate_bytes)}")
+        if request.body:
+            sb.append("> [raw stream]" if self._raw else self._redacted_dump("> ", request.body))
+        sb.append(f"< {self._response.status_code} {self._response.reason}")
+        if self._raw and self._response.headers.get("Content-Type", None) != "application/json":
+            # Raw streams with `Transfer-Encoding: chunked` do not have `Content-Type` header
+            sb.append("< [raw stream]")
+        elif self._response.content:
+            decoded = self._response.content.decode("utf-8", errors="replace")
+            sb.append(self._redacted_dump("< ", decoded))
+        return "\n".join(sb)
+
+    @staticmethod
+    def _mask(m: Dict[str, any]):
+        for k in m:
+            if k in {
+                "bytes_value",
+                "string_value",
+                "token_value",
+                "value",
+                "content",
+            }:
+                m[k] = "**REDACTED**"
+
+    @staticmethod
+    def _map_keys(m: Dict[str, any]) -> List[str]:
+        keys = list(m.keys())
+        keys.sort()
+        return keys
+
+    @staticmethod
+    def _only_n_bytes(j: str, num_bytes: int = 96) -> str:
+        diff = len(j.encode("utf-8")) - num_bytes
+        if diff > 0:
+            return f"{j[:num_bytes]}... ({diff} more bytes)"
+        return j
+
+    def _recursive_marshal_dict(self, m, budget) -> dict:
+        out = {}
+        self._mask(m)
+        for k in sorted(m.keys()):
+            raw = self._recursive_marshal(m[k], budget)
+            out[k] = raw
+            budget -= len(str(raw))
+        return out
+
+    def _recursive_marshal_list(self, s, budget) -> list:
+        out = []
+        for i in range(len(s)):
+            if i > 0 >= budget:
+                out.append("... (%d additional elements)" % (len(s) - len(out)))
+                break
+            raw = self._recursive_marshal(s[i], budget)
+            out.append(raw)
+            budget -= len(str(raw))
+        return out
+
+    def _recursive_marshal(self, v: Any, budget: int) -> Any:
+        if isinstance(v, dict):
+            return self._recursive_marshal_dict(v, budget)
+        elif isinstance(v, list):
+            return self._recursive_marshal_list(v, budget)
+        elif isinstance(v, str):
+            return self._only_n_bytes(v, self._debug_truncate_bytes)
+        else:
+            return v
+
+    def _redacted_dump(self, prefix: str, body: str) -> str:
+        if len(body) == 0:
+            return ""
+        try:
+            # Unmarshal body into primitive types.
+            tmp = json.loads(body)
+            max_bytes = 96
+            if self._debug_truncate_bytes > max_bytes:
+                max_bytes = self._debug_truncate_bytes
+            # Re-marshal body taking redaction and character limit into account.
+            raw = self._recursive_marshal(tmp, max_bytes)
+            return "\n".join([f"{prefix}{line}" for line in json.dumps(raw, indent=2).split("\n")])
+        except json.JSONDecodeError:
+            to_log = self._only_n_bytes(body, self._debug_truncate_bytes)
+            log_lines = [prefix + x.strip("\r") for x in to_log.split("\n")]
+            return "\n".join(log_lines)

venv/lib/python3.9/site-packages/databricks/sdk/mixins/compute.py

@@ -0,0 +1,280 @@
+import datetime
+import logging
+import re
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+from databricks.sdk.core import DatabricksError
+from databricks.sdk.errors import OperationFailed
+from databricks.sdk.service import compute
+
+_LOG = logging.getLogger("databricks.sdk")
+
+
+@dataclass
+class SemVer:
+    major: int
+    minor: int
+    patch: int
+    pre_release: Optional[str] = None
+    build: Optional[str] = None
+
+    # official https://semver.org/ recommendation: https://regex101.com/r/Ly7O1x/
+    # with addition of "x" wildcards for minor/patch versions. Also, patch version may be omitted.
+    _pattern = re.compile(
+        r"^"
+        r"(?P<major>0|[1-9]\d*)\.(?P<minor>x|0|[1-9]\d*)(\.(?P<patch>x|0|[1-9x]\d*))?"
+        r"(?:-(?P<pre_release>(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)"
+        r"(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?"
+        r"(?:\+(?P<build>[0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$"
+    )
+
+    @classmethod
+    def parse(cls, v: str) -> "SemVer":
+        if not v:
+            raise ValueError(f"Not a valid SemVer: {v}")
+        if v[0] != "v":
+            v = f"v{v}"
+        m = cls._pattern.match(v[1:])
+        if not m:
+            raise ValueError(f"Not a valid SemVer: {v}")
+        # patch and/or minor versions may be wildcards.
+        # for now, we're converting wildcards to zeroes.
+        minor = m.group("minor")
+        try:
+            patch = m.group("patch")
+        except IndexError:
+            patch = 0
+        return SemVer(
+            major=int(m.group("major")),
+            minor=0 if minor == "x" else int(minor),
+            patch=0 if patch == "x" or patch is None else int(patch),
+            pre_release=m.group("pre_release"),
+            build=m.group("build"),
+        )
+
+    def __lt__(self, other: "SemVer"):
+        if not other:
+            return False
+        if self.major != other.major:
+            return self.major < other.major
+        if self.minor != other.minor:
+            return self.minor < other.minor
+        if self.patch != other.patch:
+            return self.patch < other.patch
+        if self.pre_release != other.pre_release:
+            return self.pre_release < other.pre_release
+        if self.build != other.build:
+            return self.build < other.build
+        return False
+
+
+class ClustersExt(compute.ClustersAPI):
+    __doc__ = compute.ClustersAPI.__doc__
+
+    def select_spark_version(
+        self,
+        long_term_support: bool = False,
+        beta: bool = False,
+        latest: bool = True,
+        ml: bool = False,
+        genomics: bool = False,
+        gpu: bool = False,
+        scala: str = "2.12",
+        spark_version: str = None,
+        photon: bool = False,
+        graviton: bool = False,
+    ) -> str:
+        """Selects the latest Databricks Runtime Version.
+
+        :param long_term_support: bool
+        :param beta: bool
+        :param latest: bool
+        :param ml: bool
+        :param genomics: bool
+        :param gpu: bool
+        :param scala: str
+        :param spark_version: str
+        :param photon: bool
+        :param graviton: bool
+
+        :returns: `spark_version` compatible string
+        """
+        # Logic ported from https://github.com/databricks/databricks-sdk-go/blob/main/service/compute/spark_version.go
+        versions = []
+        sv = self.spark_versions()
+        for version in sv.versions:
+            if "-scala" + scala not in version.key:
+                continue
+            matches = (
+                ("apache-spark-" not in version.key)
+                and (("-ml-" in version.key) == ml)
+                and (("-hls-" in version.key) == genomics)
+                and (("-gpu-" in version.key) == gpu)
+                and (("-photon-" in version.key) == photon)
+                and (("-aarch64-" in version.key) == graviton)
+                and (("Beta" in version.name) == beta)
+            )
+            if matches and long_term_support:
+                matches = matches and (("LTS" in version.name) or ("-esr-" in version.key))
+            if matches and spark_version:
+                matches = matches and ("Apache Spark " + spark_version in version.name)
+            if matches:
+                versions.append(version.key)
+        if len(versions) < 1:
+            raise ValueError("spark versions query returned no results")
+        if len(versions) > 1:
+            if not latest:
+                raise ValueError("spark versions query returned multiple results")
+            versions = sorted(versions, key=SemVer.parse, reverse=True)
+        return versions[0]
+
+    @staticmethod
+    def _node_sorting_tuple(item: compute.NodeType) -> tuple:
+        local_disks = local_disk_size_gb = local_nvme_disk = local_nvme_disk_size_gb = 0
+        if item.node_instance_type is not None:
+            local_disks = item.node_instance_type.local_disks
+            local_nvme_disk = item.node_instance_type.local_nvme_disks
+            local_disk_size_gb = item.node_instance_type.local_disk_size_gb
+            local_nvme_disk_size_gb = item.node_instance_type.local_nvme_disk_size_gb
+        return (
+            item.is_deprecated,
+            item.num_cores,
+            item.memory_mb,
+            local_disks,
+            local_disk_size_gb,
+            local_nvme_disk,
+            local_nvme_disk_size_gb,
+            item.num_gpus,
+            item.instance_type_id,
+        )
+
+    @staticmethod
+    def _should_node_be_skipped(nt: compute.NodeType) -> bool:
+        if not nt.node_info:
+            return False
+        if not nt.node_info.status:
+            return False
+        val = compute.CloudProviderNodeStatus
+        for st in nt.node_info.status:
+            if st in (
+                val.NOT_AVAILABLE_IN_REGION,
+                val.NOT_ENABLED_ON_SUBSCRIPTION,
+            ):
+                return True
+        return False
+
+    def select_node_type(
+        self,
+        min_memory_gb: int = None,
+        gb_per_core: int = None,
+        min_cores: int = None,
+        min_gpus: int = None,
+        local_disk: bool = None,
+        local_disk_min_size: int = None,
+        category: str = None,
+        photon_worker_capable: bool = None,
+        photon_driver_capable: bool = None,
+        graviton: bool = None,
+        is_io_cache_enabled: bool = None,
+        support_port_forwarding: bool = None,
+        fleet: str = None,
+    ) -> str:
+        """Selects smallest available node type given the conditions.
+
+        :param min_memory_gb: int
+        :param gb_per_core: int
+        :param min_cores: int
+        :param min_gpus: int
+        :param local_disk: bool
+        :param local_disk_min_size: bool
+        :param category: bool
+        :param photon_worker_capable: bool
+        :param photon_driver_capable: bool
+        :param graviton: bool
+        :param is_io_cache_enabled: bool
+        :param support_port_forwarding: bool
+        :param fleet: bool
+
+        :returns: `node_type` compatible string
+        """
+        # Logic ported from https://github.com/databricks/databricks-sdk-go/blob/main/service/clusters/node_type.go
+        res = self.list_node_types()
+        types = sorted(res.node_types, key=self._node_sorting_tuple)
+        for nt in types:
+            if self._should_node_be_skipped(nt):
+                continue
+            gbs = nt.memory_mb // 1024
+            if fleet is not None and fleet not in nt.node_type_id:
+                continue
+            if min_memory_gb is not None and gbs < min_memory_gb:
+                continue
+            if gb_per_core is not None and gbs // nt.num_cores < gb_per_core:
+                continue
+            if min_cores is not None and nt.num_cores < min_cores:
+                continue
+            if (min_gpus is not None and nt.num_gpus < min_gpus) or (min_gpus == 0 and nt.num_gpus > 0):
+                continue
+            if local_disk or local_disk_min_size is not None:
+                instance_type = nt.node_instance_type
+                local_disks = int(instance_type.local_disks) if instance_type.local_disks else 0
+                local_nvme_disks = int(instance_type.local_nvme_disks) if instance_type.local_nvme_disks else 0
+                if instance_type is None or (local_disks < 1 and local_nvme_disks < 1):
+                    continue
+                local_disk_size_gb = instance_type.local_disk_size_gb if instance_type.local_disk_size_gb else 0
+                local_nvme_disk_size_gb = (
+                    instance_type.local_nvme_disk_size_gb if instance_type.local_nvme_disk_size_gb else 0
+                )
+                all_disks_size = local_disk_size_gb + local_nvme_disk_size_gb
+                if local_disk_min_size is not None and all_disks_size < local_disk_min_size:
+                    continue
+            if category is not None and not nt.category.lower() == category.lower():
+                continue
+            if is_io_cache_enabled and not nt.is_io_cache_enabled:
+                continue
+            if support_port_forwarding and not nt.support_port_forwarding:
+                continue
+            if photon_driver_capable and not nt.photon_driver_capable:
+                continue
+            if photon_worker_capable and not nt.photon_worker_capable:
+                continue
+            if graviton and nt.is_graviton != graviton:
+                continue
+            return nt.node_type_id
+        raise ValueError("cannot determine smallest node type")
+
+    def ensure_cluster_is_running(self, cluster_id: str) -> None:
+        """Ensures that given cluster is running, regardless of the current state"""
+        timeout = datetime.timedelta(minutes=20)
+        deadline = time.time() + timeout.total_seconds()
+        while time.time() < deadline:
+            try:
+                state = compute.State
+                info = self.get(cluster_id)
+                if info.state == state.RUNNING:
+                    return
+                elif info.state == state.TERMINATED:
+                    self.start(cluster_id).result()
+                    return
+                elif info.state == state.TERMINATING:
+                    self.wait_get_cluster_terminated(cluster_id)
+                    self.start(cluster_id).result()
+                    return
+                elif info.state in (
+                    state.PENDING,
+                    state.RESIZING,
+                    state.RESTARTING,
+                ):
+                    self.wait_get_cluster_running(cluster_id)
+                    return
+                elif info.state in (state.ERROR, state.UNKNOWN):
+                    raise RuntimeError(f"Cluster {info.cluster_name} is {info.state}: {info.state_message}")
+            except DatabricksError as e:
+                if e.error_code == "INVALID_STATE":
+                    _LOG.debug(f"Cluster was started by other process: {e} Retrying.")
+                    continue
+                raise e
+            except OperationFailed as e:
+                _LOG.debug("Operation failed, retrying", exc_info=e)
+        raise TimeoutError(f"timed out after {timeout}")

venv/lib/python3.9/site-packages/databricks/sdk/mixins/files_utils.py

@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+import os
+import threading
+from dataclasses import dataclass
+from typing import Any, BinaryIO, Callable, Iterable, Optional
+
+
+@dataclass
+class CreateDownloadUrlResponse:
+    """Response from the download URL API call."""
+
+    url: str
+    """The presigned URL to download the file."""
+    headers: dict[str, str]
+    """Headers to use when making the download request."""
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> CreateDownloadUrlResponse:
+        """Create an instance from a dictionary."""
+        if "url" not in data:
+            raise ValueError("Missing 'url' in response data")
+        headers = data["headers"] if "headers" in data else {}
+        parsed_headers = {x["name"]: x["value"] for x in headers}
+        return cls(url=data["url"], headers=parsed_headers)
+
+
+class _ConcatenatedInputStream(BinaryIO):
+    """This class joins two input streams into one."""
+
+    def __init__(self, head_stream: BinaryIO, tail_stream: BinaryIO):
+        if not head_stream.readable():
+            raise ValueError("head_stream is not readable")
+        if not tail_stream.readable():
+            raise ValueError("tail_stream is not readable")
+
+        self._head_stream = head_stream
+        self._tail_stream = tail_stream
+        self._head_size = None
+        self._tail_size = None
+
+    def close(self) -> None:
+        try:
+            self._head_stream.close()
+        finally:
+            self._tail_stream.close()
+
+    def fileno(self) -> int:
+        raise AttributeError()
+
+    def flush(self) -> None:
+        raise NotImplementedError("Stream is not writable")
+
+    def isatty(self) -> bool:
+        raise NotImplementedError()
+
+    def read(self, __n: int = -1) -> bytes:
+        head = self._head_stream.read(__n)
+        remaining_bytes = __n - len(head) if __n >= 0 else __n
+        tail = self._tail_stream.read(remaining_bytes)
+        return head + tail
+
+    def readable(self) -> bool:
+        return True
+
+    def readline(self, __limit: int = -1) -> bytes:
+        # Read and return one line from the stream.
+        # If __limit is specified, at most __limit bytes will be read.
+        # The line terminator is always b'\n' for binary files.
+        head = self._head_stream.readline(__limit)
+        if len(head) > 0 and head[-1:] == b"\n":
+            # end of line happened before (or at) the limit
+            return head
+
+        # if __limit >= 0, len(head) can't exceed limit
+        remaining_bytes = __limit - len(head) if __limit >= 0 else __limit
+        tail = self._tail_stream.readline(remaining_bytes)
+        return head + tail
+
+    def readlines(self, __hint: int = -1) -> list[bytes]:
+        # Read and return a list of lines from the stream.
+        # Hint can be specified to control the number of lines read: no more lines will be read
+        # If the total size (in bytes/characters) of all lines so far exceeds hint.
+
+        # In fact, BytesIO(bytes) will not read next line if total size of all lines
+        # *equals or* exceeds hint.
+
+        head_result = self._head_stream.readlines(__hint)
+        head_total_bytes = sum(len(line) for line in head_result)
+
+        if 0 < __hint <= head_total_bytes and head_total_bytes > 0:
+            # We reached (or passed) the hint by reading from head_stream, or exhausted head_stream.
+
+            if head_result[-1][-1:] == b"\n":
+                # If we reached/passed the hint and also stopped at the line break, return.
+                return head_result
+
+            # Reading from head_stream could have stopped only because the stream was exhausted
+            if len(self._head_stream.read(1)) > 0:
+                raise ValueError(
+                    f"Stream reading finished prematurely after reading {head_total_bytes} bytes, reaching or exceeding hint {__hint}"
+                )
+
+            # We need to finish reading the current line, now from tail_stream.
+
+            tail_result = self._tail_stream.readlines(1)  # We will only read the first line from tail_stream.
+            assert len(tail_result) <= 1
+            if len(tail_result) > 0:
+                # We will then append the tail as the last line of the result.
+                return head_result[:-1] + [head_result[-1] + tail_result[0]]
+            else:
+                return head_result
+
+        # We did not reach the hint by reading head_stream but exhausted it, continue reading from tail_stream
+        # with an adjusted hint
+        if __hint >= 0:
+            remaining_bytes = __hint - head_total_bytes
+        else:
+            remaining_bytes = __hint
+
+        tail_result = self._tail_stream.readlines(remaining_bytes)
+
+        if head_total_bytes > 0 and head_result[-1][-1:] != b"\n" and len(tail_result) > 0:
+            # If head stream does not end with the line break, we need to concatenate
+            # the last line of the head result and the first line of tail result
+            return head_result[:-1] + [head_result[-1] + tail_result[0]] + tail_result[1:]
+        else:
+            # Otherwise, just append two lists of lines.
+            return head_result + tail_result
+
+    def _get_stream_size(self, stream: BinaryIO) -> int:
+        prev_offset = stream.tell()
+        try:
+            stream.seek(0, os.SEEK_END)
+            return stream.tell()
+        finally:
+            stream.seek(prev_offset, os.SEEK_SET)
+
+    def _get_head_size(self) -> int:
+        if self._head_size is None:
+            self._head_size = self._get_stream_size(self._head_stream)
+        return self._head_size
+
+    def _get_tail_size(self) -> int:
+        if self._tail_size is None:
+            self._tail_size = self._get_stream_size(self._tail_stream)
+        return self._tail_size
+
+    def seek(self, __offset: int, __whence: int = os.SEEK_SET) -> int:
+        if not self.seekable():
+            raise NotImplementedError("Stream is not seekable")
+
+        if __whence == os.SEEK_SET:
+            if __offset < 0:
+                # Follow native buffer behavior
+                raise ValueError(f"Negative seek value: {__offset}")
+
+            head_size = self._get_head_size()
+
+            if __offset <= head_size:
+                self._head_stream.seek(__offset, os.SEEK_SET)
+                self._tail_stream.seek(0, os.SEEK_SET)
+            else:
+                self._head_stream.seek(0, os.SEEK_END)  # move head stream to the end
+                self._tail_stream.seek(__offset - head_size, os.SEEK_SET)
+
+        elif __whence == os.SEEK_CUR:
+            current_offset = self.tell()
+            new_offset = current_offset + __offset
+            if new_offset < 0:
+                # gracefully don't seek before start
+                new_offset = 0
+            self.seek(new_offset, os.SEEK_SET)
+
+        elif __whence == os.SEEK_END:
+            if __offset > 0:
+                # Python allows to seek beyond the end of stream.
+
+                # Move head to EOF and tail to (EOF + offset), so subsequent tell()
+                # returns len(head) + len(tail) + offset, same as for native buffer
+                self._head_stream.seek(0, os.SEEK_END)
+                self._tail_stream.seek(__offset, os.SEEK_END)
+            else:
+                self._tail_stream.seek(__offset, os.SEEK_END)
+                tail_pos = self._tail_stream.tell()
+                if tail_pos > 0:
+                    # target position lies within the tail, move head to EOF
+                    self._head_stream.seek(0, os.SEEK_END)
+                else:
+                    tail_size = self._get_tail_size()
+                    self._head_stream.seek(__offset + tail_size, os.SEEK_END)
+        else:
+            raise ValueError(__whence)
+        return self.tell()
+
+    def seekable(self) -> bool:
+        return self._head_stream.seekable() and self._tail_stream.seekable()
+
+    def __getattribute__(self, name: str) -> Any:
+        if name == "fileno":
+            raise AttributeError()
+        elif name in ["tell", "seek"] and not self.seekable():
+            raise AttributeError()
+
+        return super().__getattribute__(name)
+
+    def tell(self) -> int:
+        if not self.seekable():
+            raise NotImplementedError()
+
+        # Assuming that tail stream stays at 0 until head stream is exhausted
+        return self._head_stream.tell() + self._tail_stream.tell()
+
+    def truncate(self, __size: Optional[int] = None) -> int:
+        raise NotImplementedError("Stream is not writable")
+
+    def writable(self) -> bool:
+        return False
+
+    def write(self, __s: bytes) -> int:
+        raise NotImplementedError("Stream is not writable")
+
+    def writelines(self, __lines: Iterable[bytes]) -> None:
+        raise NotImplementedError("Stream is not writable")
+
+    def __next__(self) -> bytes:
+        # IOBase [...] supports the iterator protocol, meaning that an IOBase object can be
+        # iterated over yielding the lines in a stream. [...] See readline().
+        result = self.readline()
+        if len(result) == 0:
+            raise StopIteration
+        return result
+
+    def __iter__(self) -> "BinaryIO":
+        return self
+
+    def __enter__(self) -> "BinaryIO":
+        self._head_stream.__enter__()
+        self._tail_stream.__enter__()
+        return self
+
+    def __exit__(self, __type, __value, __traceback) -> None:
+        self._head_stream.__exit__(__type, __value, __traceback)
+        self._tail_stream.__exit__(__type, __value, __traceback)
+
+    def __str__(self) -> str:
+        return f"Concat: {self._head_stream}, {self._tail_stream}]"
+
+
+class _PresignedUrlDistributor:
+    """
+    Distributes and manages presigned URLs for downloading files.
+
+    This class ensures thread-safe access to a presigned URL, allowing retrieval and invalidation.
+    When the URL is invalidated, a new one will be fetched using the provided function.
+    """
+
+    def __init__(self, get_new_url_func: Callable[[], CreateDownloadUrlResponse]):
+        """
+        Initialize the distributor.
+
+        Args:
+            get_new_url_func: A callable that returns a new presigned URL response.
+        """
+        self._get_new_url_func = get_new_url_func
+        self._current_url = None
+        self.current_version = 0
+        self.lock = threading.RLock()
+
+    def get_url(self) -> tuple[CreateDownloadUrlResponse, int]:
+        """
+        Get the current presigned URL and its version.
+
+        Returns:
+            A tuple containing the current presigned URL response and its version.
+        """
+        with self.lock:
+            if self._current_url is None:
+                self._current_url = self._get_new_url_func()
+            return self._current_url, self.current_version
+
+    def invalidate_url(self, version: int) -> None:
+        """
+        Invalidate the current presigned URL if the version matches. If the version does not match,
+        the URL remains unchanged. This ensures that only the most recent version can invalidate the URL.
+
+        Args:
+            version: The version to check before invalidating the URL.
+        """
+        with self.lock:
+            if version == self.current_version:
+                self._current_url = None
+                self.current_version += 1

venv/lib/python3.9/site-packages/databricks/sdk/mixins/jobs.py

@@ -0,0 +1,230 @@
+from typing import Iterator, Optional
+
+from databricks.sdk.service import jobs
+from databricks.sdk.service.jobs import BaseJob, BaseRun, Job, RunType
+
+
+class JobsExt(jobs.JobsAPI):
+
+    def list(
+        self,
+        *,
+        expand_tasks: Optional[bool] = None,
+        limit: Optional[int] = None,
+        name: Optional[str] = None,
+        offset: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> Iterator[BaseJob]:
+        """List jobs.
+
+        Retrieves a list of jobs. If the job has multiple pages of tasks, job_clusters, parameters or environments,
+        it will paginate through all pages and aggregate the results.
+
+        :param expand_tasks: bool (optional)
+          Whether to include task and cluster details in the response. Note that in API 2.2, only the first
+          100 elements will be shown. Use :method:jobs/get to paginate through all tasks and clusters.
+        :param limit: int (optional)
+          The number of jobs to return. This value must be greater than 0 and less or equal to 100. The
+          default value is 20.
+        :param name: str (optional)
+          A filter on the list based on the exact (case insensitive) job name.
+        :param offset: int (optional)
+          The offset of the first job to return, relative to the most recently created job. Deprecated since
+          June 2023. Use `page_token` to iterate through the pages instead.
+        :param page_token: str (optional)
+          Use `next_page_token` or `prev_page_token` returned from the previous request to list the next or
+          previous page of jobs respectively.
+
+        :returns: Iterator over :class:`BaseJob`
+        """
+        # fetch jobs with limited elements in top level arrays
+        jobs_list = super().list(
+            expand_tasks=expand_tasks,
+            limit=limit,
+            name=name,
+            offset=offset,
+            page_token=page_token,
+        )
+        if not expand_tasks:
+            yield from jobs_list
+
+        # fully fetch all top level arrays for each job in the list
+        for job in jobs_list:
+            if job.has_more:
+                job_from_get_call = self.get(job.job_id)
+                job.settings.tasks = job_from_get_call.settings.tasks
+                job.settings.job_clusters = job_from_get_call.settings.job_clusters
+                job.settings.parameters = job_from_get_call.settings.parameters
+                job.settings.environments = job_from_get_call.settings.environments
+            # Remove has_more fields for each job in the list.
+            # This field in Jobs API 2.2 is useful for pagination. It indicates if there are more than 100 tasks or job_clusters in the job.
+            # This function hides pagination details from the user. So the field does not play useful role here.
+            if hasattr(job, "has_more"):
+                delattr(job, "has_more")
+            yield job
+
+    def list_runs(
+        self,
+        *,
+        active_only: Optional[bool] = None,
+        completed_only: Optional[bool] = None,
+        expand_tasks: Optional[bool] = None,
+        job_id: Optional[int] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        page_token: Optional[str] = None,
+        run_type: Optional[RunType] = None,
+        start_time_from: Optional[int] = None,
+        start_time_to: Optional[int] = None,
+    ) -> Iterator[BaseRun]:
+        """List job runs.
+
+        List runs in descending order by start time. If the job has multiple pages of tasks, job_clusters, parameters or repair history,
+        it will paginate through all pages and aggregate the results.
+
+        :param active_only: bool (optional)
+          If active_only is `true`, only active runs are included in the results; otherwise, lists both active
+          and completed runs. An active run is a run in the `QUEUED`, `PENDING`, `RUNNING`, or `TERMINATING`.
+          This field cannot be `true` when completed_only is `true`.
+        :param completed_only: bool (optional)
+          If completed_only is `true`, only completed runs are included in the results; otherwise, lists both
+          active and completed runs. This field cannot be `true` when active_only is `true`.
+        :param expand_tasks: bool (optional)
+          Whether to include task and cluster details in the response. Note that in API 2.2, only the first
+          100 elements will be shown. Use :method:jobs/getrun to paginate through all tasks and clusters.
+        :param job_id: int (optional)
+          The job for which to list runs. If omitted, the Jobs service lists runs from all jobs.
+        :param limit: int (optional)
+          The number of runs to return. This value must be greater than 0 and less than 25. The default value
+          is 20. If a request specifies a limit of 0, the service instead uses the maximum limit.
+        :param offset: int (optional)
+          The offset of the first run to return, relative to the most recent run. Deprecated since June 2023.
+          Use `page_token` to iterate through the pages instead.
+        :param page_token: str (optional)
+          Use `next_page_token` or `prev_page_token` returned from the previous request to list the next or
+          previous page of runs respectively.
+        :param run_type: :class:`RunType` (optional)
+          The type of runs to return. For a description of run types, see :method:jobs/getRun.
+        :param start_time_from: int (optional)
+          Show runs that started _at or after_ this value. The value must be a UTC timestamp in milliseconds.
+          Can be combined with _start_time_to_ to filter by a time range.
+        :param start_time_to: int (optional)
+          Show runs that started _at or before_ this value. The value must be a UTC timestamp in milliseconds.
+          Can be combined with _start_time_from_ to filter by a time range.
+
+        :returns: Iterator over :class:`BaseRun`
+        """
+        # fetch runs with limited elements in top level arrays
+        runs_list = super().list_runs(
+            active_only=active_only,
+            completed_only=completed_only,
+            expand_tasks=expand_tasks,
+            job_id=job_id,
+            limit=limit,
+            offset=offset,
+            page_token=page_token,
+            run_type=run_type,
+            start_time_from=start_time_from,
+            start_time_to=start_time_to,
+        )
+
+        if not expand_tasks:
+            yield from runs_list
+
+        # fully fetch all top level arrays for each run in the list
+        for run in runs_list:
+            if run.has_more:
+                run_from_get_call = self.get_run(run.run_id)
+                run.tasks = run_from_get_call.tasks
+                run.job_clusters = run_from_get_call.job_clusters
+                run.job_parameters = run_from_get_call.job_parameters
+                run.repair_history = run_from_get_call.repair_history
+            # Remove has_more fields for each run in the list.
+            # This field in Jobs API 2.2 is useful for pagination. It indicates if there are more than 100 tasks or job_clusters in the run.
+            # This function hides pagination details from the user. So the field does not play useful role here.
+            if hasattr(run, "has_more"):
+                delattr(run, "has_more")
+            yield run
+
+    def get_run(
+        self,
+        run_id: int,
+        *,
+        include_history: Optional[bool] = None,
+        include_resolved_values: Optional[bool] = None,
+        page_token: Optional[str] = None,
+    ) -> jobs.Run:
+        """Get a single job run.
+
+        Retrieve the metadata of a run. If a run has multiple pages of tasks, it will paginate through all pages of tasks, iterations, job_clusters, job_parameters, and repair history.
+
+        :param run_id: int
+          The canonical identifier of the run for which to retrieve the metadata. This field is required.
+        :param include_history: bool (optional)
+          Whether to include the repair history in the response.
+        :param include_resolved_values: bool (optional)
+          Whether to include resolved parameter values in the response.
+        :param page_token: str (optional)
+          To list the next page of job tasks, set this field to the value of the `next_page_token` returned in
+          the GetJob response.
+
+        :returns: :class:`Run`
+        """
+        run = super().get_run(
+            run_id,
+            include_history=include_history,
+            include_resolved_values=include_resolved_values,
+            page_token=page_token,
+        )
+
+        # When querying a Job run, a page token is returned when there are more than 100 tasks. No iterations are defined for a Job run. Therefore, the next page in the response only includes the next page of tasks.
+        # When querying a ForEach task run, a page token is returned when there are more than 100 iterations. Only a single task is returned, corresponding to the ForEach task itself. Therefore, the client only reads the iterations from the next page and not the tasks.
+        is_paginating_iterations = run.iterations is not None and len(run.iterations) > 0
+
+        # runs/get response includes next_page_token as long as there are more pages to fetch.
+        while run.next_page_token is not None:
+            next_run = super().get_run(
+                run_id,
+                include_history=include_history,
+                include_resolved_values=include_resolved_values,
+                page_token=run.next_page_token,
+            )
+            if is_paginating_iterations:
+                run.iterations.extend(next_run.iterations)
+            else:
+                run.tasks.extend(next_run.tasks)
+            # Each new page of runs/get response includes the next page of the job_clusters, job_parameters, and repair history.
+            run.job_clusters.extend(next_run.job_clusters)
+            run.job_parameters.extend(next_run.job_parameters)
+            run.repair_history.extend(next_run.repair_history)
+            run.next_page_token = next_run.next_page_token
+
+        return run
+
+    def get(self, job_id: int, *, page_token: Optional[str] = None) -> Job:
+        """Get a single job.
+
+        Retrieves the details for a single job. If the job has multiple pages of tasks, job_clusters, parameters or environments,
+        it will paginate through all pages and aggregate the results.
+
+        :param job_id: int
+          The canonical identifier of the job to retrieve information about. This field is required.
+        :param page_token: str (optional)
+          Use `next_page_token` returned from the previous GetJob to request the next page of the job's
+          sub-resources.
+
+        :returns: :class:`Job`
+        """
+        job = super().get(job_id, page_token=page_token)
+
+        # jobs/get response includes next_page_token as long as there are more pages to fetch.
+        while job.next_page_token is not None:
+            next_job = super().get(job_id, page_token=job.next_page_token)
+            # Each new page of jobs/get response includes the next page of the tasks, job_clusters, job_parameters, and environments.
+            job.settings.tasks.extend(next_job.settings.tasks)
+            job.settings.job_clusters.extend(next_job.settings.job_clusters)
+            job.settings.parameters.extend(next_job.settings.parameters)
+            job.settings.environments.extend(next_job.settings.environments)
+            job.next_page_token = next_job.next_page_token
+
+        return job

venv/lib/python3.9/site-packages/databricks/sdk/mixins/open_ai_client.py

@@ -0,0 +1,209 @@
+import json as js
+import warnings
+from typing import Dict, Optional
+
+from requests import Response
+
+from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
+                                            HttpRequestResponse,
+                                            ServingEndpointsAPI)
+
+
+class ServingEndpointsExt(ServingEndpointsAPI):
+
+    # Using the HTTP Client to pass in the databricks authorization
+    # This method will be called on every invocation, so when using with model serving will always get the refreshed token
+    def _get_authorized_http_client(self):
+        import httpx
+
+        class BearerAuth(httpx.Auth):
+
+            def __init__(self, get_headers_func):
+                self.get_headers_func = get_headers_func
+
+            def auth_flow(self, request: httpx.Request) -> httpx.Request:
+                auth_headers = self.get_headers_func()
+                request.headers["Authorization"] = auth_headers["Authorization"]
+                yield request
+
+        databricks_token_auth = BearerAuth(self._api._cfg.authenticate)
+
+        # Create an HTTP client with Bearer Token authentication
+        http_client = httpx.Client(auth=databricks_token_auth)
+        return http_client
+
+    def get_open_ai_client(self, **kwargs):
+        """Create an OpenAI client configured for Databricks Model Serving.
+
+        .. deprecated::
+            This method is deprecated. Please install the `databricks-openai` package
+            and use `from databricks_openai import DatabricksOpenAI` instead.
+            See https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_openai.html for more information.
+
+        Returns an OpenAI client instance that is pre-configured to send requests to
+        Databricks Model Serving endpoints. The client uses Databricks authentication
+        to query endpoints within the workspace associated with the current WorkspaceClient
+        instance.
+
+        Args:
+            **kwargs: Additional parameters to pass to the OpenAI client constructor.
+                Common parameters include:
+                - timeout (float): Request timeout in seconds (e.g., 30.0)
+                - max_retries (int): Maximum number of retries for failed requests (e.g., 3)
+                - default_headers (dict): Additional headers to include with requests
+                - default_query (dict): Additional query parameters to include with requests
+
+                Any parameter accepted by the OpenAI client constructor can be passed here,
+                except for the following parameters which are reserved for Databricks integration:
+                base_url, api_key, http_client
+
+        Returns:
+            OpenAI: An OpenAI client instance configured for Databricks Model Serving.
+
+        Raises:
+            ImportError: If the OpenAI library is not installed.
+            ValueError: If any reserved Databricks parameters are provided in kwargs.
+
+        Example:
+            >>> client = workspace_client.serving_endpoints.get_open_ai_client()
+            >>> # With custom timeout and retries
+            >>> client = workspace_client.serving_endpoints.get_open_ai_client(
+            ...     timeout=30.0,
+            ...     max_retries=5
+            ... )
+        """
+        warnings.warn(
+            "get_open_ai_client() is deprecated. Please install the databricks-openai package "
+            "and use 'from databricks_openai import DatabricksOpenAI' instead. "
+            "See https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_openai.html for more information.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        try:
+            from openai import OpenAI
+        except Exception:
+            raise ImportError(
+                "Open AI is not installed. Please install the Databricks SDK with the following command `pip install databricks-sdk[openai]`"
+            )
+
+        # Check for reserved parameters that should not be overridden
+        reserved_params = {"base_url", "api_key", "http_client"}
+        conflicting_params = reserved_params.intersection(kwargs.keys())
+        if conflicting_params:
+            raise ValueError(
+                f"Cannot override reserved Databricks parameters: {', '.join(sorted(conflicting_params))}. "
+                f"These parameters are automatically configured for Databricks Model Serving."
+            )
+
+        # Default parameters that are required for Databricks integration
+        client_params = {
+            "base_url": self._api._cfg.host + "/serving-endpoints",
+            "api_key": "no-token",  # Passing in a placeholder to pass validations, this will not be used
+            "http_client": self._get_authorized_http_client(),
+        }
+
+        # Update with any additional parameters passed by the user
+        client_params.update(kwargs)
+
+        return OpenAI(**client_params)
+
+    def get_langchain_chat_open_ai_client(self, model):
+        """Create a LangChain ChatOpenAI client configured for Databricks Model Serving.
+
+        .. deprecated::
+            This method is deprecated. Please install the `databricks-langchain` package
+            and use `from databricks_langchain import ChatDatabricks` instead.
+            See https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html for more information.
+        """
+        warnings.warn(
+            "get_langchain_chat_open_ai_client() is deprecated. Please install the databricks-langchain package "
+            "and use 'from databricks_langchain import ChatDatabricks' instead. "
+            "See https://pypi.org/project/databricks-langchain/ for more information.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        try:
+            from langchain_openai import ChatOpenAI
+        except Exception:
+            raise ImportError(
+                "Langchain Open AI is not installed. Please install the Databricks SDK with the following command `pip install databricks-sdk[openai]` and ensure you are using python>3.7"
+            )
+
+        return ChatOpenAI(
+            model=model,
+            openai_api_base=self._api._cfg.host + "/serving-endpoints",
+            api_key="no-token",  # Passing in a placeholder to pass validations, this will not be used
+            http_client=self._get_authorized_http_client(),
+        )
+
+    def http_request(
+        self,
+        conn: str,
+        method: ExternalFunctionRequestHttpMethod,
+        path: str,
+        *,
+        headers: Optional[Dict[str, str]] = None,
+        json: Optional[Dict[str, str]] = None,
+        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
+          The connection name to use. This is required to identify the external connection.
+        :param method: :class:`ExternalFunctionRequestHttpMethod`
+          The HTTP method to use (e.g., 'GET', 'POST'). This is required.
+        :param path: str
+          The relative path for the API endpoint. This is required.
+        :param headers: Dict[str,str] (optional)
+          Additional headers for the request. If not provided, only auth headers from connections would be
+          passed.
+        :param json: Dict[str,str] (optional)
+          JSON payload for the request.
+        :param params: Dict[str,str] (optional)
+          Query parameters for the request.
+        :returns: :class:`Response`
+        """
+        response = Response()
+        response.status_code = 200
+
+        # 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
+        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.")
+
+        # 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

venv/lib/python3.9/site-packages/databricks/sdk/mixins/sharing.py

@@ -0,0 +1,44 @@
+from typing import Iterator, Optional
+
+from databricks.sdk.service import sharing
+from databricks.sdk.service.sharing import ShareInfo
+
+
+class SharesExt(sharing.SharesAPI):
+    def list(self, *, max_results: Optional[int] = None, page_token: Optional[str] = None) -> Iterator[ShareInfo]:
+        """Gets an array of data object shares from the metastore. The caller must be a metastore admin or the
+        owner of the share. There is no guarantee of a specific ordering of the elements in the array.
+
+        :param max_results: int (optional)
+          Maximum number of shares to return. - when set to 0, the page length is set to a server configured
+          value (recommended); - when set to a value greater than 0, the page length is the minimum of this
+          value and a server configured value; - when set to a value less than 0, an invalid parameter error
+          is returned; - If not set, all valid shares are returned (not recommended). - Note: The number of
+          returned shares might be less than the specified max_results size, even zero. The only definitive
+          indication that no further shares can be fetched is when the next_page_token is unset from the
+          response.
+        :param page_token: str (optional)
+          Opaque pagination token to go to next page based on previous query.
+
+        :returns: Iterator over :class:`ShareInfo`
+        """
+
+        query = {}
+        if max_results is not None:
+            query["max_results"] = max_results
+        if page_token is not None:
+            query["page_token"] = page_token
+        headers = {
+            "Accept": "application/json",
+        }
+
+        if "max_results" not in query:
+            query["max_results"] = 0
+        while True:
+            json = self._api.do("GET", "/api/2.1/unity-catalog/shares", query=query, headers=headers)
+            if "shares" in json:
+                for v in json["shares"]:
+                    yield ShareInfo.from_dict(v)
+            if "next_page_token" not in json or not json["next_page_token"]:
+                return
+            query["page_token"] = json["next_page_token"]

venv/lib/python3.9/site-packages/databricks/sdk/mixins/workspace.py

@@ -0,0 +1,117 @@
+from typing import Any, BinaryIO, Iterator, Optional, Union
+
+from ..core import DatabricksError
+from ..service.workspace import (ExportFormat, ImportFormat, Language,
+                                 ObjectInfo, ObjectType, WorkspaceAPI)
+
+
+def _fqcn(x: Any) -> str:
+    return f"{x.__module__}.{x.__name__}"
+
+
+class WorkspaceExt(WorkspaceAPI):
+    __doc__ = WorkspaceAPI.__doc__
+
+    def list(
+        self,
+        path: str,
+        *,
+        notebooks_modified_after: Optional[int] = None,
+        recursive: Optional[bool] = False,
+        **kwargs,
+    ) -> Iterator[ObjectInfo]:
+        """List workspace objects
+
+        :param recursive: bool
+            Optionally invoke recursive traversal
+
+        :returns: Iterator of workspaceObjectInfo
+        """
+        parent_list = super().list
+        queue = [path]
+        while queue:
+            path, queue = queue[0], queue[1:]
+            for object_info in parent_list(path, notebooks_modified_after=notebooks_modified_after):
+                if recursive and object_info.object_type == ObjectType.DIRECTORY:
+                    queue.append(object_info.path)
+                    continue
+                yield object_info
+
+    def upload(
+        self,
+        path: str,
+        content: Union[bytes, BinaryIO],
+        *,
+        format: Optional[ImportFormat] = None,
+        language: Optional[Language] = None,
+        overwrite: Optional[bool] = False,
+    ) -> None:
+        """
+        Uploads a workspace object (for example, a notebook or file) or the contents of an entire
+        directory (`DBC` format).
+
+        Errors:
+         * `RESOURCE_ALREADY_EXISTS`: if `path` already exists no `overwrite=True`.
+         * `INVALID_PARAMETER_VALUE`: if `format` and `content` values are not compatible.
+
+        :param path:     target location of the file on workspace.
+        :param content:  the contents as either raw binary data `bytes` or a file-like the file-like `io.BinaryIO` of the `path` contents.
+        :param format:   By default, `ImportFormat.SOURCE`. If using `ImportFormat.AUTO` the `path`
+                         is imported or exported as either a workspace file or a notebook, depending
+                         on an analysis of the `item`’s extension and the header content provided in
+                         the request. In addition, if the `path` is imported as a notebook, then
+                         the `item`’s extension is automatically removed.
+        :param language: Only required if using `ExportFormat.SOURCE`.
+        """
+        if format is not None and not isinstance(format, ImportFormat):
+            raise ValueError(f"format is expected to be {_fqcn(ImportFormat)}, but got {_fqcn(format.__class__)}")
+        if (not format or format == ImportFormat.SOURCE) and not language:
+            suffixes = {
+                ".py": Language.PYTHON,
+                ".sql": Language.SQL,
+                ".scala": Language.SCALA,
+                ".R": Language.R,
+            }
+            for sfx, lang in suffixes.items():
+                if path.endswith(sfx):
+                    language = lang
+                    break
+        if language is not None and not isinstance(language, Language):
+            raise ValueError(f"language is expected to be {_fqcn(Language)}, but got {_fqcn(language.__class__)}")
+        data = {"path": path}
+        if format:
+            data["format"] = format.value
+        if language:
+            data["language"] = language.value
+        if overwrite:
+            data["overwrite"] = "true"
+        try:
+            return self._api.do(
+                "POST",
+                "/api/2.0/workspace/import",
+                files={"content": content},
+                data=data,
+            )
+        except DatabricksError as e:
+            if e.error_code == "INVALID_PARAMETER_VALUE":
+                msg = f"Perhaps you forgot to specify the `format=ImportFormat.AUTO`. {e}"
+                raise DatabricksError(message=msg, error_code=e.error_code)
+            else:
+                raise e
+
+    def download(self, path: str, *, format: Optional[ExportFormat] = None) -> BinaryIO:
+        """
+        Downloads notebook or file from the workspace
+
+        :param path:     location of the file or notebook on workspace.
+        :param format:   By default, `ExportFormat.SOURCE`. If using `ExportFormat.AUTO` the `path`
+                         is imported or exported as either a workspace file or a notebook, depending
+                         on an analysis of the `item`’s extension and the header content provided in
+                         the request.
+        :return:         file-like `io.BinaryIO` of the `path` contents.
+        """
+        query = {"path": path, "direct_download": "true"}
+        if format:
+            query["format"] = format.value
+        response = self._api.do("GET", "/api/2.0/workspace/export", query=query, raw=True)
+        return response["contents"]

venv/lib/python3.9/site-packages/databricks/sdk/oauth.py

@@ -0,0 +1,945 @@
+import base64
+import functools
+import hashlib
+import json
+import logging
+import os
+import secrets
+import threading
+import urllib.parse
+import webbrowser
+from abc import abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from enum import Enum
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any, Callable, Dict, List, Optional
+
+import requests
+import requests.auth
+
+from ._base_client import _BaseClient, _fix_host_if_needed
+
+# Error code for PKCE flow in Azure Active Directory, that gets additional retry.
+# See https://stackoverflow.com/a/75466778/277035 for more info
+NO_ORIGIN_FOR_SPA_CLIENT_ERROR = "AADSTS9002327"
+
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AuthorizationDetail:
+    type: str
+    object_type: str
+    object_path: str
+    actions: List[str]
+
+    def as_dict(self) -> dict:
+        return {
+            "type": self.type,
+            "object_type": self.object_type,
+            "object_path": self.object_path,
+            "actions": self.actions,
+        }
+
+    def from_dict(self, d: dict) -> "AuthorizationDetail":
+        return AuthorizationDetail(
+            type=d.get("type"),
+            object_type=d.get("object_type"),
+            object_path=d.get("object_path"),
+            actions=d.get("actions"),
+        )
+
+
+class IgnoreNetrcAuth(requests.auth.AuthBase):
+    """This auth method is a no-op.
+
+    We use it to force requestslib to not use .netrc to write auth headers
+    when making .post() requests to the oauth token endpoints, since these
+    don't require authentication.
+
+    In cases where .netrc is outdated or corrupt, these requests will fail.
+
+    See issue #121
+    """
+
+    def __call__(self, r):
+        return r
+
+
+@dataclass
+class OidcEndpoints:
+    """
+    The endpoints used for OAuth-based authentication in Databricks.
+    """
+
+    authorization_endpoint: str  # ../v1/authorize
+    """The authorization endpoint for the OAuth flow. The user-agent should be directed to this endpoint in order for
+    the user to login and authorize the client for user-to-machine (U2M) flows."""
+
+    token_endpoint: str  # ../v1/token
+    """The token endpoint for the OAuth flow."""
+
+    @staticmethod
+    def from_dict(d: dict) -> "OidcEndpoints":
+        return OidcEndpoints(
+            authorization_endpoint=d.get("authorization_endpoint"),
+            token_endpoint=d.get("token_endpoint"),
+        )
+
+    def as_dict(self) -> dict:
+        return {
+            "authorization_endpoint": self.authorization_endpoint,
+            "token_endpoint": self.token_endpoint,
+        }
+
+
+@dataclass
+class Token:
+    access_token: str
+    token_type: Optional[str] = None
+    refresh_token: Optional[str] = None
+    expiry: Optional[datetime] = None
+
+    @property
+    def expired(self):
+        if not self.expiry:
+            return False
+        # Azure Databricks rejects tokens that expire in 30 seconds or less,
+        # so we refresh the token 40 seconds before it expires.
+        potentially_expired = self.expiry - timedelta(seconds=40)
+        now = datetime.now(tz=potentially_expired.tzinfo)
+        is_expired = potentially_expired < now
+        return is_expired
+
+    @property
+    def valid(self):
+        return self.access_token and not self.expired
+
+    def as_dict(self) -> dict:
+        raw = {
+            "access_token": self.access_token,
+            "token_type": self.token_type,
+        }
+        if self.expiry:
+            raw["expiry"] = self.expiry.isoformat()
+        if self.refresh_token:
+            raw["refresh_token"] = self.refresh_token
+        return raw
+
+    @staticmethod
+    def from_dict(raw: dict) -> "Token":
+        return Token(
+            access_token=raw["access_token"],
+            token_type=raw["token_type"],
+            expiry=datetime.fromisoformat(raw["expiry"]),
+            refresh_token=raw.get("refresh_token"),
+        )
+
+    def jwt_claims(self) -> Dict[str, str]:
+        """Get claims from the access token or return an empty dictionary if it is not a JWT token.
+
+        All refreshable tokens we're dealing with are JSON Web Tokens (JWT).
+
+        The common claims are:
+        - 'aud' represents the intended recipient of the token. In case of Azure, this is an app's Application ID
+                assigned within the Azure portal.
+        - 'iss' serves to identify the security token service (STS) responsible for creating and delivering the token.
+                In case of Azure, it includes the Azure AD tenant where user authentication occurred.
+        - 'appid' stands for the application ID of the client utilizing this token. This application can operate either
+                autonomously or on behalf of a user. The application ID commonly represents an application object but
+                may also denote a service principal object in case of Azure.
+        - 'idp' is used to document the identity provider that authenticated the subject of the token.
+        - 'oid' is the unchanging identifier for an entity within the identity system.
+        - 'sub' identifies the primary entity for the token, such as the user of an app. This value is specific to
+                a particular application ID. If a single user logs into two different apps using distinct client IDs,
+                these apps will receive different values for the subject claim.
+        - 'tid' In case of Azure, this value represents Azure Tenant ID.
+
+        See https://datatracker.ietf.org/doc/html/rfc7519 for specification.
+        See https://jwt.ms for debugger.
+        """
+        try:
+            jwt_split = self.access_token.split(".")
+            if len(jwt_split) != 3:
+                logger.debug(f"Tried to decode access token as JWT, but failed: {len(jwt_split)} components")
+                return {}
+            payload_with_padding = jwt_split[1] + "=="
+            payload_bytes = base64.standard_b64decode(payload_with_padding)
+            payload_json = payload_bytes.decode("utf8")
+            claims = json.loads(payload_json)
+            return claims
+        except ValueError as err:
+            logger.debug(f"Tried to decode access token as JWT, but failed: {err}")
+            return {}
+
+
+class TokenSource:
+    @abstractmethod
+    def token(self) -> Token:
+        pass
+
+
+def retrieve_token(
+    client_id,
+    client_secret,
+    token_url,
+    params,
+    use_params=False,
+    use_header=False,
+    headers=None,
+) -> Token:
+    logger.debug(f"Retrieving token for {client_id}")
+    if use_params:
+        if client_id:
+            params["client_id"] = client_id
+        if client_secret:
+            params["client_secret"] = client_secret
+    auth = None
+    if use_header:
+        auth = requests.auth.HTTPBasicAuth(client_id, client_secret)
+    else:
+        auth = IgnoreNetrcAuth()
+    resp = requests.post(token_url, params, auth=auth, headers=headers)
+    if not resp.ok:
+        if resp.headers["Content-Type"].startswith("application/json"):
+            err = resp.json()
+            code = err.get("errorCode", err.get("error", "unknown"))
+            summary = err.get("errorSummary", err.get("error_description", "unknown"))
+            summary = summary.replace("\r\n", " ")
+            raise ValueError(f"{code}: {summary}")
+        raise ValueError(resp.content)
+    try:
+        j = resp.json()
+        expires_in = int(j["expires_in"])
+        expiry = datetime.now() + timedelta(seconds=expires_in)
+        return Token(
+            access_token=j["access_token"],
+            refresh_token=j.get("refresh_token"),
+            token_type=j["token_type"],
+            expiry=expiry,
+        )
+    except Exception as e:
+        raise NotImplementedError(f"Not supported yet: {e}")
+
+
+class _TokenState(Enum):
+    """
+    Represents the state of a token. Each token can be in one of
+    the following three states:
+      - FRESH: The token is valid.
+      - STALE: The token is valid but will expire soon.
+      - EXPIRED: The token has expired and cannot be used.
+    """
+
+    FRESH = 1  # The token is valid.
+    STALE = 2  # The token is valid but will expire soon.
+    EXPIRED = 3  # The token has expired and cannot be used.
+
+
+class Refreshable(TokenSource):
+    """A token source that supports refreshing expired tokens."""
+
+    _EXECUTOR = None
+    _EXECUTOR_LOCK = threading.Lock()
+    # Default duration for the stale period. This value is chosen to cover the
+    # maximum monthly downtime allowed by a 99.99% uptime SLA (~4.38 minutes).
+    _DEFAULT_STALE_DURATION = timedelta(minutes=5)
+
+    @classmethod
+    def _get_executor(cls):
+        """Lazy initialization of the ThreadPoolExecutor."""
+        if cls._EXECUTOR is None:
+            with cls._EXECUTOR_LOCK:
+                if cls._EXECUTOR is None:
+                    # This thread pool has multiple workers because it is shared by all instances of Refreshable.
+                    cls._EXECUTOR = ThreadPoolExecutor(max_workers=10)
+        return cls._EXECUTOR
+
+    def __init__(
+        self,
+        token: Optional[Token] = None,
+        disable_async: bool = True,
+        stale_duration: timedelta = _DEFAULT_STALE_DURATION,
+    ):
+        # Config properties
+        self._stale_duration = stale_duration
+        self._disable_async = disable_async
+        # Lock
+        self._lock = threading.Lock()
+        # Non Thread safe properties. They should be accessed only when protected by the lock above.
+        self._token = token or Token("")
+        self._is_refreshing = False
+        self._refresh_err = False
+
+    # This is the main entry point for the Token. Do not access the token
+    # using any of the internal functions.
+    def token(self) -> Token:
+        """Returns a valid token, blocking if async refresh is disabled."""
+        with self._lock:
+            if self._disable_async:
+                return self._blocking_token()
+            return self._async_token()
+
+    def _async_token(self) -> Token:
+        """
+        Returns a token.
+        If the token is stale, triggers an asynchronous refresh.
+        If the token is expired, refreshes it synchronously, blocking until the refresh is complete.
+        """
+        state = self._token_state()
+        token = self._token
+
+        if state == _TokenState.FRESH:
+            return token
+        if state == _TokenState.STALE:
+            self._trigger_async_refresh()
+            return token
+        return self._blocking_token()
+
+    def _token_state(self) -> _TokenState:
+        """Returns the current state of the token."""
+        if not self._token or not self._token.valid:
+            return _TokenState.EXPIRED
+        if not self._token.expiry:
+            return _TokenState.FRESH
+
+        lifespan = self._token.expiry - datetime.now()
+        if lifespan < timedelta(seconds=0):
+            return _TokenState.EXPIRED
+        if lifespan < self._stale_duration:
+            return _TokenState.STALE
+        return _TokenState.FRESH
+
+    def _blocking_token(self) -> Token:
+        """Returns a token, blocking if necessary to refresh it."""
+        state = self._token_state()
+        # This is important to recover from potential previous failed attempts
+        # to refresh the token asynchronously.
+        self._refresh_err = False
+        self._is_refreshing = False
+
+        # It's possible that the token got refreshed (either by a _blocking_refresh or
+        # an _async_refresh call) while this particular call was waiting to acquire
+        # the lock. This check avoids refreshing the token again in such cases.
+        if state != _TokenState.EXPIRED:
+            return self._token
+
+        self._token = self.refresh()
+        return self._token
+
+    def _trigger_async_refresh(self):
+        """Starts an asynchronous refresh if none is in progress."""
+
+        def _refresh_internal():
+            new_token = None
+            try:
+                new_token = self.refresh()
+            except Exception as e:
+                # This happens on a thread, so we don't want to propagate the error.
+                # Instead, if there is no new_token for any reason, we will disable async refresh below
+                # But we will do it inside the lock.
+                logger.warning(f"Tried to refresh token asynchronously, but failed: {e}")
+
+            with self._lock:
+                if new_token is not None:
+                    self._token = new_token
+                else:
+                    self._refresh_err = True
+                self._is_refreshing = False
+
+        # The token may have been refreshed by another thread.
+        if self._token_state() == _TokenState.FRESH:
+            return
+        if not self._is_refreshing and not self._refresh_err:
+            self._is_refreshing = True
+            Refreshable._get_executor().submit(_refresh_internal)
+
+    @abstractmethod
+    def refresh(self) -> Token:
+        pass
+
+
+class _OAuthCallback(BaseHTTPRequestHandler):
+    def __init__(self, feedback: list, *args):
+        self._feedback = feedback
+        super().__init__(*args)
+
+    def log_message(self, fmt: str, *args: Any) -> None:
+        logger.debug(fmt, *args)
+
+    def do_GET(self):
+        from urllib.parse import parse_qsl
+
+        parts = self.path.split("?")
+        if len(parts) != 2:
+            self.send_error(400, "Missing Query")
+            return
+
+        query = dict(parse_qsl(parts[1]))
+        self._feedback.append(query)
+
+        if "error" in query:
+            self.send_error(400, query["error"], query.get("error_description"))
+            return
+
+        self.send_response(200)
+        self.send_header("Content-type", "text/html")
+        self.end_headers()
+        # TODO: show better message
+        self.wfile.write(b"You can close this tab.")
+
+
+@dataclass
+class HostMetadata:
+    """Parsed response from the /.well-known/databricks-config discovery endpoint."""
+
+    oidc_endpoint: str
+    account_id: Optional[str] = None
+    workspace_id: Optional[str] = None
+
+    @staticmethod
+    def from_dict(d: dict) -> "HostMetadata":
+        return HostMetadata(
+            oidc_endpoint=d.get("oidc_endpoint", ""),
+            account_id=d.get("account_id"),
+            workspace_id=d.get("workspace_id"),
+        )
+
+    def as_dict(self) -> dict:
+        return {
+            "oidc_endpoint": self.oidc_endpoint,
+            "account_id": self.account_id,
+            "workspace_id": self.workspace_id,
+        }
+
+
+def get_host_metadata(host: str, client: _BaseClient = _BaseClient()) -> HostMetadata:
+    """
+    [Experimental] Fetch the raw Databricks well-known configuration from {host}/.well-known/databricks-config.
+
+    :param host: The Databricks host (workspace or account console).
+    :return: Parsed :class:`HostMetadata` as returned by the server.
+    """
+    host = _fix_host_if_needed(host)
+    try:
+        resp = client.do("GET", f"{host}/.well-known/databricks-config")
+    except Exception as e:
+        raise ValueError(f"Failed to fetch host metadata from {host}/.well-known/databricks-config: {e}") from e
+    return HostMetadata.from_dict(resp)
+
+
+def get_endpoints_from_url(url: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
+    """
+    Fetch OIDC endpoints directly from a discovery URL.
+
+    :param url: Full URL to the OIDC discovery document (e.g. the value of discovery_url config).
+    :return: Parsed :class:`OidcEndpoints`.
+    """
+    resp = client.do("GET", url)
+    return OidcEndpoints.from_dict(resp)
+
+
+def get_account_endpoints(host: str, account_id: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
+    """
+    Get the OIDC endpoints for a given account.
+    :param host: The Databricks account host.
+    :param account_id: The account ID.
+    :return: The account's OIDC endpoints.
+    """
+    host = _fix_host_if_needed(host)
+    oidc = f"{host}/oidc/accounts/{account_id}/.well-known/oauth-authorization-server"
+    resp = client.do("GET", oidc)
+    return OidcEndpoints.from_dict(resp)
+
+
+def get_workspace_endpoints(host: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
+    """
+    Get the OIDC endpoints for a given workspace.
+    :param host: The Databricks workspace host.
+    :return: The workspace's OIDC endpoints.
+    """
+    host = _fix_host_if_needed(host)
+    oidc = f"{host}/oidc/.well-known/oauth-authorization-server"
+    resp = client.do("GET", oidc)
+    return OidcEndpoints.from_dict(resp)
+
+
+def get_unified_endpoints(host: str, account_id: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
+    """
+    Get the OIDC endpoints for a unified host.
+    :param host: The Databricks unified host.
+    :param account_id: The account ID.
+    :return: The OIDC endpoints for the unified host.
+    """
+    host = _fix_host_if_needed(host)
+    oidc = f"{host}/oidc/accounts/{account_id}/.well-known/oauth-authorization-server"
+    resp = client.do("GET", oidc)
+    return OidcEndpoints.from_dict(resp)
+
+
+def get_azure_entra_id_workspace_endpoints(
+    host: str,
+) -> Optional[OidcEndpoints]:
+    """
+    Get the Azure Entra ID endpoints for a given workspace. Can only be used when authenticating to Azure Databricks
+    using an application registered in Azure Entra ID.
+    :param host: The Databricks workspace host.
+    :return: The OIDC endpoints for the workspace's Azure Entra ID tenant.
+    """
+    # In Azure, this workspace endpoint redirects to the Entra ID authorization endpoint
+    host = _fix_host_if_needed(host)
+    res = requests.get(f"{host}/oidc/oauth2/v2.0/authorize", allow_redirects=False)
+    real_auth_url = res.headers.get("location")
+    if not real_auth_url:
+        return None
+    return OidcEndpoints(
+        authorization_endpoint=real_auth_url,
+        token_endpoint=real_auth_url.replace("/authorize", "/token"),
+    )
+
+
+class SessionCredentials(Refreshable):
+    def __init__(
+        self,
+        token: Token,
+        token_endpoint: str,
+        client_id: str,
+        client_secret: str = None,
+        redirect_url: str = None,
+        disable_async: bool = True,
+    ):
+        self._token_endpoint = token_endpoint
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._redirect_url = redirect_url
+        super().__init__(
+            token=token,
+            disable_async=disable_async,
+        )
+
+    def as_dict(self) -> dict:
+        return {"token": self.token().as_dict()}
+
+    @staticmethod
+    def from_dict(
+        raw: dict,
+        token_endpoint: str,
+        client_id: str,
+        client_secret: str = None,
+        redirect_url: str = None,
+    ) -> "SessionCredentials":
+        return SessionCredentials(
+            token=Token.from_dict(raw["token"]),
+            token_endpoint=token_endpoint,
+            client_id=client_id,
+            client_secret=client_secret,
+            redirect_url=redirect_url,
+        )
+
+    def auth_type(self):
+        """Implementing CredentialsProvider protocol"""
+        # TODO: distinguish between Databricks IDP and Azure AD
+        return "oauth"
+
+    def __call__(self, *args, **kwargs):
+        """Implementing CredentialsProvider protocol"""
+
+        def inner() -> Dict[str, str]:
+            return {"Authorization": f"Bearer {self.token().access_token}"}
+
+        return inner
+
+    def refresh(self) -> Token:
+        refresh_token = self._token.refresh_token
+        if not refresh_token:
+            raise ValueError("oauth2: token expired and refresh token is not set")
+        params = {
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+        }
+        headers = {}
+        if "microsoft" in self._token_endpoint:
+            # Tokens issued for the 'Single-Page Application' client-type may
+            # only be redeemed via cross-origin requests
+            headers = {"Origin": self._redirect_url}
+        return retrieve_token(
+            client_id=self._client_id,
+            client_secret=self._client_secret,
+            token_url=self._token_endpoint,
+            params=params,
+            use_params=True,
+            headers=headers,
+        )
+
+
+class Consent:
+    def __init__(
+        self,
+        state: str,
+        verifier: str,
+        authorization_url: str,
+        redirect_url: str,
+        token_endpoint: str,
+        client_id: str,
+        client_secret: str = None,
+    ) -> None:
+        self._verifier = verifier
+        self._state = state
+        self._authorization_url = authorization_url
+        self._redirect_url = redirect_url
+        self._token_endpoint = token_endpoint
+        self._client_id = client_id
+        self._client_secret = client_secret
+
+    def as_dict(self) -> dict:
+        return {
+            "state": self._state,
+            "verifier": self._verifier,
+            "authorization_url": self._authorization_url,
+            "redirect_url": self._redirect_url,
+            "token_endpoint": self._token_endpoint,
+            "client_id": self._client_id,
+        }
+
+    @property
+    def authorization_url(self) -> str:
+        return self._authorization_url
+
+    @staticmethod
+    def from_dict(raw: dict, client_secret: str = None) -> "Consent":
+        return Consent(
+            raw["state"],
+            raw["verifier"],
+            authorization_url=raw["authorization_url"],
+            redirect_url=raw["redirect_url"],
+            token_endpoint=raw["token_endpoint"],
+            client_id=raw["client_id"],
+            client_secret=client_secret,
+        )
+
+    def launch_external_browser(self) -> SessionCredentials:
+        redirect_url = urllib.parse.urlparse(self._redirect_url)
+        if redirect_url.hostname not in ("localhost", "127.0.0.1"):
+            raise ValueError(f"cannot listen on {redirect_url.hostname}")
+        feedback = []
+        logger.info(f"Opening {self._authorization_url} in a browser")
+        webbrowser.open_new(self._authorization_url)
+        port = redirect_url.port
+        handler_factory = functools.partial(_OAuthCallback, feedback)
+        with HTTPServer(("localhost", port), handler_factory) as httpd:
+            logger.info(f"Waiting for redirect to http://localhost:{port}")
+            httpd.handle_request()
+        if not feedback:
+            raise ValueError("No data received in callback")
+        query = feedback.pop()
+        return self.exchange_callback_parameters(query)
+
+    def exchange_callback_parameters(self, query: Dict[str, str]) -> SessionCredentials:
+        if "error" in query:
+            raise ValueError("{error}: {error_description}".format(**query))
+        if "code" not in query or "state" not in query:
+            raise ValueError("No code returned in callback")
+        return self.exchange(query["code"], query["state"])
+
+    def exchange(self, code: str, state: str) -> SessionCredentials:
+        if self._state != state:
+            raise ValueError("state mismatch")
+        params = {
+            "redirect_uri": self._redirect_url,
+            "grant_type": "authorization_code",
+            "code_verifier": self._verifier,
+            "code": code,
+        }
+        headers = {}
+        while True:
+            try:
+                token = retrieve_token(
+                    client_id=self._client_id,
+                    client_secret=self._client_secret,
+                    token_url=self._token_endpoint,
+                    params=params,
+                    headers=headers,
+                    use_params=True,
+                )
+                return SessionCredentials(
+                    token,
+                    self._token_endpoint,
+                    self._client_id,
+                    self._client_secret,
+                    self._redirect_url,
+                )
+            except ValueError as e:
+                if NO_ORIGIN_FOR_SPA_CLIENT_ERROR in str(e):
+                    # Retry in cases of 'Single-Page Application' client-type with
+                    # 'Origin' header equal to client's redirect URL.
+                    headers["Origin"] = self._redirect_url
+                    msg = f"Retrying OAuth token exchange with {self._redirect_url} origin"
+                    logger.debug(msg)
+                    continue
+                raise e
+
+
+class OAuthClient:
+    """Enables 3-legged OAuth2 flow with PKCE
+
+    For a regular web app running on a server, it's recommended to use
+    the Authorization Code Flow to obtain an Access Token and a Refresh
+    Token. This method is considered safe because the Access Token is
+    transmitted directly to the server hosting the app, without passing
+    through the user's web browser and risking exposure.
+
+    To enhance the security of the Authorization Code Flow, the PKCE
+    (Proof Key for Code Exchange) mechanism can be employed. With PKCE,
+    the calling application generates a secret called the Code Verifier,
+    which is verified by the authorization server. The app also creates
+    a transform value of the Code Verifier, called the Code Challenge,
+    and sends it over HTTPS to obtain an Authorization Code.
+    By intercepting the Authorization Code, a malicious attacker cannot
+    exchange it for a token without possessing the Code Verifier.
+    """
+
+    def __init__(
+        self,
+        oidc_endpoints: OidcEndpoints,
+        redirect_url: str,
+        client_id: str,
+        scopes: List[str] = None,
+        client_secret: str = None,
+    ):
+        if not scopes:
+            # Default for direct OAuthClient users (e.g., via from_host()).
+            # When used via credentials_provider.external_browser(), scopes are always
+            # passed explicitly from Config.get_scopes(), with offline_access handling
+            # controlled by the disable_oauth_refresh_token flag.
+            scopes = ["all-apis", "offline_access"]
+
+        self.redirect_url = redirect_url
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._oidc_endpoints = oidc_endpoints
+        self._scopes = scopes
+
+    @staticmethod
+    def from_host(
+        host: str,
+        client_id: str,
+        redirect_url: str,
+        *,
+        scopes: List[str] = None,
+        client_secret: str = None,
+    ) -> "OAuthClient":
+        from .core import Config
+        from .credentials_provider import credentials_strategy
+
+        @credentials_strategy("noop", [])
+        def noop_credentials(_: any):
+            return lambda: {}
+
+        config = Config(host=host, credentials_strategy=noop_credentials)
+        oidc = config.databricks_oidc_endpoints
+        if not oidc:
+            raise ValueError(f"{host} does not support OAuth")
+        return OAuthClient(oidc, redirect_url, client_id, scopes, client_secret)
+
+    def initiate_consent(self) -> Consent:
+        state = secrets.token_urlsafe(16)
+
+        # token_urlsafe() already returns base64-encoded string
+        verifier = secrets.token_urlsafe(32)
+        digest = hashlib.sha256(verifier.encode("UTF-8")).digest()
+        challenge = base64.urlsafe_b64encode(digest).decode("UTF-8").replace("=", "")
+
+        params = {
+            "response_type": "code",
+            "client_id": self._client_id,
+            "redirect_uri": self.redirect_url,
+            "scope": " ".join(self._scopes),
+            "state": state,
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+        }
+        auth_url = f"{self._oidc_endpoints.authorization_endpoint}?{urllib.parse.urlencode(params)}"
+        return Consent(
+            state,
+            verifier,
+            authorization_url=auth_url,
+            redirect_url=self.redirect_url,
+            token_endpoint=self._oidc_endpoints.token_endpoint,
+            client_id=self._client_id,
+            client_secret=self._client_secret,
+        )
+
+    def __repr__(self) -> str:
+        return f"<OAuthClient client_id={self._client_id} token_url={self._oidc_endpoints.token_endpoint} auth_url={self._oidc_endpoints.authorization_endpoint}>"
+
+
+@dataclass
+class ClientCredentials(Refreshable):
+    """Enables client credentials 2-legged OAuth2 flow
+
+    When it comes to authorizing machine-to-machine interactions,
+    the need for end-user authorization is eliminated because the SDK
+    functions as both the Resource Owner and Client. Typical example is
+    the CI/CD process or any other automated job. In this scenario,
+    the background job uses the Client ID and Client Secret to obtain
+    an Access Token from the Authorization Server.
+    """
+
+    client_id: str
+    client_secret: str
+    token_url: str
+    endpoint_params: dict = None
+    scopes: str = None
+    use_params: bool = False
+    use_header: bool = False
+    disable_async: bool = True
+    authorization_details: str = None
+
+    def __post_init__(self):
+        super().__init__(disable_async=self.disable_async)
+
+    def refresh(self) -> Token:
+        params = {"grant_type": "client_credentials"}
+        if self.scopes:
+            params["scope"] = self.scopes
+        if self.authorization_details:
+            params["authorization_details"] = self.authorization_details
+        if self.endpoint_params:
+            for k, v in self.endpoint_params.items():
+                params[k] = v
+        return retrieve_token(
+            self.client_id,
+            self.client_secret,
+            self.token_url,
+            params,
+            use_params=self.use_params,
+            use_header=self.use_header,
+        )
+
+
+@dataclass
+class PATOAuthTokenExchange(Refreshable):
+    """Performs OAuth token exchange using a Personal Access Token (PAT) as the subject token.
+
+    This class implements the OAuth 2.0 Token Exchange flow (RFC 8693) to exchange a Databricks
+    Internal PAT Token for an access token with specific scopes and authorization details.
+
+    Args:
+        get_original_token: A callable that returns the PAT to be exchanged. This is a callable
+            rather than a string value to ensure that a fresh Internal PAT Token is retrieved
+            at the time of refresh.
+        host: The Databricks workspace URL (e.g., "https://my-workspace.cloud.databricks.com").
+        scopes: Space-delimited string of OAuth scopes to request (e.g., "all-apis offline_access").
+        authorization_details: Optional JSON string containing authorization details as defined in
+            AuthorizationDetail class above.
+        disable_async: Whether to disable asynchronous token refresh. Defaults to True.
+    """
+
+    get_original_token: Callable[[], Optional[str]]
+    host: str
+    scopes: str
+    authorization_details: str = None
+    disable_async: bool = True
+
+    def __post_init__(self):
+        super().__init__(disable_async=self.disable_async)
+
+    def refresh(self) -> Token:
+        token_exchange_url = f"{self.host}/oidc/v1/token"
+        params = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            "subject_token": self.get_original_token(),
+            "subject_token_type": "urn:databricks:params:oauth:token-type:personal-access-token",
+            "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
+            "scope": self.scopes,
+        }
+        if self.authorization_details:
+            params["authorization_details"] = self.authorization_details
+
+        resp = requests.post(token_exchange_url, params)
+        if not resp.ok:
+            if resp.headers["Content-Type"].startswith("application/json"):
+                err = resp.json()
+                code = err.get("errorCode", err.get("error", "unknown"))
+                summary = err.get("errorSummary", err.get("error_description", "unknown"))
+                summary = summary.replace("\r\n", " ")
+                raise ValueError(f"{code}: {summary}")
+            raise ValueError(resp.content)
+        try:
+            j = resp.json()
+            expires_in = int(j["expires_in"])
+            expiry = datetime.now() + timedelta(seconds=expires_in)
+            return Token(
+                access_token=j["access_token"],
+                expiry=expiry,
+                token_type=j["token_type"],
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to exchange PAT for OAuth token: {e}")
+
+
+class TokenCache:
+    BASE_PATH = "~/.config/databricks-sdk-py/oauth"
+
+    def __init__(
+        self,
+        host: str,
+        oidc_endpoints: OidcEndpoints,
+        client_id: str,
+        redirect_url: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        scopes: Optional[List[str]] = None,
+    ) -> None:
+        self._host = host
+        self._client_id = client_id
+        self._oidc_endpoints = oidc_endpoints
+        self._redirect_url = redirect_url
+        self._client_secret = client_secret
+        self._scopes = scopes or []
+
+    @property
+    def filename(self) -> str:
+        # Include host, client_id, and scopes in the cache filename to make it unique.
+        hash = hashlib.sha256()
+        for chunk in [
+            self._host,
+            self._client_id,
+            ",".join(self._scopes),
+        ]:
+            hash.update(chunk.encode("utf-8"))
+        return os.path.expanduser(os.path.join(self.__class__.BASE_PATH, hash.hexdigest() + ".json"))
+
+    def load(self) -> Optional[SessionCredentials]:
+        """
+        Load credentials from cache file. Return None if the cache file does not exist or is invalid.
+        """
+        if not os.path.exists(self.filename):
+            return None
+
+        try:
+            with open(self.filename, "r") as f:
+                raw = json.load(f)
+                return SessionCredentials.from_dict(
+                    raw,
+                    token_endpoint=self._oidc_endpoints.token_endpoint,
+                    client_id=self._client_id,
+                    client_secret=self._client_secret,
+                    redirect_url=self._redirect_url,
+                )
+        except Exception:
+            return None
+
+    def save(self, credentials: SessionCredentials) -> None:
+        """
+        Save credentials to cache file.
+        """
+        os.makedirs(os.path.dirname(self.filename), exist_ok=True)
+        with open(self.filename, "w") as f:
+            json.dump(credentials.as_dict(), f)
+        os.chmod(self.filename, 0o600)

venv/lib/python3.9/site-packages/databricks/sdk/oidc.py

@@ -0,0 +1,212 @@
+"""
+Package oidc provides utilities for working with OIDC ID tokens.
+
+This package is experimental and subject to change.
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from . import oauth
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class IdToken:
+    """Represents an OIDC ID token that can be exchanged for a Databricks access token.
+
+    Parameters
+    ----------
+    jwt : str
+        The signed JWT token string.
+    """
+
+    jwt: str
+
+
+class IdTokenSource(ABC):
+    """Abstract base class representing anything that returns an IDToken.
+
+    This class defines the interface for token sources that can provide OIDC ID tokens.
+    """
+
+    @abstractmethod
+    def id_token(self) -> IdToken:
+        """Get an ID token.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        Exception
+            Implementation specific exceptions.
+        """
+
+
+class EnvIdTokenSource(IdTokenSource):
+    """IDTokenSource that reads the ID token from an environment variable.
+
+    Parameters
+    ----------
+    env_var : str
+        The name of the environment variable containing the ID token.
+    """
+
+    def __init__(self, env_var: str):
+        self.env_var = env_var
+
+    def id_token(self) -> IdToken:
+        """Get an ID token from an environment variable.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        ValueError
+            If the environment variable is not set.
+        """
+        token = os.getenv(self.env_var)
+        if not token:
+            raise ValueError(f"Missing env var {self.env_var!r}")
+        return IdToken(jwt=token)
+
+
+class FileIdTokenSource(IdTokenSource):
+    """IDTokenSource that reads the ID token from a file.
+
+    Parameters
+    ----------
+    path : str
+        The path to the file containing the ID token.
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+
+    def id_token(self) -> IdToken:
+        """Get an ID token from a file.
+
+        Returns
+        -------
+        IdToken
+            An ID token.
+
+        Raises
+        ------
+        ValueError
+            If the file is empty, does not exist, or cannot be read.
+        """
+        if not self.path:
+            raise ValueError("Missing path")
+
+        token = None
+        try:
+            with open(self.path, "r") as f:
+                token = f.read().strip()
+        except FileNotFoundError:
+            raise ValueError(f"File {self.path!r} does not exist")
+        except Exception as e:
+            raise ValueError(f"Error reading token file: {str(e)}")
+
+        if not token:
+            raise ValueError(f"File {self.path!r} is empty")
+        return IdToken(jwt=token)
+
+
+class DatabricksOidcTokenSource(oauth.TokenSource):
+    """A TokenSource which exchanges a token using Workload Identity Federation.
+
+    Parameters
+    ----------
+    host : str
+        The host of the Databricks account or workspace.
+    id_token_source : IdTokenSource
+        IDTokenSource that returns the IDToken to be used for the token exchange.
+    token_endpoint_provider : Callable[[], dict]
+        Returns the token endpoint for the Databricks OIDC application.
+    client_id : Optional[str], optional
+        ClientID of the Databricks OIDC application. It corresponds to the
+        Application ID of the Databricks Service Principal. Only required for
+        Workload Identity Federation and should be empty for Account-wide token
+        federation.
+    account_id : Optional[str], optional
+        The account ID of the Databricks Account. Only required for
+        Account-wide token federation.
+    audience : Optional[str], optional
+        The audience of the Databricks OIDC application. Only used for
+        Workspace level tokens.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        token_endpoint: str,
+        id_token_source: IdTokenSource,
+        client_id: Optional[str] = None,
+        account_id: Optional[str] = None,
+        audience: Optional[str] = None,
+        disable_async: bool = False,
+        scopes: Optional[str] = None,
+    ):
+        self._host = host
+        self._id_token_source = id_token_source
+        self._token_endpoint = token_endpoint
+        self._client_id = client_id
+        self._account_id = account_id
+        self._audience = audience
+        self._disable_async = disable_async
+        self._scopes = scopes
+
+    def token(self) -> oauth.Token:
+        """Get a token by exchanging the ID token.
+
+        Returns
+        -------
+        dict
+            The exchanged token.
+
+        Raises
+        ------
+        ValueError
+            If the host is missing or other configuration errors occur.
+        """
+        if not self._host:
+            logger.debug("Missing Host")
+            raise ValueError("missing Host")
+
+        if not self._client_id:
+            logger.debug("No ClientID provided, authenticating with Account-wide token federation")
+        else:
+            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="",  # 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.jwt,
+                "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            },
+            scopes=self._scopes,
+            use_params=True,
+            disable_async=self._disable_async,
+        )
+
+        return client.token()

venv/lib/python3.9/site-packages/databricks/sdk/oidc_token_supplier.py

@@ -0,0 +1,108 @@
+import logging
+import os
+from typing import Optional
+
+import requests
+
+logger = logging.getLogger("databricks.sdk")
+
+
+# TODO: Check the required environment variables while creating the instance rather than in the get_oidc_token method to allow early return.
+class GitHubOIDCTokenSupplier:
+    """
+    Supplies OIDC tokens from GitHub Actions.
+    """
+
+    def get_oidc_token(self, audience: str) -> Optional[str]:
+        if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ or "ACTIONS_ID_TOKEN_REQUEST_URL" not in os.environ:
+            # not in GitHub actions
+            return None
+        # See https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers
+        headers = {"Authorization": f"Bearer {os.environ['ACTIONS_ID_TOKEN_REQUEST_TOKEN']}"}
+        endpoint = f"{os.environ['ACTIONS_ID_TOKEN_REQUEST_URL']}&audience={audience}"
+        response = requests.get(endpoint, headers=headers)
+        if not response.ok:
+            return None
+
+        # get the ID Token with aud=api://AzureADTokenExchange sub=repo:org/repo:environment:name
+        response_json = response.json()
+        if "value" not in response_json:
+            return None
+
+        return response_json["value"]
+
+
+class AzureDevOpsOIDCTokenSupplier:
+    """
+    Supplies OIDC tokens from Azure DevOps pipelines.
+
+    Constructs the OIDC token request URL using official Azure DevOps predefined variables.
+    See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables
+    """
+
+    def __init__(self):
+        """Initialize and validate Azure DevOps environment variables."""
+        # Get Azure DevOps environment variables.
+        self.access_token = os.environ.get("SYSTEM_ACCESSTOKEN")
+        self.collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI")
+        self.project_id = os.environ.get("SYSTEM_TEAMPROJECTID")
+        self.plan_id = os.environ.get("SYSTEM_PLANID")
+        self.job_id = os.environ.get("SYSTEM_JOBID")
+        self.hub_name = os.environ.get("SYSTEM_HOSTTYPE")
+
+        # Check for required variables with specific error messages.
+        missing_vars = []
+        if not self.access_token:
+            missing_vars.append("SYSTEM_ACCESSTOKEN")
+        if not self.collection_uri:
+            missing_vars.append("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI")
+        if not self.project_id:
+            missing_vars.append("SYSTEM_TEAMPROJECTID")
+        if not self.plan_id:
+            missing_vars.append("SYSTEM_PLANID")
+        if not self.job_id:
+            missing_vars.append("SYSTEM_JOBID")
+        if not self.hub_name:
+            missing_vars.append("SYSTEM_HOSTTYPE")
+
+        if missing_vars:
+            if "SYSTEM_ACCESSTOKEN" in missing_vars:
+                error_msg = "Azure DevOps OIDC: SYSTEM_ACCESSTOKEN env var not found. If calling from Azure DevOps Pipeline, please set this env var following https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken"
+            else:
+                error_msg = f"Azure DevOps OIDC: missing required environment variables: {', '.join(missing_vars)}"
+            raise ValueError(error_msg)
+
+    def get_oidc_token(self, audience: str) -> Optional[str]:
+        # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange".
+        # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers.
+
+        try:
+            # Construct the OIDC token request URL.
+            # Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken.
+            request_url = f"{self.collection_uri}{self.project_id}/_apis/distributedtask/hubs/{self.hub_name}/plans/{self.plan_id}/jobs/{self.job_id}/oidctoken"
+
+            # Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps).
+            endpoint = f"{request_url}?api-version=7.2-preview.1"
+            headers = {
+                "Authorization": f"Bearer {self.access_token}",
+                "Content-Type": "application/json",
+                "Content-Length": "0",
+            }
+
+            # Azure DevOps OIDC endpoint requires POST request with empty body.
+            response = requests.post(endpoint, headers=headers)
+            if not response.ok:
+                logger.debug(f"Azure DevOps OIDC: token request failed with status {response.status_code}")
+                return None
+
+            # Azure DevOps returns the token in 'oidcToken' field.
+            response_json = response.json()
+            if "oidcToken" not in response_json:
+                logger.debug("Azure DevOps OIDC: response missing 'oidcToken' field")
+                return None
+
+            logger.debug("Azure DevOps OIDC: successfully obtained token")
+            return response_json["oidcToken"]
+        except Exception as e:
+            logger.debug(f"Azure DevOps OIDC: failed to get token: {e}")
+            return None

venv/lib/python3.9/site-packages/databricks/sdk/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561.  The databricks-sdk package uses inline types.

venv/lib/python3.9/site-packages/databricks/sdk/retries.py

@@ -0,0 +1,174 @@
+import functools
+import logging
+from datetime import timedelta
+from random import random, uniform
+from typing import Callable, Optional, Sequence, Tuple, Type, TypeVar
+
+from .clock import Clock, RealClock
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def retried(
+    *,
+    on: Optional[Sequence[Type[BaseException]]] = None,
+    is_retryable: Optional[Callable[[BaseException], Optional[str]]] = None,
+    timeout=timedelta(minutes=20),
+    clock: Optional[Clock] = None,
+    before_retry: Optional[Callable] = None,
+    max_attempts: Optional[int] = 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):
+        raise SyntaxError("either on=[Exception] or callback=lambda x: .. is required")
+    if clock is None:
+        clock = RealClock()
+
+    def decorator(func):
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            deadline = clock.time() + timeout.total_seconds()
+            attempt = 1
+            last_err = None
+            while clock.time() < deadline and (max_attempts is None or attempt <= max_attempts):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as err:
+                    last_err = err
+                    retry_reason = None
+                    # sleep 10s max per attempt, unless it's HTTP 429 or 503
+                    sleep = min(10, attempt)
+                    retry_after_secs = getattr(err, "retry_after_secs", None)
+                    if retry_after_secs is not None:
+                        # cannot depend on DatabricksError directly because of circular dependency
+                        sleep = retry_after_secs
+                        retry_reason = "throttled by platform"
+                    elif is_retryable is not None:
+                        retry_reason = is_retryable(err)
+                    elif on is not None:
+                        for err_type in on:
+                            if not isinstance(err, err_type):
+                                continue
+                            retry_reason = f"{type(err).__name__} is allowed to retry"
+
+                    if retry_reason is None:
+                        # raise if exception is not retryable
+                        raise err
+
+                    logger.debug(f"Retrying: {retry_reason} (sleeping ~{sleep}s)")
+                    if before_retry:
+                        before_retry()
+
+                    clock.sleep(sleep + random())
+                    attempt += 1
+
+            # Determine which limit was hit
+            if max_attempts is not None and attempt > max_attempts:
+                raise RuntimeError(f"Exceeded max retry attempts ({max_attempts})") from last_err
+            raise TimeoutError(f"Timed out after {timeout}") from last_err
+
+        return wrapper
+
+    return decorator
+
+
+class RetryError(Exception):
+    """Error that can be returned from poll functions to control retry behavior."""
+
+    def __init__(self, err: Exception, halt: bool = False):
+        self.err = err
+        self.halt = halt
+        super().__init__(str(err))
+
+    @staticmethod
+    def continues(msg: str) -> "RetryError":
+        """Create a non-halting retry error with a message."""
+        return RetryError(Exception(msg), halt=False)
+
+    @staticmethod
+    def halt(err: Exception) -> "RetryError":
+        """Create a halting retry error."""
+        return RetryError(err, halt=True)
+
+
+def _backoff(attempt: int) -> float:
+    """Calculate backoff time with jitter.
+
+    Linear backoff: attempt * 1 second, capped at 10 seconds
+    Plus random jitter between 50ms and 750ms.
+    """
+    wait = min(10, attempt)
+    jitter = uniform(0.05, 0.75)
+    return wait + jitter
+
+
+def poll(
+    fn: Callable[[], Tuple[Optional[T], Optional[RetryError]]],
+    timeout: Optional[timedelta] = None,
+    clock: Optional[Clock] = None,
+) -> T:
+    """Poll a function until it succeeds or times out.
+
+    The backoff is linear backoff and jitter.
+
+    This function is not meant to be used directly by users.
+    It is used internally by the SDK to poll for the result of an operation.
+    It can be changed in the future without any notice.
+
+    :param fn: Function that returns (result, error).
+               Return (None, RetryError.continues("msg")) to continue polling.
+               Return (None, RetryError.halt(err)) to stop with error.
+               Return (result, None) on success.
+    :param timeout: Maximum time to poll. If None, polls indefinitely.
+    :param clock: Clock implementation for testing (default: RealClock)
+    :returns: The result of the successful function call
+    :raises TimeoutError: If the timeout is reached
+    :raises Exception: If a halting error is encountered
+
+    Example:
+        def check_operation():
+            op = get_operation()
+            if not op.done:
+                return None, RetryError.continues("operation still in progress")
+            if op.error:
+                return None, RetryError.halt(Exception(f"operation failed: {op.error}"))
+            return op.result, None
+
+        result = poll(check_operation, timeout=timedelta(minutes=5))
+    """
+    if clock is None:
+        clock = RealClock()
+
+    deadline = float("inf") if timeout is None else clock.time() + timeout.total_seconds()
+    attempt = 0
+    last_err = None
+
+    while clock.time() < deadline:
+        attempt += 1
+
+        try:
+            result, err = fn()
+
+            if err is None:
+                return result
+
+            if err.halt:
+                raise err.err
+
+            # Continue polling.
+            last_err = err.err
+            wait = _backoff(attempt)
+            logger.debug(f"{str(err.err).rstrip('.')}. Sleeping {wait:.3f}s")
+            clock.sleep(wait)
+
+        except RetryError:
+            raise
+        except Exception as e:
+            # Unexpected error, halt immediately.
+            raise e
+
+    raise TimeoutError(f"Timed out after {timeout}") from last_err

venv/lib/python3.9/site-packages/databricks/sdk/runtime/__init__.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import logging
+from typing import Dict, Optional, Union, cast
+
+logger = logging.getLogger("databricks.sdk")
+is_local_implementation = True
+
+# All objects that are injected into the Notebook's user namespace should also be made
+# available to be imported from databricks.sdk.runtime.globals. This import can be used
+# in Python modules so users can access these objects from Files more easily.
+dbruntime_objects = [
+    "display",
+    "displayHTML",
+    "dbutils",
+    "table",
+    "sql",
+    "udf",
+    "getArgument",
+    "sc",
+    "sqlContext",
+    "spark",
+]
+
+# DO NOT MOVE THE TRY-CATCH BLOCK BELOW AND DO NOT ADD THINGS BEFORE IT! WILL MAKE TEST FAIL.
+try:
+    # We don't want to expose additional entity to user namespace, so
+    # a workaround here for exposing required information in notebook environment
+    from dbruntime.sdk_credential_provider import init_runtime_native_auth
+
+    logger.debug("runtime SDK credential provider available")
+    dbruntime_objects.append("init_runtime_native_auth")
+except ImportError:
+    init_runtime_native_auth = None
+
+globals()["init_runtime_native_auth"] = init_runtime_native_auth
+
+
+def init_runtime_repl_auth():
+    try:
+        from dbruntime.databricks_repl_context import get_context
+
+        ctx = get_context()
+        if ctx is None:
+            logger.debug("Empty REPL context returned, skipping runtime auth")
+            return None, None
+        if ctx.workspaceUrl is None:
+            logger.debug("Workspace URL is not available, skipping runtime auth")
+            return None, None
+        host = f"https://{ctx.workspaceUrl}"
+
+        def inner() -> Dict[str, str]:
+            ctx = get_context()
+            return {"Authorization": f"Bearer {ctx.apiToken}"}
+
+        return host, inner
+    except ImportError:
+        return None, None
+
+
+def init_runtime_legacy_auth():
+    try:
+        import IPython
+
+        ip_shell = IPython.get_ipython()
+        if ip_shell is None:
+            return None, None
+        global_ns = ip_shell.ns_table["user_global"]
+        if "dbutils" not in global_ns:
+            return None, None
+        dbutils = global_ns["dbutils"].notebook.entry_point.getDbutils()
+        if dbutils is None:
+            return None, None
+        ctx = dbutils.notebook().getContext()
+        if ctx is None:
+            return None, None
+        host = getattr(ctx, "apiUrl")().get()
+
+        def inner() -> Dict[str, str]:
+            ctx = dbutils.notebook().getContext()
+            return {"Authorization": f'Bearer {getattr(ctx, "apiToken")().get()}'}
+
+        return host, inner
+    except ImportError:
+        return None, None
+
+
+try:
+    # Internal implementation
+    # Separated from above for backward compatibility
+    from dbruntime import UserNamespaceInitializer
+
+    userNamespaceGlobals = UserNamespaceInitializer.getOrCreate().get_namespace_globals()
+    _globals = globals()
+    for var in dbruntime_objects:
+        if var not in userNamespaceGlobals:
+            continue
+        _globals[var] = userNamespaceGlobals[var]
+    is_local_implementation = False
+except ImportError:
+    # OSS implementation
+    is_local_implementation = True
+
+    for var in dbruntime_objects:
+        globals()[var] = None
+
+    # The next few try-except blocks are for initialising globals in a best effort
+    # mannaer. We separate them to try to get as many of them working as possible
+    try:
+        # We expect this to fail and only do this for providing types
+        from pyspark.sql.context import SQLContext
+
+        sqlContext: SQLContext = None  # type: ignore
+        table = sqlContext.table
+    except Exception as e:
+        logging.debug(f"Failed to initialize globals 'sqlContext' and 'table', continuing. Cause: {e}")
+
+    try:
+        from pyspark.sql.functions import udf  # type: ignore
+    except ImportError as e:
+        logging.debug(f"Failed to initialise udf global: {e}")
+
+    try:
+        from databricks.connect import DatabricksSession  # type: ignore
+
+        spark = DatabricksSession.builder.getOrCreate()
+        sql = spark.sql  # type: ignore
+    except Exception as e:
+        # We are ignoring all failures here because user might want to initialize
+        # spark session themselves and we don't want to interfere with that
+        logging.debug(f"Failed to initialize globals 'spark' and 'sql', continuing. Cause: {e}")
+
+    try:
+        # We expect this to fail locally since dbconnect does not support sparkcontext. This is just for typing
+        sc = spark.sparkContext  # type: ignore
+    except Exception as e:
+        logging.debug(f"Failed to initialize global 'sc', continuing. Cause: {e}")
+
+    def display(input=None, *args, **kwargs) -> None:  # type: ignore
+        """
+        Display plots or data.
+        Display plot:
+                        - display() # no-op
+                        - display(matplotlib.figure.Figure)
+        Display dataset:
+                        - display(spark.DataFrame)
+                        - display(list) # if list can be converted to DataFrame, e.g., list of named tuples
+                        - display(pandas.DataFrame)
+                        - display(koalas.DataFrame)
+                        - display(pyspark.pandas.DataFrame)
+        Display any other value that has a _repr_html_() method
+        For Spark 2.0 and 2.1:
+                        - display(DataFrame, streamName='optional', trigger=optional pyspark.sql.streaming.Trigger,
+                                                        checkpointLocation='optional')
+        For Spark 2.2+:
+                        - display(DataFrame, streamName='optional', trigger=optional interval like '1 second',
+                                                        checkpointLocation='optional')
+        """
+        # Import inside the function so that imports are only triggered on usage.
+        from IPython import display as IPDisplay
+
+        return IPDisplay.display(input, *args, **kwargs)  # type: ignore
+
+    def displayHTML(html) -> None:  # type: ignore
+        """
+        Display HTML data.
+        Parameters
+        ----------
+        data : URL or HTML string
+                        If data is a URL, display the resource at that URL, the resource is loaded dynamically by the browser.
+                        Otherwise data should be the HTML to be displayed.
+        See also:
+        IPython.display.HTML
+        IPython.display.display_html
+        """
+        # Import inside the function so that imports are only triggered on usage.
+        from IPython import display as IPDisplay
+
+        return IPDisplay.display_html(html, raw=True)  # type: ignore
+
+    # We want to propagate the error in initialising dbutils because this is a core
+    # functionality of the sdk
+    from databricks.sdk.dbutils import RemoteDbUtils
+
+    from . import dbutils_stub
+
+    dbutils_type = Union[dbutils_stub.dbutils, RemoteDbUtils]
+
+    dbutils = RemoteDbUtils()
+    dbutils = cast(dbutils_type, dbutils)
+
+    # We do this to prevent importing widgets implementation prematurely
+    # The widget import should prompt users to use the implementation
+    # which has ipywidget support.
+    def getArgument(name: str, defaultValue: Optional[str] = None):
+        return dbutils.widgets.getArgument(name, defaultValue)
+
+
+__all__ = dbruntime_objects

venv/lib/python3.9/site-packages/databricks/sdk/runtime/dbutils_stub.py

@@ -0,0 +1,373 @@
+import typing
+from collections import namedtuple
+
+
+class FileInfo(namedtuple("FileInfo", ["path", "name", "size", "modificationTime"])):
+    pass
+
+
+class MountInfo(namedtuple("MountInfo", ["mountPoint", "source", "encryptionType"])):
+    pass
+
+
+class SecretScope(namedtuple("SecretScope", ["name"])):
+
+    def getName(self):
+        return self.name
+
+
+class SecretMetadata(namedtuple("SecretMetadata", ["key"])):
+    pass
+
+
+class dbutils:
+
+    class credentials:
+        """
+        Utilities for interacting with credentials within notebooks
+        """
+
+        @staticmethod
+        def assumeRole(role: str) -> bool:
+            """
+            Sets the role ARN to assume when looking for credentials to authenticate with S3
+            """
+            ...
+
+        @staticmethod
+        def showCurrentRole() -> typing.List[str]:
+            """
+            Shows the currently set role
+            """
+            ...
+
+        @staticmethod
+        def showRoles() -> typing.List[str]:
+            """
+            Shows the set of possibly assumed roles
+            """
+            ...
+
+        @staticmethod
+        def getCurrentCredentials() -> typing.Mapping[str, str]: ...
+
+    class data:
+        """
+        Utilities for understanding and interacting with datasets (EXPERIMENTAL)
+        """
+
+        @staticmethod
+        def summarize(df: any, precise: bool = False) -> None:
+            """Summarize a Spark/pandas/Koalas DataFrame and visualize the statistics to get quick insights.
+
+            Example: dbutils.data.summarize(df)
+
+            :param df: A pyspark.sql.DataFrame, pyspark.pandas.DataFrame, databricks.koalas.DataFrame
+            or pandas.DataFrame object to summarize. Streaming dataframes are not supported.
+            :param precise: If false, percentiles, distinct item counts, and frequent item counts
+            will be computed approximately to reduce the run time.
+            If true, distinct item counts and frequent item counts will be computed exactly,
+            and percentiles will be computed with high precision.
+
+            :return: visualization of the computed summmary statistics.
+            """
+            ...
+
+    class fs:
+        """
+        Manipulates the Databricks filesystem (DBFS) from the console
+        """
+
+        @staticmethod
+        def cp(source: str, dest: str, recurse: bool = False) -> bool:
+            """
+            Copies a file or directory, possibly across FileSystems
+            """
+            ...
+
+        @staticmethod
+        def head(file: str, max_bytes: int = 65536) -> str:
+            """
+            Returns up to the first 'maxBytes' bytes of the given file as a String encoded in UTF-8
+            """
+            ...
+
+        @staticmethod
+        def ls(path: str) -> typing.List[FileInfo]:
+            """
+            Lists the contents of a directory
+            """
+            ...
+
+        @staticmethod
+        def mkdirs(dir: str) -> bool:
+            """
+            Creates the given directory if it does not exist, also creating any necessary parent directories
+            """
+            ...
+
+        @staticmethod
+        def mv(source: str, dest: str, recurse: bool = False) -> bool:
+            """
+            Moves a file or directory, possibly across FileSystems
+            """
+            ...
+
+        @staticmethod
+        def put(file: str, contents: str, overwrite: bool = False) -> bool:
+            """
+            Writes the given String out to a file, encoded in UTF-8
+            """
+            ...
+
+        @staticmethod
+        def rm(dir: str, recurse: bool = False) -> bool:
+            """
+            Removes a file or directory
+            """
+            ...
+
+        @staticmethod
+        def cacheFiles(*files): ...
+
+        @staticmethod
+        def cacheTable(name: str): ...
+
+        @staticmethod
+        def uncacheFiles(*files): ...
+
+        @staticmethod
+        def uncacheTable(name: str): ...
+
+        @staticmethod
+        def mount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
+            """
+            Mounts the given source directory into DBFS at the given mount point
+            """
+            ...
+
+        @staticmethod
+        def updateMount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
+            """
+            Similar to mount(), but updates an existing mount point (if present) instead of creating a new one
+            """
+            ...
+
+        @staticmethod
+        def mounts() -> typing.List[MountInfo]:
+            """
+            Displays information about what is mounted within DBFS
+            """
+            ...
+
+        @staticmethod
+        def refreshMounts() -> bool:
+            """
+            Forces all machines in this cluster to refresh their mount cache, ensuring they receive the most recent information
+            """
+            ...
+
+        @staticmethod
+        def unmount(mount_point: str) -> bool:
+            """
+            Deletes a DBFS mount point
+            """
+            ...
+
+    class jobs:
+        """
+        Utilities for leveraging jobs features
+        """
+
+        class taskValues:
+            """
+            Provides utilities for leveraging job task values
+            """
+
+            @staticmethod
+            def get(
+                taskKey: str,
+                key: str,
+                default: any = None,
+                debugValue: any = None,
+            ) -> None:
+                """
+                Returns the latest task value that belongs to the current job run
+                """
+                ...
+
+            @staticmethod
+            def set(key: str, value: any) -> None:
+                """
+                Sets a task value on the current task run
+                """
+                ...
+
+    class library:
+        """
+        Utilities for session isolated libraries
+        """
+
+        @staticmethod
+        def restartPython() -> None:
+            """
+            Restart python process for the current notebook session
+            """
+            ...
+
+    class notebook:
+        """
+        Utilities for the control flow of a notebook (EXPERIMENTAL)
+        """
+
+        @staticmethod
+        def exit(value: str) -> None:
+            """
+            This method lets you exit a notebook with a value
+            """
+            ...
+
+        @staticmethod
+        def run(
+            path: str,
+            timeout_seconds: int,
+            arguments: typing.Mapping[str, str],
+        ) -> str:
+            """
+            This method runs a notebook and returns its exit value
+            """
+            ...
+
+    class secrets:
+        """
+        Provides utilities for leveraging secrets within notebooks
+        """
+
+        @staticmethod
+        def get(scope: str, key: str) -> str:
+            """
+            Gets the string representation of a secret value with scope and key
+            """
+            ...
+
+        @staticmethod
+        def getBytes(self, scope: str, key: str) -> bytes:
+            """Gets the bytes representation of a secret value for the specified scope and key."""
+
+        @staticmethod
+        def list(scope: str) -> typing.List[SecretMetadata]:
+            """
+            Lists secret metadata for secrets within a scope
+            """
+            ...
+
+        @staticmethod
+        def listScopes() -> typing.List[SecretScope]:
+            """
+            Lists secret scopes
+            """
+            ...
+
+    class widgets:
+        """
+        provides utilities for working with notebook widgets. You can create different types of widgets and get their bound value
+        """
+
+        @staticmethod
+        def get(name: str) -> str:
+            """Returns the current value of a widget with give name.
+            :param name: Name of the argument to be accessed
+            :return: Current value of the widget or default value
+            """
+            ...
+
+        @staticmethod
+        def getArgument(name: str, defaultValue: typing.Optional[str] = None) -> typing.Optional[str]:
+            """Returns the current value of a widget with give name.
+            :param name: Name of the argument to be accessed
+            :param defaultValue: (Deprecated) default value
+            :return: Current value of the widget or default value
+            """
+            ...
+
+        @staticmethod
+        def text(name: str, defaultValue: str, label: str = None):
+            """Creates a text input widget with given name, default value and optional label for
+            display
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def dropdown(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: str = None,
+        ):
+            """Creates a dropdown input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget (must be one of choices)
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def combobox(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
+            """Creates a combobox input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def multiselect(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
+            """Creates a multiselect input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget (must be one of choices)
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def remove(name: str):
+            """Removes given input widget. If widget does not exist it will throw an error.
+            :param name: Name of argument associated with input widget to be removed
+            """
+            ...
+
+        @staticmethod
+        def removeAll():
+            """Removes all input widgets in the notebook."""
+            ...
+
+
+getArgument = dbutils.widgets.getArgument