diff --git a/datastore/google/cloud/datastore/connection.py b/datastore/google/cloud/datastore/connection.py index c2701c218bef..e84457a011ea 100644 --- a/datastore/google/cloud/datastore/connection.py +++ b/datastore/google/cloud/datastore/connection.py @@ -276,7 +276,13 @@ def run_query(self, project, request_pb): :returns: The returned protobuf response object. """ request_pb.project_id = project - return self._stub.RunQuery(request_pb) + try: + return self._stub.RunQuery(request_pb) + except GrpcRendezvous as exc: + error_code = exc.code() + if error_code == StatusCode.INVALID_ARGUMENT: + raise BadRequest(exc.details()) + raise def begin_transaction(self, project, request_pb): """Perform a ``beginTransaction`` request. diff --git a/datastore/unit_tests/test_connection.py b/datastore/unit_tests/test_connection.py index 6791a59cde76..cbafc72ac3f0 100644 --- a/datastore/unit_tests/test_connection.py +++ b/datastore/unit_tests/test_connection.py @@ -200,6 +200,43 @@ def test_run_query(self): self.assertEqual(stub.method_calls, [(request_pb, 'RunQuery')]) + def _run_query_failure_helper(self, exc, err_class): + stub = _GRPCStub(side_effect=exc) + datastore_api = self._makeOne(stub=stub) + + request_pb = _RequestPB() + project = 'PROJECT' + with self.assertRaises(err_class): + datastore_api.run_query(project, request_pb) + + self.assertEqual(request_pb.project_id, project) + self.assertEqual(stub.method_calls, + [(request_pb, 'RunQuery')]) + + @unittest.skipUnless(_HAVE_GRPC, 'No gRPC') + def test_run_query_invalid_argument(self): + from grpc import StatusCode + from grpc._channel import _RPCState + from google.cloud.exceptions import BadRequest + from google.cloud.exceptions import GrpcRendezvous + + details = ('Cannot have inequality filters on multiple ' + 'properties: [created, priority]') + exc_state = _RPCState((), None, None, + StatusCode.INVALID_ARGUMENT, details) + exc = GrpcRendezvous(exc_state, None, None, None) + self._run_query_failure_helper(exc, BadRequest) + + @unittest.skipUnless(_HAVE_GRPC, 'No gRPC') + def test_run_query_cancelled(self): + from grpc import StatusCode + from grpc._channel import _RPCState + from google.cloud.exceptions import GrpcRendezvous + + exc_state = _RPCState((), None, None, StatusCode.CANCELLED, None) + exc = GrpcRendezvous(exc_state, None, None, None) + self._run_query_failure_helper(exc, GrpcRendezvous) + def test_begin_transaction(self): return_val = object() stub = _GRPCStub(return_val) @@ -1130,7 +1167,11 @@ def Lookup(self, request_pb): return self._method(request_pb, 'Lookup') def RunQuery(self, request_pb): - return self._method(request_pb, 'RunQuery') + result = self._method(request_pb, 'RunQuery') + if self.side_effect is Exception: + return result + else: + raise self.side_effect def BeginTransaction(self, request_pb): return self._method(request_pb, 'BeginTransaction')