star

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"]