diff --git a/packages/google-api-core/google/api_core/exceptions.py b/packages/google-api-core/google/api_core/exceptions.py index aa9898c5b9a6..7235da756d30 100644 --- a/packages/google-api-core/google/api_core/exceptions.py +++ b/packages/google-api-core/google/api_core/exceptions.py @@ -93,6 +93,12 @@ def cause(self): def __str__(self): return "{}, last exception: {}".format(self.message, self.cause) + def __reduce__(self): + # ``cause`` is a required constructor argument that isn't stored in + # ``self.args``, so the default ``BaseException.__reduce__`` can't + # rebuild the exception (e.g. when it's pickled across processes). + return (self.__class__, (self.message, self._cause), self.__dict__) + class _GoogleAPICallErrorMeta(type): """Metaclass for registering GoogleAPICallError subclasses.""" diff --git a/packages/google-api-core/tests/unit/test_exceptions.py b/packages/google-api-core/tests/unit/test_exceptions.py index 6a71ddeb9261..04eaefbc0e7b 100644 --- a/packages/google-api-core/tests/unit/test_exceptions.py +++ b/packages/google-api-core/tests/unit/test_exceptions.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import http.client import json +import pickle from unittest import mock import pytest @@ -54,6 +56,50 @@ def test_create_google_cloud_error_with_args(): assert exception.response == response +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_retry_error_pickle_round_trip(protocol): + cause = exceptions.ServiceUnavailable("backend unavailable") + error = exceptions.RetryError("Timeout of 1.0s exceeded", cause) + + restored = pickle.loads(pickle.dumps(error, protocol=protocol)) + + assert isinstance(restored, exceptions.RetryError) + assert restored.message == error.message + assert isinstance(restored.cause, exceptions.ServiceUnavailable) + assert restored.cause.message == cause.message + assert str(restored) == str(error) + + +def test_retry_error_pickle_round_trip_without_cause(): + error = exceptions.RetryError("Timeout of 1.0s exceeded", None) + + restored = pickle.loads(pickle.dumps(error)) + + assert restored.cause is None + assert str(restored) == str(error) + + +def test_retry_error_pickle_preserves_instance_attributes(): + error = exceptions.RetryError("Timeout of 1.0s exceeded", None) + error.custom_attribute = "value" + + restored = pickle.loads(pickle.dumps(error)) + + assert restored.custom_attribute == "value" + + +@pytest.mark.parametrize("copy_fn", [copy.copy, copy.deepcopy]) +def test_retry_error_copy(copy_fn): + cause = exceptions.ServiceUnavailable("backend unavailable") + error = exceptions.RetryError("Timeout of 1.0s exceeded", cause) + + copied = copy_fn(error) + + assert isinstance(copied, exceptions.RetryError) + assert copied.message == error.message + assert str(copied) == str(error) + + def test_from_http_status(): message = "message" exception = exceptions.from_http_status(http.client.NOT_FOUND, message)