Merge branch 'main' into update-api-spec

Author: mgyucht

NEXT_CHANGELOG.md

@@ -4,9 +4,10 @@
 
 ### New Features and Improvements
 
-* Add native support for authentication through Azure DevOps OIDC
+* Add native support for authentication through Azure DevOps OIDC.
 
 ### Bug Fixes
+* Fix a security issue that resulted in bearer tokens being logged in exception messages.
 
 ### Documentation
 

databricks/sdk/_base_client.py

@@ -99,7 +99,10 @@ def __init__(
         # Default to 60 seconds
         self._http_timeout_seconds = http_timeout_seconds or 60
 
-        self._error_parser = _Parser(extra_error_customizers=extra_error_customizers)
+        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:

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, the default timeout is 20 minutes.
+        """
+        self.timeout = timeout or timedelta(minutes=20)

databricks/sdk/common/types/__init__.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})"

databricks/sdk/common/types/fieldmask.py

@@ -210,7 +210,11 @@ def __init__(self) -> None:
 class RemoteDbUtils:
 
     def __init__(self, config: "Config" = None):
-        self._config = Config() if not config else config
+        # 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)

databricks/sdk/dbutils.py

@@ -31,12 +31,15 @@
 ]
 
 
-def _unknown_error(response: requests.Response) -> str:
+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=True, debug_truncate_bytes=10 * 1024).generate()
+    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 "
@@ -56,11 +59,13 @@ 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]:
         """
@@ -84,7 +89,7 @@ def get_api_error(self, response: requests.Response) -> Optional[DatabricksError
                     )
             return _error_mapper(
                 response,
-                {"message": "unable to parse response. " + _unknown_error(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

databricks/sdk/errors/parser.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 wrapper(*args, **kwargs):
         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/retries.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/_internal.py

@@ -27,6 +27,7 @@ classifiers = [
 dependencies = [
     "requests>=2.28.1,<3",
     "google-auth~=2.0",
+    "protobuf>=4.21.0,<7.0",
 ]
 
 [project.urls]