Introduce more specific exceptions, like NotFound, AlreadyExists,…

… `BadRequest`, `PermissionDenied`, `InternalError`, and others (#376)

Improve the ergonomics of SDK, where instead of `except DatabricksError
as err: if err.error_code != 'NOT_FOUND': raise err else: do_stuff()` we
could do `except NotFound: do_stuff()`, like in [this
example](https://github.com/databrickslabs/ucx/blob/main/src/databricks/labs/ucx/workspace_access/generic.py#L71-L84).

Additionally, it'll make it easier to read stack traces, as they will
contain specific exception class name. Examples of unclear stack traces
are: databrickslabs/ucx#359,
databrickslabs/ucx#353,
databrickslabs/ucx#347,

# First principles
- ~~do not override `builtins.NotImplemented` for `NOT_IMPLEMENTED`
error code~~
- assume that platform error_code/HTTP status code mapping is not
perfect and in the state of transition
- we do represent reasonable subset of error codes as specific
exceptions
- it's still possible to access `error_code` from specific exceptions
like `NotFound` or `AlreadyExists`.

# Proposal
- have hierarchical exceptions, also inheriting from Python's built-in
exceptions
- more specific error codes override more generic HTTP status codes
- more generic HTTP status codes matched after more specific error
codes, where there's a default exception class per HTTP status code, and
we do rely on Databricks platform exception mapper to do the right
thing.
- have backward-compatible error creation for cases like using older
versions of the SDK on the way never releases of the platform.


![image](https://github.com/databricks/databricks-sdk-py/assets/259697/a4519f76-0778-468c-9bf5-6133984b5af7)

### Naming conflict resolution

We have four sources of naming and this is a proposed order of naming
conflict resolution:
1. Databricks `error_code`, that is surfaced in our API documentation,
known by Databricks users
2. HTTP Status codes, known by some developers
3. Python builtin exceptions, known by some developers
4. grpc error codes
https://github.com/googleapis/googleapis/blob/master/google/rpc/code.proto#L38-L185,
know by some developers

---------

Co-authored-by: Miles Yucht <miles@databricks.com>

.codegen.json

@@ -4,7 +4,8 @@
     ".codegen/service.py.tmpl": "databricks/sdk/service/{{.Name}}.py"
   },
   "batch": {
-    ".codegen/__init__.py.tmpl": "databricks/sdk/__init__.py"
+    ".codegen/__init__.py.tmpl": "databricks/sdk/__init__.py",
+    ".codegen/error_mapping.py.tmpl": "databricks/sdk/errors/mapping.py"
   },
   "samples": {
     ".codegen/example.py.tmpl": "examples/{{.Service.SnakeName}}/{{.Method.SnakeName}}_{{.SnakeName}}.py"

.codegen/error_mapping.py.tmpl

@@ -0,0 +1,20 @@
+# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
+
+from .base import DatabricksError
+
+{{range .ExceptionTypes}}
+class {{.PascalName}}({{if .Inherit -}}
+    {{.Inherit.PascalName}}
+  {{- else -}}
+    DatabricksError
+  {{- end -}}):
+  """{{.Comment "    " 100 | trimSuffix "\"" }}"""
+{{end}}
+
+STATUS_CODE_MAPPING = { {{range .ErrorStatusCodeMapping}}
+   {{.StatusCode}}: {{.PascalName}},{{- end}}
+}
+
+ERROR_CODE_MAPPING = { {{range .ErrorCodeMapping}}
+    '{{.ErrorCode}}': {{.PascalName}},{{- end}}
+}

databricks/sdk/core.py

@@ -24,6 +24,7 @@
 
 from .azure import (ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment,
                     add_sp_management_token, add_workspace_id_header)
+from .errors import DatabricksError, error_mapper
 from .oauth import (ClientCredentials, OAuthClient, OidcEndpoints, Refreshable,
                     Token, TokenCache, TokenSource)
 from .retries import retried
@@ -963,70 +964,6 @@ def copy(self):
         return cpy
 
 
-class ErrorDetail:
-
-    def __init__(self,
-                 type: str = None,
-                 reason: str = None,
-                 domain: str = None,
-                 metadata: dict = None,
-                 **kwargs):
-        self.type = type
-        self.reason = reason
-        self.domain = domain
-        self.metadata = metadata
-
-    @classmethod
-    def from_dict(cls, d: Dict[str, any]) -> 'ErrorDetail':
-        if '@type' in d:
-            d['type'] = d['@type']
-        return cls(**d)
-
-
-class DatabricksError(IOError):
-    """ Generic error from Databricks REST API """
-    # Known ErrorDetail types
-    _error_info_type = "type.googleapis.com/google.rpc.ErrorInfo"
-
-    def __init__(self,
-                 message: str = None,
-                 *,
-                 error_code: str = None,
-                 detail: str = None,
-                 status: str = None,
-                 scimType: str = None,
-                 error: str = None,
-                 retry_after_secs: int = None,
-                 details: List[Dict[str, any]] = None,
-                 **kwargs):
-        if error:
-            # API 1.2 has different response format, let's adapt
-            message = error
-        if detail:
-            # Handle SCIM error message details
-            # @see https://tools.ietf.org/html/rfc7644#section-3.7.3
-            if detail == "null":
-                message = "SCIM API Internal Error"
-            else:
-                message = detail
-            # add more context from SCIM responses
-            message = f"{scimType} {message}".strip(" ")
-            error_code = f"SCIM_{status}"
-        super().__init__(message if message else error)
-        self.error_code = error_code
-        self.retry_after_secs = retry_after_secs
-        self.details = [ErrorDetail.from_dict(detail) for detail in details] if details else []
-        self.kwargs = kwargs
-
-    def get_error_info(self) -> List[ErrorDetail]:
-        return self._get_details_by_type(DatabricksError._error_info_type)
-
-    def _get_details_by_type(self, error_type) -> List[ErrorDetail]:
-        if self.details == None:
-            return []
-        return [detail for detail in self.details if detail.type == error_type]
-
-
 class ApiClient:
     _cfg: Config
     _RETRY_AFTER_DEFAULT: int = 1
@@ -1255,7 +1192,7 @@ def _make_nicer_error(self, *, response: requests.Response, **kwargs) -> Databri
         if is_too_many_requests_or_unavailable:
             kwargs['retry_after_secs'] = self._parse_retry_after(response)
         kwargs['message'] = message
-        return DatabricksError(**kwargs)
+        return error_mapper(status_code, kwargs)
 
     def _record_request_log(self, response: requests.Response, raw=False):
         if not logger.isEnabledFor(logging.DEBUG):

databricks/sdk/errors/__init__.py

@@ -0,0 +1,4 @@
+from .base import DatabricksError, ErrorDetail
+from .mapper import error_mapper
+from .mapping import *
+from .sdk import *

databricks/sdk/errors/base.py

@@ -0,0 +1,65 @@
+from typing import Dict, List
+
+
+class ErrorDetail:
+
+    def __init__(self,
+                 type: str = None,
+                 reason: str = None,
+                 domain: str = None,
+                 metadata: dict = None,
+                 **kwargs):
+        self.type = type
+        self.reason = reason
+        self.domain = domain
+        self.metadata = metadata
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> 'ErrorDetail':
+        if '@type' in d:
+            d['type'] = d['@type']
+        return cls(**d)
+
+
+class DatabricksError(IOError):
+    """ Generic error from Databricks REST API """
+    # Known ErrorDetail types
+    _error_info_type = "type.googleapis.com/google.rpc.ErrorInfo"
+
+    def __init__(self,
+                 message: str = None,
+                 *,
+                 error_code: str = None,
+                 detail: str = None,
+                 status: str = None,
+                 scimType: str = None,
+                 error: str = None,
+                 retry_after_secs: int = None,
+                 details: List[Dict[str, any]] = None,
+                 **kwargs):
+        if error:
+            # API 1.2 has different response format, let's adapt
+            message = error
+        if detail:
+            # Handle SCIM error message details
+            # @see https://tools.ietf.org/html/rfc7644#section-3.7.3
+            if detail == "null":
+                message = "SCIM API Internal Error"
+            else:
+                message = detail
+            # add more context from SCIM responses
+            message = f"{scimType} {message}".strip(" ")
+            error_code = f"SCIM_{status}"
+        super().__init__(message if message else error)
+        self.error_code = error_code
+        self.retry_after_secs = retry_after_secs
+        self.details = [ErrorDetail.from_dict(detail) for detail in details] if details else []
+        self.kwargs = kwargs
+
+    def get_error_info(self) -> List[ErrorDetail]:
+        return self._get_details_by_type(DatabricksError._error_info_type)
+
+    def _get_details_by_type(self, error_type) -> List[ErrorDetail]:
+        if self.details == None:
+            return []
+        return [detail for detail in self.details if detail.type == error_type]

databricks/sdk/errors/mapper.py

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

databricks/sdk/errors/mapping.py

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

databricks/sdk/errors/sdk.py

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

tests/test_errors.py

@@ -0,0 +1,53 @@
+import pytest
+
+from databricks.sdk import errors
+
+
+def test_error_code_has_precedence_over_http_status():
+    err = errors.error_mapper(400, {'error_code': 'INVALID_PARAMETER_VALUE', 'message': 'nope'})
+    assert errors.InvalidParameterValue == type(err)
+
+
+def test_http_status_code_maps_fine():
+    err = errors.error_mapper(400, {'error_code': 'MALFORMED_REQUEST', 'message': 'nope'})
+    assert errors.BadRequest == type(err)
+
+
+def test_other_errors_also_map_fine():
+    err = errors.error_mapper(417, {'error_code': 'WHOOPS', 'message': 'nope'})
+    assert errors.DatabricksError == type(err)
+
+
+def test_missing_error_code():
+    err = errors.error_mapper(522, {'message': 'nope'})
+    assert errors.DatabricksError == type(err)
+
+
+@pytest.mark.parametrize('status_code, error_code, klass',
+                         [(400, ..., errors.BadRequest), (400, 'INVALID_PARAMETER_VALUE', errors.BadRequest),
+                          (400, 'INVALID_PARAMETER_VALUE', errors.InvalidParameterValue),
+                          (400, 'REQUEST_LIMIT_EXCEEDED', errors.TooManyRequests), (400, ..., IOError),
+                          (401, ..., errors.Unauthenticated), (401, ..., IOError),
+                          (403, ..., errors.PermissionDenied),
+                          (403, ..., IOError), (404, ..., errors.NotFound), (404, ..., IOError),
+                          (409, ..., errors.ResourceConflict), (409, 'ABORTED', errors.Aborted),
+                          (409, 'ABORTED', errors.ResourceConflict),
+                          (409, 'ALREADY_EXISTS', errors.AlreadyExists),
+                          (409, 'ALREADY_EXISTS', errors.ResourceConflict), (409, ..., IOError),
+                          (429, ..., errors.TooManyRequests),
+                          (429, 'REQUEST_LIMIT_EXCEEDED', errors.TooManyRequests),
+                          (429, 'REQUEST_LIMIT_EXCEEDED', errors.RequestLimitExceeded),
+                          (429, 'RESOURCE_EXHAUSTED', errors.TooManyRequests),
+                          (429, 'RESOURCE_EXHAUSTED', errors.ResourceExhausted), (429, ..., IOError),
+                          (499, ..., errors.Cancelled), (499, ..., IOError), (500, ..., errors.InternalError),
+                          (500, 'UNKNOWN', errors.InternalError), (500, 'UNKNOWN', errors.Unknown),
+                          (500, 'DATA_LOSS', errors.InternalError), (500, 'DATA_LOSS', errors.DataLoss),
+                          (500, ..., IOError), (501, ..., errors.NotImplemented), (501, ..., IOError),
+                          (503, ..., errors.TemporarilyUnavailable), (503, ..., IOError),
+                          (504, ..., errors.DeadlineExceeded), (504, ..., IOError),
+                          (444, ..., errors.DatabricksError), (444, ..., IOError), ])
+def test_subclasses(status_code, error_code, klass):
+    try:
+        raise errors.error_mapper(status_code, {'error_code': error_code, 'message': 'nope'})
+    except klass:
+        return