Override INVALID_PARAMETER_VALUE on fetching non-existent job/cluster

README.md

@@ -30,6 +30,7 @@ The SDK's internal HTTP client is robust and handles failures on different level
 - [Long-running operations](#long-running-operations)
 - [Paginated responses](#paginated-responses)
 - [Single-sign-on with OAuth](#single-sign-on-sso-with-oauth)
+- [Error handling](#error-handling)
 - [Logging](#logging)
 - [Integration with `dbutils`](#interaction-with-dbutils)
 - [Interface stability](#interface-stability)
@@ -507,6 +508,23 @@ logging.info(f'Created new custom app: '
              f'--client_secret {custom_app.client_secret}')
 ```
 
+## Error handling<a id="error-handling"></a>
+
+The Databricks SDK for Python provides a robust error-handling mechanism that allows developers to catch and handle API errors. When an error occurs, the SDK will raise an exception that contains information about the error, such as the HTTP status code, error message, and error details. Developers can catch these exceptions and handle them appropriately in their code.
+
+```python
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import ResourceDoesNotExist
+
+w = WorkspaceClient()
+try:
+    w.clusters.get(cluster_id='1234-5678-9012')
+except ResourceDoesNotExist as e:
+    print(f'Cluster not found: {e}')
+```
+
+The SDK handles inconsistencies in error responses amongst the different services, providing a consistent interface for developers to work with. Simply catch the appropriate exception type and handle the error as needed.
+
 ## Logging<a id="logging"></a>
 
 The Databricks SDK for Python seamlessly integrates with the standard [Logging facility for Python](https://docs.python.org/3/library/logging.html).

databricks/sdk/core.py

@@ -264,7 +264,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(response, kwargs)
 
     def _record_request_log(self, response: requests.Response, raw=False):
         if not logger.isEnabledFor(logging.DEBUG):

databricks/sdk/errors/mapper.py

@@ -1,8 +1,72 @@
 from databricks.sdk.errors import platform
 from databricks.sdk.errors.base import DatabricksError
+from dataclasses import dataclass
+import re
+from typing import Optional
+import requests
+from .platform import ResourceDoesNotExist
 
 
-def error_mapper(status_code: int, raw: dict) -> DatabricksError:
+@dataclass
+class _ErrorOverride:
+    # The name of the override. Used for logging purposes.
+    debug_name: str
+
+    # A regex that must match the path of the request for this override to be applied.
+    path_regex: re.Pattern
+
+    # The HTTP method of the request for the override to apply
+    verb: str
+
+    # The custom error class to use for this override.
+    error_class: type
+
+    # A regular expression that must match the error code for this override to be applied. If None,
+    # this field is ignored.
+    error_code_regex: Optional[re.Pattern] = None
+
+    # A regular expression that must match the message for this override to be applied. If None,
+    # this field is ignored.
+    message_regex: Optional[re.Pattern] = None
+
+    def matches(self, response: requests.Response, raw_error: dict):
+        if response.request.method != self.verb:
+            return False
+        if not self.path_regex.match(response.request.path_url):
+            return False
+        if self.error_code_regex and not self.error_code_regex.match(raw_error.get('error_code', '')):
+            return False
+        if self.message_regex and not self.message_regex.match(raw_error.get('message', '')):
+            return False
+        return True
+
+
+_INVALID_PARAMETER_VALUE = 'INVALID_PARAMETER_VALUE'
+
+_all_overrides = [
+    _ErrorOverride(
+        debug_name='Clusters InvalidParameterValue => ResourceDoesNotExist',
+        path_regex=re.compile(r'/api/2\.0/clusters/get'),
+        verb='GET',
+        error_code_regex=re.compile(_INVALID_PARAMETER_VALUE),
+        message_regex=re.compile(r'Cluster .+ does not exist'),
+        error_class=ResourceDoesNotExist,
+    ),
+    _ErrorOverride(
+        debug_name='Jobs InvalidParameterValue => ResourceDoesNotExist',
+        path_regex=re.compile(r'/api/2\.\d/jobs/get'),
+        verb='GET',
+        error_code_regex=re.compile(_INVALID_PARAMETER_VALUE),
+        message_regex=re.compile(r'Job .+ does not exist'),
+        error_class=ResourceDoesNotExist,
+    ),
+]
+
+def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
+    for override in _all_overrides:
+        if override.matches(response, raw):
+            return override.error_class(**raw)
+    status_code = response.status_code
     error_code = raw.get('error_code', None)
     if error_code in platform.ERROR_CODE_MAPPING:
         # more specific error codes override more generic HTTP status codes

tests/test_core.py

@@ -22,6 +22,7 @@
                                                  DatabricksCliTokenSource,
                                                  HeaderFactory, databricks_cli)
 from databricks.sdk.environments import Cloud, DatabricksEnvironment
+from databricks.sdk.errors import ResourceDoesNotExist
 from databricks.sdk.service.catalog import PermissionsChange
 from databricks.sdk.service.iam import AccessControlRequest
 from databricks.sdk.version import __version__

tests/test_errors.py

@@ -1,25 +1,31 @@
 import pytest
+import requests
 
 from databricks.sdk import errors
 
+def fake_response(status_code: int) -> requests.Response:
+    resp = requests.Response()
+    resp.status_code = status_code
+    resp.request = requests.Request('GET', 'https://databricks.com/api/2.0/service').prepare()
+    return resp
 
 def test_error_code_has_precedence_over_http_status():
-    err = errors.error_mapper(400, {'error_code': 'INVALID_PARAMETER_VALUE', 'message': 'nope'})
+    err = errors.error_mapper(fake_response(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'})
+    err = errors.error_mapper(fake_response(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'})
+    err = errors.error_mapper(fake_response(417), {'error_code': 'WHOOPS', 'message': 'nope'})
     assert errors.DatabricksError == type(err)
 
 
 def test_missing_error_code():
-    err = errors.error_mapper(522, {'message': 'nope'})
+    err = errors.error_mapper(fake_response(522), {'message': 'nope'})
     assert errors.DatabricksError == type(err)
 
 
@@ -48,6 +54,21 @@ def test_missing_error_code():
                           (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'})
+        raise errors.error_mapper(fake_response(status_code), {'error_code': error_code, 'message': 'nope'})
     except klass:
         return
+
+@pytest.mark.parametrize(
+    'verb, path, status_code, error_code, message, expected_error',
+    [
+        ['GET', '/api/2.0/clusters/get', 400, 'INVALID_PARAMETER_VALUE', 'Cluster abcde does not exist', errors.ResourceDoesNotExist],
+        ['GET', '/api/2.0/jobs/get', 400, 'INVALID_PARAMETER_VALUE', 'Job abcde does not exist', errors.ResourceDoesNotExist],
+        ['GET', '/api/2.1/jobs/get', 400, 'INVALID_PARAMETER_VALUE', 'Job abcde does not exist', errors.ResourceDoesNotExist],
+        ['GET', '/api/2.1/jobs/get', 400, 'INVALID_PARAMETER_VALUE', 'Invalid spark version', errors.InvalidParameterValue],
+    ])
+def test_error_overrides(verb, path, status_code, error_code, message, expected_error):
+    resp = requests.Response()
+    resp.status_code = status_code
+    resp.request = requests.Request(verb, f'https://databricks.com{path}').prepare()
+    with pytest.raises(expected_error):
+        raise errors.error_mapper(resp, {'error_code': error_code, 'message': message})