Skip to content

Commit 6f2d0af

Browse files
dhermesdhermes
authored
Merge pull request googleapis#2590 from dhermes/fix-2497
Remapping (almost) all RPC status codes to our exceptions in datastore.
2 parents 52d1aa3 + fde56c1 commit 6f2d0af

2 files changed

Lines changed: 125 additions & 96 deletions

File tree

‎datastore/google/cloud/datastore/connection.py‎

Lines changed: 53 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""Connections to Google Cloud Datastore API servers."""
1616

17+
import contextlib
1718
import os
1819

1920
from google.rpc import status_pb2
@@ -23,19 +24,35 @@
2324
from google.cloud import connection as connection_module
2425
from google.cloud.environment_vars import DISABLE_GRPC
2526
from google.cloud.environment_vars import GCD_HOST
26-
from google.cloud.exceptions import BadRequest
27-
from google.cloud.exceptions import Conflict
28-
from google.cloud.exceptions import GrpcRendezvous
29-
from google.cloud.exceptions import make_exception
27+
from google.cloud import exceptions
3028
from google.cloud.datastore._generated import datastore_pb2 as _datastore_pb2
3129
try:
3230
from grpc import StatusCode
3331
from google.cloud.datastore._generated import datastore_grpc_pb2
3432
except ImportError: # pragma: NO COVER
33+
_GRPC_ERROR_MAPPING = {}
3534
_HAVE_GRPC = False
3635
datastore_grpc_pb2 = None
3736
StatusCode = None
3837
else:
38+
# NOTE: We don't include OK -> 200 or CANCELLED -> 499
39+
_GRPC_ERROR_MAPPING = {
40+
StatusCode.UNKNOWN: exceptions.InternalServerError,
41+
StatusCode.INVALID_ARGUMENT: exceptions.BadRequest,
42+
StatusCode.DEADLINE_EXCEEDED: exceptions.GatewayTimeout,
43+
StatusCode.NOT_FOUND: exceptions.NotFound,
44+
StatusCode.ALREADY_EXISTS: exceptions.Conflict,
45+
StatusCode.PERMISSION_DENIED: exceptions.Forbidden,
46+
StatusCode.UNAUTHENTICATED: exceptions.Unauthorized,
47+
StatusCode.RESOURCE_EXHAUSTED: exceptions.TooManyRequests,
48+
StatusCode.FAILED_PRECONDITION: exceptions.PreconditionFailed,
49+
StatusCode.ABORTED: exceptions.Conflict,
50+
StatusCode.OUT_OF_RANGE: exceptions.BadRequest,
51+
StatusCode.UNIMPLEMENTED: exceptions.MethodNotImplemented,
52+
StatusCode.INTERNAL: exceptions.InternalServerError,
53+
StatusCode.UNAVAILABLE: exceptions.ServiceUnavailable,
54+
StatusCode.DATA_LOSS: exceptions.InternalServerError,
55+
}
3956
_HAVE_GRPC = True
4057

4158

@@ -93,7 +110,8 @@ def _request(self, project, method, data):
93110
status = headers['status']
94111
if status != '200':
95112
error_status = status_pb2.Status.FromString(content)
96-
raise make_exception(headers, error_status.message, use_json=False)
113+
raise exceptions.make_exception(
114+
headers, error_status.message, use_json=False)
97115

98116
return content
99117

@@ -220,6 +238,28 @@ def allocate_ids(self, project, request_pb):
220238
_datastore_pb2.AllocateIdsResponse)
221239

222240

241+
@contextlib.contextmanager
242+
def _grpc_catch_rendezvous():
243+
"""Re-map gRPC exceptions that happen in context.
244+
245+
.. _code.proto: https://github.com/googleapis/googleapis/blob/\
246+
master/google/rpc/code.proto
247+
248+
Remaps gRPC exceptions to the classes defined in
249+
:mod:`~google.cloud.exceptions` (according to the description
250+
in `code.proto`_).
251+
"""
252+
try:
253+
yield
254+
except exceptions.GrpcRendezvous as exc:
255+
error_code = exc.code()
256+
error_class = _GRPC_ERROR_MAPPING.get(error_code)
257+
if error_class is None:
258+
raise
259+
else:
260+
raise error_class(exc.details())
261+
262+
223263
class _DatastoreAPIOverGRPC(object):
224264
"""Helper mapping datastore API methods.
225265
@@ -276,13 +316,8 @@ def run_query(self, project, request_pb):
276316
:returns: The returned protobuf response object.
277317
"""
278318
request_pb.project_id = project
279-
try:
319+
with _grpc_catch_rendezvous():
280320
return self._stub.RunQuery(request_pb)
281-
except GrpcRendezvous as exc:
282-
error_code = exc.code()
283-
if error_code == StatusCode.INVALID_ARGUMENT:
284-
raise BadRequest(exc.details())
285-
raise
286321

287322
def begin_transaction(self, project, request_pb):
288323
"""Perform a ``beginTransaction`` request.
@@ -299,7 +334,8 @@ def begin_transaction(self, project, request_pb):
299334
:returns: The returned protobuf response object.
300335
"""
301336
request_pb.project_id = project
302-
return self._stub.BeginTransaction(request_pb)
337+
with _grpc_catch_rendezvous():
338+
return self._stub.BeginTransaction(request_pb)
303339

304340
def commit(self, project, request_pb):
305341
"""Perform a ``commit`` request.
@@ -315,15 +351,8 @@ def commit(self, project, request_pb):
315351
:returns: The returned protobuf response object.
316352
"""
317353
request_pb.project_id = project
318-
try:
354+
with _grpc_catch_rendezvous():
319355
return self._stub.Commit(request_pb)
320-
except GrpcRendezvous as exc:
321-
error_code = exc.code()
322-
if error_code == StatusCode.ABORTED:
323-
raise Conflict(exc.details())
324-
if error_code == StatusCode.INVALID_ARGUMENT:
325-
raise BadRequest(exc.details())
326-
raise
327356

328357
def rollback(self, project, request_pb):
329358
"""Perform a ``rollback`` request.
@@ -339,7 +368,8 @@ def rollback(self, project, request_pb):
339368
:returns: The returned protobuf response object.
340369
"""
341370
request_pb.project_id = project
342-
return self._stub.Rollback(request_pb)
371+
with _grpc_catch_rendezvous():
372+
return self._stub.Rollback(request_pb)
343373

344374
def allocate_ids(self, project, request_pb):
345375
"""Perform an ``allocateIds`` request.
@@ -355,7 +385,8 @@ def allocate_ids(self, project, request_pb):
355385
:returns: The returned protobuf response object.
356386
"""
357387
request_pb.project_id = project
358-
return self._stub.AllocateIds(request_pb)
388+
with _grpc_catch_rendezvous():
389+
return self._stub.AllocateIds(request_pb)
359390

360391

361392
class Connection(connection_module.Connection):

‎datastore/unit_tests/test_connection.py‎

Lines changed: 72 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,72 @@ def test__request_not_200(self):
106106
[{'method': METHOD, 'project': PROJECT}])
107107

108108

109+
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
110+
class Test__grpc_catch_rendezvous(unittest.TestCase):
111+
112+
def _callFUT(self):
113+
from google.cloud.datastore.connection import _grpc_catch_rendezvous
114+
return _grpc_catch_rendezvous()
115+
116+
@staticmethod
117+
def _fake_method(exc, result=None):
118+
if exc is None:
119+
return result
120+
else:
121+
raise exc
122+
123+
def test_success(self):
124+
expected = object()
125+
with self._callFUT():
126+
result = self._fake_method(None, expected)
127+
self.assertIs(result, expected)
128+
129+
def test_failure_aborted(self):
130+
from grpc import StatusCode
131+
from grpc._channel import _RPCState
132+
from google.cloud.exceptions import Conflict
133+
from google.cloud.exceptions import GrpcRendezvous
134+
135+
details = 'Bad things.'
136+
exc_state = _RPCState((), None, None, StatusCode.ABORTED, details)
137+
exc = GrpcRendezvous(exc_state, None, None, None)
138+
with self.assertRaises(Conflict):
139+
with self._callFUT():
140+
self._fake_method(exc)
141+
142+
def test_failure_invalid_argument(self):
143+
from grpc import StatusCode
144+
from grpc._channel import _RPCState
145+
from google.cloud.exceptions import BadRequest
146+
from google.cloud.exceptions import GrpcRendezvous
147+
148+
details = ('Cannot have inequality filters on multiple '
149+
'properties: [created, priority]')
150+
exc_state = _RPCState((), None, None,
151+
StatusCode.INVALID_ARGUMENT, details)
152+
exc = GrpcRendezvous(exc_state, None, None, None)
153+
with self.assertRaises(BadRequest):
154+
with self._callFUT():
155+
self._fake_method(exc)
156+
157+
def test_failure_cancelled(self):
158+
from grpc import StatusCode
159+
from grpc._channel import _RPCState
160+
from google.cloud.exceptions import GrpcRendezvous
161+
162+
exc_state = _RPCState((), None, None, StatusCode.CANCELLED, None)
163+
exc = GrpcRendezvous(exc_state, None, None, None)
164+
with self.assertRaises(GrpcRendezvous):
165+
with self._callFUT():
166+
self._fake_method(exc)
167+
168+
def test_commit_failure_non_grpc_err(self):
169+
exc = RuntimeError('Not a gRPC error')
170+
with self.assertRaises(RuntimeError):
171+
with self._callFUT():
172+
self._fake_method(exc)
173+
174+
109175
class Test_DatastoreAPIOverGRPC(unittest.TestCase):
110176

111177
def _getTargetClass(self):
@@ -227,16 +293,6 @@ def test_run_query_invalid_argument(self):
227293
exc = GrpcRendezvous(exc_state, None, None, None)
228294
self._run_query_failure_helper(exc, BadRequest)
229295

230-
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
231-
def test_run_query_cancelled(self):
232-
from grpc import StatusCode
233-
from grpc._channel import _RPCState
234-
from google.cloud.exceptions import GrpcRendezvous
235-
236-
exc_state = _RPCState((), None, None, StatusCode.CANCELLED, None)
237-
exc = GrpcRendezvous(exc_state, None, None, None)
238-
self._run_query_failure_helper(exc, GrpcRendezvous)
239-
240296
def test_begin_transaction(self):
241297
return_val = object()
242298
stub = _GRPCStub(return_val)
@@ -264,59 +320,6 @@ def test_commit_success(self):
264320
self.assertEqual(stub.method_calls,
265321
[(request_pb, 'Commit')])
266322

267-
def _commit_failure_helper(self, exc, err_class):
268-
stub = _GRPCStub(side_effect=exc)
269-
datastore_api = self._makeOne(stub=stub)
270-
271-
request_pb = _RequestPB()
272-
project = 'PROJECT'
273-
with self.assertRaises(err_class):
274-
datastore_api.commit(project, request_pb)
275-
276-
self.assertEqual(request_pb.project_id, project)
277-
self.assertEqual(stub.method_calls,
278-
[(request_pb, 'Commit')])
279-
280-
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
281-
def test_commit_failure_aborted(self):
282-
from grpc import StatusCode
283-
from grpc._channel import _RPCState
284-
from google.cloud.exceptions import Conflict
285-
from google.cloud.exceptions import GrpcRendezvous
286-
287-
details = 'Bad things.'
288-
exc_state = _RPCState((), None, None, StatusCode.ABORTED, details)
289-
exc = GrpcRendezvous(exc_state, None, None, None)
290-
self._commit_failure_helper(exc, Conflict)
291-
292-
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
293-
def test_commit_failure_invalid_argument(self):
294-
from grpc import StatusCode
295-
from grpc._channel import _RPCState
296-
from google.cloud.exceptions import BadRequest
297-
from google.cloud.exceptions import GrpcRendezvous
298-
299-
details = 'Too long content.'
300-
exc_state = _RPCState((), None, None,
301-
StatusCode.INVALID_ARGUMENT, details)
302-
exc = GrpcRendezvous(exc_state, None, None, None)
303-
self._commit_failure_helper(exc, BadRequest)
304-
305-
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
306-
def test_commit_failure_cancelled(self):
307-
from grpc import StatusCode
308-
from grpc._channel import _RPCState
309-
from google.cloud.exceptions import GrpcRendezvous
310-
311-
exc_state = _RPCState((), None, None, StatusCode.CANCELLED, None)
312-
exc = GrpcRendezvous(exc_state, None, None, None)
313-
self._commit_failure_helper(exc, GrpcRendezvous)
314-
315-
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
316-
def test_commit_failure_non_grpc_err(self):
317-
exc = RuntimeError('Not a gRPC error')
318-
self._commit_failure_helper(exc, RuntimeError)
319-
320323
def test_rollback(self):
321324
return_val = object()
322325
stub = _GRPCStub(return_val)
@@ -1161,27 +1164,22 @@ def __init__(self, return_val=None, side_effect=Exception):
11611164

11621165
def _method(self, request_pb, name):
11631166
self.method_calls.append((request_pb, name))
1164-
return self.return_val
1167+
if self.side_effect is Exception:
1168+
return self.return_val
1169+
else:
1170+
raise self.side_effect
11651171

11661172
def Lookup(self, request_pb):
11671173
return self._method(request_pb, 'Lookup')
11681174

11691175
def RunQuery(self, request_pb):
1170-
result = self._method(request_pb, 'RunQuery')
1171-
if self.side_effect is Exception:
1172-
return result
1173-
else:
1174-
raise self.side_effect
1176+
return self._method(request_pb, 'RunQuery')
11751177

11761178
def BeginTransaction(self, request_pb):
11771179
return self._method(request_pb, 'BeginTransaction')
11781180

11791181
def Commit(self, request_pb):
1180-
result = self._method(request_pb, 'Commit')
1181-
if self.side_effect is Exception:
1182-
return result
1183-
else:
1184-
raise self.side_effect
1182+
return self._method(request_pb, 'Commit')
11851183

11861184
def Rollback(self, request_pb):
11871185
return self._method(request_pb, 'Rollback')

0 commit comments

Comments
 (0)