Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions NEXT_CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
## Release v0.71.0

### New Features and Improvements
* Add a new `_ProtobufErrorDeserializer` for handling Protobuf response errors.

### Bug Fixes

Expand Down
20 changes: 20 additions & 0 deletions databricks/sdk/errors/deserializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,3 +117,23 @@ def deserialize_error(self, response: requests.Response, response_body: bytes) -
}
logging.debug("_HtmlErrorParser: no <pre> tag found in error response")
return None


class _ProtobufErrorDeserializer(_ErrorDeserializer):
"""
Parses errors from the Databricks REST API in Protobuf format.
"""

def deserialize_error(self, response: requests.Response, response_body: bytes) -> Optional[dict]:
try:
from google.rpc import status_pb2

status = status_pb2.Status()
status.ParseFromString(response_body)
return {
"message": status.message,
"error_code": response.status_code,
}
except Exception as e:
logging.debug("_ProtobufErrorParser: unable to parse response as Protobuf", exc_info=e)
return None
4 changes: 3 additions & 1 deletion databricks/sdk/errors/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
from .base import DatabricksError
from .customizer import _ErrorCustomizer, _RetryAfterCustomizer
from .deserializer import (_EmptyDeserializer, _ErrorDeserializer,
_HtmlErrorDeserializer, _StandardErrorDeserializer,
_HtmlErrorDeserializer, _ProtobufErrorDeserializer,
_StandardErrorDeserializer,
_StringErrorDeserializer)
from .mapper import _error_mapper
from .private_link import (_get_private_link_validation_error,
Expand All @@ -21,6 +22,7 @@
_StandardErrorDeserializer(),
_StringErrorDeserializer(),
_HtmlErrorDeserializer(),
_ProtobufErrorDeserializer(),
]

# A list of _ErrorCustomizers that are applied to the error arguments after they are parsed. Customizers can modify the
Expand Down
2 changes: 1 addition & 1 deletion databricks/sdk/logger/round_trip_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def generate(self) -> str:
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 self._raw else self._redacted_dump("> ", request.body))
sb.append("> [raw stream]" if self._raw else self._redacted_dump("> ", str(request.body)))
sb.append(f"< {self._response.status_code} {self._response.reason}")
if self._raw and self._response.headers.get("Content-Type", None) != "application/json":
# Raw streams with `Transfer-Encoding: chunked` do not have `Content-Type` header
Expand Down
92 changes: 92 additions & 0 deletions tests/test_errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import pytest
import requests
from google.rpc import status_pb2

from databricks.sdk import errors
from databricks.sdk.errors import details
Expand Down Expand Up @@ -415,3 +416,94 @@ def test_debug_headers_enabled_shows_headers():
assert "debug-token-12345" in error_message
assert "X-Databricks-Azure-SP-Management-Token" in error_message
assert "debug-azure-token-67890" in error_message


def test_protobuf_error_deserializer_valid_protobuf():
# Create a valid protobuf Status message
status = status_pb2.Status()
status.code = 3 # INVALID_ARGUMENT
status.message = "Invalid parameter provided"
serialized_status = status.SerializeToString()

resp = fake_raw_response(
method="POST",
status_code=400,
response_body=serialized_status,
)

parser = errors._Parser()
error = parser.get_api_error(resp)

assert isinstance(error, errors.BadRequest)
assert str(error) == "Invalid parameter provided"


def test_protobuf_error_deserializer_invalid_protobuf():
# Create a response with invalid protobuf data that should fall through to other parsers
resp = fake_raw_response(
method="POST",
status_code=400,
response_body=b"\x00\x01\x02\x03\x04\x05", # Invalid protobuf
)

parser = errors._Parser()
error = parser.get_api_error(resp)

# Should fall back to the generic error handler
assert isinstance(error, errors.BadRequest)
assert "unable to parse response" in str(error)


def test_protobuf_error_deserializer_empty_message():
# Create a protobuf Status message with empty message
status = status_pb2.Status()
status.code = 5 # NOT_FOUND
status.message = ""
serialized_status = status.SerializeToString()

resp = fake_raw_response(
method="GET",
status_code=404,
response_body=serialized_status,
)

parser = errors._Parser()
error = parser.get_api_error(resp)

assert isinstance(error, errors.NotFound)
assert str(error) == "None"


def test_protobuf_error_deserializer_with_details():
# Create a protobuf Status message with details
status = status_pb2.Status()
status.code = 9 # FAILED_PRECONDITION
status.message = "Resource is in an invalid state"
serialized_status = status.SerializeToString()

resp = fake_raw_response(
method="POST",
status_code=400,
response_body=serialized_status,
)

parser = errors._Parser()
error = parser.get_api_error(resp)

assert isinstance(error, errors.BadRequest)
assert str(error) == "Resource is in an invalid state"


def test_protobuf_error_deserializer_priority():
resp = fake_valid_response(
method="POST",
status_code=400,
error_code="INVALID_REQUEST",
message="Invalid request body",
)

parser = errors._Parser()
error = parser.get_api_error(resp)

assert isinstance(error, errors.BadRequest)
assert str(error) == "Invalid request body"
Loading