generate exceptions

.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,49 @@
+# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
+
+from .base import DatabricksError
+
+__all__ = ['error_mapper'{{range .ExceptionTypes}}, '{{.PascalName}}'{{end}}]
+
+{{$builtinErrors := dict
+    "DEADLINE_EXCEEDED" "TimeoutError"
+    "NOT_FOUND" "LookupError"
+    "UNAUTHENTICATED" "PermissionError"
+    "PERMISSION_DENIED" "PermissionError"
+    "TOO_MANY_REQUESTS" "ResourceWarning"
+    "NOT_IMPLEMENTED" "NotImplementedError"
+}}
+
+{{range .ExceptionTypes}}{{- $builtin := getOrDefault $builtinErrors .Name "" -}}
+  class {{.PascalName}}({{if .Inherit -}}
+    {{.Inherit.PascalName}}
+  {{- else -}}
+    DatabricksError
+  {{- end -}}{{if ne "" $builtin -}}
+    , {{ $builtin }}
+  {{- end}}):
+    """{{.Comment "    " 100 | trimSuffix "\"" }}"""
+{{end}}
+
+_STATUS_CODE_MAPPING = { {{range .ErrorStatusCodeMapping}}
+   {{.StatusCode}}: {{.PascalName}},{{- end}}
+}
+
+_ERROR_CODE_MAPPING = { {{range .ErrorCodeMapping}}
+    '{{.ErrorCode}}': {{.PascalName}},{{- end}}
+}
+
+def error_mapper(status_code: int, raw: dict) -> DatabricksError:
+    error_code = raw.get('error_code', None)
+    if error_code in _ERROR_CODE_MAPPING:
+        # more specific error codes override more generic HTTP status codes
+        return _ERROR_CODE_MAPPING[error_code](**raw)
+
+    if status_code in _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 _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/core.py

@@ -24,7 +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 .errors import DatabricksError, error_mapper
 from .oauth import (ClientCredentials, OAuthClient, OidcEndpoints, Refreshable,
                     Token, TokenCache, TokenSource)
 from .retries import retried
@@ -1117,7 +1117,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 _error_mapper(status_code, 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,3 @@
+from .base import DatabricksError, ErrorDetail
+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]