diff --git a/python/python/lancedb/remote/errors.py b/python/python/lancedb/remote/errors.py index 76e3b5a8f..2f2dad9b4 100644 --- a/python/python/lancedb/remote/errors.py +++ b/python/python/lancedb/remote/errors.py @@ -27,6 +27,9 @@ class LanceDBClientError(RuntimeError): self.request_id = request_id self.status_code = status_code + def __reduce__(self) -> tuple[type, tuple]: + return (self.__class__, (str(self), self.request_id, self.status_code)) + class HttpError(LanceDBClientError): """An error that occurred during an HTTP request. @@ -101,3 +104,19 @@ class RetryError(LanceDBClientError): self.max_request_failures = max_request_failures self.max_connect_failures = max_connect_failures self.max_read_failures = max_read_failures + + def __reduce__(self) -> tuple[type, tuple]: + return ( + self.__class__, + ( + str(self), + self.request_id, + self.request_failures, + self.connect_failures, + self.read_failures, + self.max_request_failures, + self.max_connect_failures, + self.max_read_failures, + self.status_code, + ), + ) diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py new file mode 100644 index 000000000..338b6a0d7 --- /dev/null +++ b/python/python/tests/test_errors.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pickle + +from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError + + +def test_pickle_lancedb_client_error(): + err = LanceDBClientError("something went wrong", "req-123", 400) + restored = pickle.loads(pickle.dumps(err)) + assert str(restored) == "something went wrong" + assert restored.request_id == "req-123" + assert restored.status_code == 400 + + +def test_pickle_lancedb_client_error_no_status_code(): + err = LanceDBClientError("fail", "req-456") + restored = pickle.loads(pickle.dumps(err)) + assert str(restored) == "fail" + assert restored.request_id == "req-456" + assert restored.status_code is None + + +def test_pickle_http_error(): + err = HttpError("not found", "req-789", 404) + restored = pickle.loads(pickle.dumps(err)) + assert isinstance(restored, HttpError) + assert str(restored) == "not found" + assert restored.request_id == "req-789" + assert restored.status_code == 404 + + +def test_pickle_retry_error(): + err = RetryError( + "max retries exceeded", + "req-abc", + request_failures=3, + connect_failures=1, + read_failures=2, + max_request_failures=5, + max_connect_failures=3, + max_read_failures=3, + status_code=503, + ) + restored = pickle.loads(pickle.dumps(err)) + assert isinstance(restored, RetryError) + assert str(restored) == "max retries exceeded" + assert restored.request_id == "req-abc" + assert restored.request_failures == 3 + assert restored.connect_failures == 1 + assert restored.read_failures == 2 + assert restored.max_request_failures == 5 + assert restored.max_connect_failures == 3 + assert restored.max_read_failures == 3 + assert restored.status_code == 503