fix errors

Signed-off-by: Serena Ruan <[email protected]>

Author: serena-ruan

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.71.0
 
 ### New Features and Improvements
+* Add a new `_ProtobufErrorDeserializer` for handling Protobuf response errors.
 
 ### Bug Fixes
 

databricks/sdk/errors/deserializer.py

@@ -117,3 +117,22 @@ 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

databricks/sdk/errors/parser.py

@@ -8,7 +8,7 @@
 from .customizer import _ErrorCustomizer, _RetryAfterCustomizer
 from .deserializer import (_EmptyDeserializer, _ErrorDeserializer,
                            _HtmlErrorDeserializer, _StandardErrorDeserializer,
-                           _StringErrorDeserializer)
+                           _StringErrorDeserializer, _ProtobufErrorDeserializer)
 from .mapper import _error_mapper
 from .private_link import (_get_private_link_validation_error,
                            _is_private_link_redirect)
@@ -21,6 +21,7 @@
     _StandardErrorDeserializer(),
     _StringErrorDeserializer(),
     _HtmlErrorDeserializer(),
+    _ProtobufErrorDeserializer(),
 ]
 
 # A list of _ErrorCustomizers that are applied to the error arguments after they are parsed. Customizers can modify the

databricks/sdk/logger/round_trip_logger.py

@@ -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

tests/test_errors.py

@@ -8,6 +8,7 @@
 
 from databricks.sdk import errors
 from databricks.sdk.errors import details
+from google.rpc import status_pb2
 
 
 def fake_response(
@@ -415,3 +416,93 @@ 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"