fix: map LRO errors to library exception types (#86)

Fixes #15 🦕

Errors raised by long running operations are currently always type GoogleAPICallError. Use the status code to create a more specific exception type.

4 files changed, 45 insertions, 3 deletions

diff --git a/google/api_core/exceptions.py b/google/api_core/exceptions.py
index b9c46ca..a29266d 100644
--- a/google/api_core/exceptions.py
+++ b/google/api_core/exceptions.py
@@ -26,6 +26,7 @@ from six.moves import http_client
 
 try:
     import grpc
+
 except ImportError:  # pragma: NO COVER
     grpc = None
 
@@ -34,6 +35,14 @@ except ImportError:  # pragma: NO COVER
 _HTTP_CODE_TO_EXCEPTION = {}
 _GRPC_CODE_TO_EXCEPTION = {}
 
+# Additional lookup table to map integer status codes to grpc status code
+# grpc does not currently support initializing enums from ints
+# i.e., grpc.StatusCode(5) raises an error
+_INT_TO_GRPC_CODE = {}
+if grpc is not None:  # pragma: no branch
+    for x in grpc.StatusCode:
+        _INT_TO_GRPC_CODE[x.value[0]] = x
+
 
 class GoogleAPIError(Exception):
     """Base class for all exceptions raised by Google API Clients."""
@@ -432,7 +441,7 @@ def from_grpc_status(status_code, message, **kwargs):
     """Create a :class:`GoogleAPICallError` from a :class:`grpc.StatusCode`.
 
     Args:
-        status_code (grpc.StatusCode): The gRPC status code.
+        status_code (Union[grpc.StatusCode, int]): The gRPC status code.
         message (str): The exception message.
         kwargs: Additional arguments passed to the :class:`GoogleAPICallError`
             constructor.
@@ -441,6 +450,10 @@ def from_grpc_status(status_code, message, **kwargs):
         GoogleAPICallError: An instance of the appropriate subclass of
             :class:`GoogleAPICallError`.
     """
+
+    if isinstance(status_code, int):
+        status_code = _INT_TO_GRPC_CODE.get(status_code, status_code)
+
     error_class = exception_class_for_grpc_status(status_code)
     error = error_class(message, **kwargs)
 
diff --git a/google/api_core/operation.py b/google/api_core/operation.py
index e6407b8..55adbdd 100644
--- a/google/api_core/operation.py
+++ b/google/api_core/operation.py
@@ -132,8 +132,9 @@ class Operation(polling.PollingFuture):
                 )
                 self.set_result(response)
             elif self._operation.HasField("error"):
-                exception = exceptions.GoogleAPICallError(
-                    self._operation.error.message,
+                exception = exceptions.from_grpc_status(
+                    status_code=self._operation.error.code,
+                    message=self._operation.error.message,
                     errors=(self._operation.error,),
                     response=self._operation,
                 )
diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py
index 040ac8a..fb29015 100644
--- a/tests/unit/test_exceptions.py
+++ b/tests/unit/test_exceptions.py
@@ -161,6 +161,17 @@ def test_from_grpc_status():
     assert exception.errors == []
 
 
+def test_from_grpc_status_as_int():
+    message = "message"
+    exception = exceptions.from_grpc_status(11, message)
+    assert isinstance(exception, exceptions.BadRequest)
+    assert isinstance(exception, exceptions.OutOfRange)
+    assert exception.code == http_client.BAD_REQUEST
+    assert exception.grpc_status_code == grpc.StatusCode.OUT_OF_RANGE
+    assert exception.message == message
+    assert exception.errors == []
+
+
 def test_from_grpc_status_with_errors_and_response():
     message = "message"
     response = mock.sentinel.response
diff --git a/tests/unit/test_operation.py b/tests/unit/test_operation.py
index 14b95cb..829a3f3 100644
--- a/tests/unit/test_operation.py
+++ b/tests/unit/test_operation.py
@@ -146,6 +146,23 @@ def test_exception():
     assert expected_exception.message in "{!r}".format(exception)
 
 
+def test_exception_with_error_code():
+    expected_exception = status_pb2.Status(message="meep", code=5)
+    responses = [
+        make_operation_proto(),
+        # Second operation response includes the error.
+        make_operation_proto(done=True, error=expected_exception),
+    ]
+    future, _, _ = make_operation_future(responses)
+
+    exception = future.exception()
+
+    assert expected_exception.message in "{!r}".format(exception)
+    # Status Code 5 maps to Not Found
+    # https://developers.google.com/maps-booking/reference/grpc-api/status_codes
+    assert isinstance(exception, exceptions.NotFound)
+
+
 def test_unexpected_result():
     responses = [
         make_operation_proto(),