databricks-sdk 0.67.0__py3-none-any.whl → 0.69.0__py3-none-any.whl

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

databricks/sdk/oidc_token_supplier.py

@@ -1,9 +1,13 @@
+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.
@@ -26,3 +30,79 @@ class GitHubOIDCTokenSupplier:
             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

databricks/sdk/retries.py

@@ -1,13 +1,15 @@
 import functools
 import logging
 from datetime import timedelta
-from random import random
-from typing import Callable, Optional, Sequence, Type
+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(
     *,
@@ -67,3 +69,101 @@ def retried(
         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: timedelta = timedelta(minutes=20),
+    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 (default: 20 minutes)
+    :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 = 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

databricks/sdk/service/_internal.py

@@ -1,6 +1,11 @@
 import datetime
 import urllib.parse
-from typing import Callable, Dict, Generic, Optional, Type, TypeVar
+from typing import Callable, Dict, Generic, List, Optional, Type, TypeVar
+
+from google.protobuf.duration_pb2 import Duration
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from databricks.sdk.common.types.fieldmask import FieldMask
 
 
 def _from_dict(d: Dict[str, any], field: str, cls: Type) -> any:
@@ -46,6 +51,93 @@ def _escape_multi_segment_path_parameter(param: str) -> str:
     return urllib.parse.quote(param)
 
 
+def _timestamp(d: Dict[str, any], field: str) -> Optional[Timestamp]:
+    """
+    Helper function to convert a timestamp string to a Timestamp object.
+    It takes a dictionary and a field name, and returns a Timestamp object.
+    The field name is the key in the dictionary that contains the timestamp string.
+    """
+    if field not in d or not d[field]:
+        return None
+    ts = Timestamp()
+    ts.FromJsonString(d[field])
+    return ts
+
+
+def _repeated_timestamp(d: Dict[str, any], field: str) -> Optional[List[Timestamp]]:
+    """
+    Helper function to convert a list of timestamp strings to a list of Timestamp objects.
+    It takes a dictionary and a field name, and returns a list of Timestamp objects.
+    The field name is the key in the dictionary that contains the list of timestamp strings.
+    """
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        ts = Timestamp()
+        ts.FromJsonString(v)
+        result.append(ts)
+    return result
+
+
+def _duration(d: Dict[str, any], field: str) -> Optional[Duration]:
+    """
+    Helper function to convert a duration string to a Duration object.
+    It takes a dictionary and a field name, and returns a Duration object.
+    The field name is the key in the dictionary that contains the duration string.
+    """
+    if field not in d or not d[field]:
+        return None
+    dur = Duration()
+    dur.FromJsonString(d[field])
+    return dur
+
+
+def _repeated_duration(d: Dict[str, any], field: str) -> Optional[List[Duration]]:
+    """
+    Helper function to convert a list of duration strings to a list of Duration objects.
+    It takes a dictionary and a field name, and returns a list of Duration objects.
+    The field name is the key in the dictionary that contains the list of duration strings.
+    """
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        dur = Duration()
+        dur.FromJsonString(v)
+        result.append(dur)
+    return result
+
+
+def _fieldmask(d: Dict[str, any], field: str) -> Optional[FieldMask]:
+    """
+    Helper function to convert a fieldmask string to a FieldMask object.
+    It takes a dictionary and a field name, and returns a FieldMask object.
+    The field name is the key in the dictionary that contains the fieldmask string.
+    """
+    if field not in d or not d[field]:
+        return None
+    fm = FieldMask()
+    fm.FromJsonString(d[field])
+    return fm
+
+
+def _repeated_fieldmask(d: Dict[str, any], field: str) -> Optional[List[FieldMask]]:
+    """
+    Helper function to convert a list of fieldmask strings to a list of FieldMask objects.
+    It takes a dictionary and a field name, and returns a list of FieldMask objects.
+    The field name is the key in the dictionary that contains the list of fieldmask strings.
+    """
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        fm = FieldMask()
+        fm.FromJsonString(v)
+        result.append(fm)
+    return result
+
+
 ReturnType = TypeVar("ReturnType")
 
 

databricks/sdk/service/agentbricks.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from ._internal import _enum, _from_dict, _repeated_dict
+from databricks.sdk.service._internal import _enum, _from_dict, _repeated_dict
 
 _LOG = logging.getLogger("databricks.sdk")