[Fix] Handle non-JSON errors gracefully (#741)

## Changes
Some errors returned by the platform are not serialized using JSON (see
databricks/databricks-sdk-go#998 for an
example). They are instead serialized in the form "<ERROR_CODE>:
<MESSAGE>". Today, the SDK cannot parse these error messages well,
resulting in a poor user experience.

This PR adds support for parsing these error messages from the platform
to the SDK. This should reduce bug reports for the SDK with respect to
unexpected response parsing. This PR also refactors the error
deserialization logic somewhat to make it more extensible in the future
for other potential error formats that are not currently handled.

As a side-effect of this change, I've refactored the structure of the
error handling in the Python SDK to more closely reflect how errors are
handled in the Go SDK. This should make maintenance more straightforward
in the future. It also introduces a new error message to the Python SDK
to refer users to our issue tracker when the SDK receives an error
response that it cannot parse, like what we do in the Go SDK.

Ports databricks/databricks-sdk-go#1031 to the
Python SDK.

## Deprecations
This PR deprecates several fields in the constructor for
DatabricksError. Going forward, SCIM-specific and API 1.2-specific
parameters should not be specified in the constructor; instead, they
will be handled in error parsers.

## Breaking Changes
The introduction of a different message for non-JSON responses may be a
breaking change if users matched on the message structure used before.

## Tests
Existing tests still pass, adding tests before merging this.

- [ ] `make test` run locally
- [ ] `make fmt` applied
- [ ] relevant integration tests applied

Author: mgyucht

databricks/sdk/core.py

@@ -1,7 +1,5 @@
 import re
-import urllib.parse
 from datetime import timedelta
-from json import JSONDecodeError
 from types import TracebackType
 from typing import Any, BinaryIO, Iterator, Type
 from urllib.parse import urlencode
@@ -12,8 +10,8 @@
 from .config import *
 # To preserve backwards compatibility (as these definitions were previously in this module)
 from .credentials_provider import *
-from .errors import DatabricksError, error_mapper
-from .errors.private_link import _is_private_link_redirect
+from .errors import DatabricksError, get_api_error
+from .logger import RoundTrip
 from .oauth import retrieve_token
 from .retries import retried
 
@@ -262,134 +260,23 @@ def _perform(self,
                                          auth=auth,
                                          stream=raw,
                                          timeout=self._http_timeout_seconds)
-        try:
-            self._record_request_log(response, raw=raw or data is not None or files is not None)
-            if not response.ok: # internally calls response.raise_for_status()
-                # TODO: experiment with traceback pruning for better readability
-                # See https://stackoverflow.com/a/58821552/277035
-                payload = response.json()
-                raise self._make_nicer_error(response=response, **payload) from None
-            # 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):
-                raise self._make_nicer_error(response=response) from None
-            return response
-        except requests.exceptions.JSONDecodeError:
-            message = self._make_sense_from_html(response.text)
-            if not message:
-                message = response.reason
-            raise self._make_nicer_error(response=response, message=message) from None
-
-    @staticmethod
-    def _make_sense_from_html(txt: str) -> str:
-        matchers = [r'<pre>(.*)</pre>', r'<title>(.*)</title>']
-        for attempt in matchers:
-            expr = re.compile(attempt, re.MULTILINE)
-            match = expr.search(txt)
-            if not match:
-                continue
-            return match.group(1).strip()
-        return txt
-
-    def _make_nicer_error(self, *, response: requests.Response, **kwargs) -> DatabricksError:
-        status_code = response.status_code
-        message = kwargs.get('message', 'request failed')
-        is_http_unauthorized_or_forbidden = status_code in (401, 403)
-        is_too_many_requests_or_unavailable = status_code in (429, 503)
-        if is_http_unauthorized_or_forbidden:
-            message = self._cfg.wrap_debug_info(message)
-        if is_too_many_requests_or_unavailable:
-            kwargs['retry_after_secs'] = self._parse_retry_after(response)
-        kwargs['message'] = message
-        return error_mapper(response, kwargs)
-
-    def _record_request_log(self, response: requests.Response, raw=False):
+        self._record_request_log(response, raw=raw or data is not None or files is not None)
+        error = get_api_error(response)
+        if error is not None:
+            status_code = response.status_code
+            is_http_unauthorized_or_forbidden = status_code in (401, 403)
+            is_too_many_requests_or_unavailable = status_code in (429, 503)
+            if is_http_unauthorized_or_forbidden:
+                error.message = self._cfg.wrap_debug_info(error.message)
+            if is_too_many_requests_or_unavailable:
+                error.retry_after_secs = self._parse_retry_after(response)
+            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
-        request = 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._cfg.debug_headers:
-            if self._cfg.host:
-                sb.append(f'> * Host: {self._cfg.host}')
-            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 raw else self._redacted_dump("> ", request.body))
-        sb.append(f'< {response.status_code} {response.reason}')
-        if raw and 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 response.content:
-            sb.append(self._redacted_dump("< ", response.content))
-        logger.debug("\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 JSONDecodeError:
-            return f'{prefix}[non-JSON document of {len(body)} bytes]'
+        logger.debug(RoundTrip(response, self._cfg.debug_headers, self._debug_truncate_bytes, raw).generate())
 
 
 class StreamingResponse(BinaryIO):

databricks/sdk/errors/__init__.py

@@ -1,5 +1,6 @@
 from .base import DatabricksError, ErrorDetail
-from .mapper import error_mapper
+from .mapper import _error_mapper
+from .parser import get_api_error
 from .platform import *
 from .private_link import PrivateLinkValidationError
 from .sdk import *

databricks/sdk/errors/base.py

@@ -1,4 +1,5 @@
 import re
+import warnings
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
@@ -41,9 +42,38 @@ def __init__(self,
                  retry_after_secs: int = None,
                  details: 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:
+        :param details:
+        :param kwargs:
+        """
+        # SCIM-specific parameters are deprecated
+        if detail:
+            warnings.warn(
+                "The 'detail' parameter of DatabricksError is deprecated and will be removed in a future version."
+            )
+        if scimType:
+            warnings.warn(
+                "The 'scimType' parameter of DatabricksError is deprecated and will be removed in a future version."
+            )
+        if status:
+            warnings.warn(
+                "The 'status' parameter of DatabricksError is deprecated and will be removed in a future version."
+            )
+
+        # API 1.2-specific parameters are deprecated
         if error:
-            # API 1.2 has different response format, let's adapt
-            message = error
+            warnings.warn(
+                "The 'error' parameter of DatabricksError is deprecated and will be removed in a future version."
+            )
+
         if detail:
             # Handle SCIM error message details
             # @see https://tools.ietf.org/html/rfc7644#section-3.7.3

databricks/sdk/errors/mapper.py

@@ -4,11 +4,9 @@
 from databricks.sdk.errors.base import DatabricksError
 
 from .overrides import _ALL_OVERRIDES
-from .private_link import (_get_private_link_validation_error,
-                           _is_private_link_redirect)
 
 
-def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
+def _error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
     for override in _ALL_OVERRIDES:
         if override.matches(response, raw):
             return override.custom_error(**raw)
@@ -23,8 +21,6 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
         # 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)
-    if _is_private_link_redirect(response):
-        return _get_private_link_validation_error(response.url)
 
     # backwards-compatible error creation for cases like using older versions of
     # the SDK on way never releases of the platform.

databricks/sdk/errors/parser.py

@@ -0,0 +1,146 @@
+import abc
+import json
+import logging
+import re
+from typing import Optional
+
+import requests
+
+from ..logger import RoundTrip
+from .base import DatabricksError
+from .mapper import _error_mapper
+from .private_link import (_get_private_link_validation_error,
+                           _is_private_link_redirect)
+
+
+class _ErrorParser(abc.ABC):
+    """A parser for errors from the Databricks REST API."""
+
+    @abc.abstractmethod
+    def parse_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 _EmptyParser(_ErrorParser):
+    """A parser that handles empty responses."""
+
+    def parse_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        if len(response_body) == 0:
+            return {'message': response.reason}
+        return None
+
+
+class _StandardErrorParser(_ErrorParser):
+    """
+    Parses errors from the Databricks REST API using the standard error format.
+    """
+
+    def parse_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
+        try:
+            payload_str = response_body.decode('utf-8')
+            resp: dict = json.loads(payload_str)
+        except json.JSONDecodeError as e:
+            logging.debug('_StandardErrorParser: unable to deserialize response as json', exc_info=e)
+            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
+            error_args[
+                'message'] = f"{scim_type} {error_args.get('message', 'SCIM API Internal Error')}".strip(" ")
+            error_args['error_code'] = f"SCIM_{status}"
+        return error_args
+
+
+class _StringErrorParser(_ErrorParser):
+    """
+    Parses errors from the Databricks REST API in the format "ERROR_CODE: MESSAGE".
+    """
+
+    __STRING_ERROR_REGEX = re.compile(r'([A-Z_]+): (.*)')
+
+    def parse_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 _HtmlErrorParser(_ErrorParser):
+    """
+    Parses errors from the Databricks REST API in HTML format.
+    """
+
+    __HTML_ERROR_REGEXES = [re.compile(r'<pre>(.*)</pre>'), re.compile(r'<title>(.*)</title>'), ]
+
+    def parse_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
+
+
+# A list of ErrorParsers that are tried in order to parse an API error from a response body. Most errors should be
+# parsable by the _StandardErrorParser, 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_parsers = [_EmptyParser(), _StandardErrorParser(), _StringErrorParser(), _HtmlErrorParser(), ]
+
+
+def _unknown_error(response: requests.Response) -> 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.
+    """
+    request_log = RoundTrip(response, debug_headers=True, 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}```')
+
+
+def get_api_error(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 _error_parsers:
+            try:
+                error_args = parser.parse_error(response, content)
+                if 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)})
+
+    # 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)