Skip to content

Commit bbeab3d

Browse files
dhermesdhermes
committed
Restoring datastore Gax exception re-mapping.
Fixes #2746.
1 parent 97379fc commit bbeab3d

File tree

2 files changed

+204
-0
lines changed

2 files changed

+204
-0
lines changed

‎datastore/google/cloud/datastore/_gax.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,64 @@ class GAPICDatastoreAPI(datastore_client.DatastoreClient):
103103
:param kwargs: Keyword arguments to pass to constructor.
104104
"""
105105

106+
def lookup(self, *args, **kwargs):
107+
"""Perform a ``lookup`` request.
108+
109+
A light wrapper around the the base method from the parent class.
110+
Intended to provide exception re-mapping (from GaxError to our
111+
native errors).
112+
113+
:type args: tuple
114+
:param args: Positional arguments to pass to base method.
115+
116+
:type kwargs: dict
117+
:param kwargs: Keyword arguments to pass to base method.
118+
119+
:rtype: :class:`.datastore_pb2.LookupResponse`
120+
:returns: The returned protobuf response object.
121+
"""
122+
with _grpc_catch_rendezvous():
123+
return super(GAPICDatastoreAPI, self).lookup(*args, **kwargs)
124+
125+
def run_query(self, *args, **kwargs):
126+
"""Perform a ``runQuery`` request.
127+
128+
A light wrapper around the the base method from the parent class.
129+
Intended to provide exception re-mapping (from GaxError to our
130+
native errors).
131+
132+
:type args: tuple
133+
:param args: Positional arguments to pass to base method.
134+
135+
:type kwargs: dict
136+
:param kwargs: Keyword arguments to pass to base method.
137+
138+
:rtype: :class:`.datastore_pb2.RunQueryResponse`
139+
:returns: The returned protobuf response object.
140+
"""
141+
with _grpc_catch_rendezvous():
142+
return super(GAPICDatastoreAPI, self).run_query(*args, **kwargs)
143+
144+
def begin_transaction(self, *args, **kwargs):
145+
"""Perform a ``beginTransaction`` request.
146+
147+
A light wrapper around the the base method from the parent class.
148+
Intended to provide exception re-mapping (from GaxError to our
149+
native errors).
150+
151+
:type args: tuple
152+
:param args: Positional arguments to pass to base method.
153+
154+
:type kwargs: dict
155+
:param kwargs: Keyword arguments to pass to base method.
156+
157+
:rtype: :class:`.datastore_pb2.BeginTransactionResponse`
158+
:returns: The returned protobuf response object.
159+
"""
160+
with _grpc_catch_rendezvous():
161+
return super(GAPICDatastoreAPI, self).begin_transaction(
162+
*args, **kwargs)
163+
106164
def commit(self, *args, **kwargs):
107165
"""Perform a ``commit`` request.
108166
@@ -122,6 +180,45 @@ def commit(self, *args, **kwargs):
122180
with _grpc_catch_rendezvous():
123181
return super(GAPICDatastoreAPI, self).commit(*args, **kwargs)
124182

183+
def rollback(self, *args, **kwargs):
184+
"""Perform a ``rollback`` request.
185+
186+
A light wrapper around the the base method from the parent class.
187+
Intended to provide exception re-mapping (from GaxError to our
188+
native errors).
189+
190+
:type args: tuple
191+
:param args: Positional arguments to pass to base method.
192+
193+
:type kwargs: dict
194+
:param kwargs: Keyword arguments to pass to base method.
195+
196+
:rtype: :class:`.datastore_pb2.RollbackResponse`
197+
:returns: The returned protobuf response object.
198+
"""
199+
with _grpc_catch_rendezvous():
200+
return super(GAPICDatastoreAPI, self).rollback(*args, **kwargs)
201+
202+
def allocate_ids(self, *args, **kwargs):
203+
"""Perform an ``allocateIds`` request.
204+
205+
A light wrapper around the the base method from the parent class.
206+
Intended to provide exception re-mapping (from GaxError to our
207+
native errors).
208+
209+
:type args: tuple
210+
:param args: Positional arguments to pass to base method.
211+
212+
:type kwargs: dict
213+
:param kwargs: Keyword arguments to pass to base method.
214+
215+
:rtype: :class:`.datastore_pb2.AllocateIdsResponse`
216+
:returns: The returned protobuf response object.
217+
"""
218+
with _grpc_catch_rendezvous():
219+
return super(GAPICDatastoreAPI, self).allocate_ids(
220+
*args, **kwargs)
221+
125222

126223
def make_datastore_api(client):
127224
"""Create an instance of the GAPIC Datastore API.

‎datastore/unit_tests/test__gax.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,69 @@ def _get_target_class():
124124
def _make_one(self, *args, **kwargs):
125125
return self._get_target_class()(*args, **kwargs)
126126

127+
def test_lookup(self):
128+
from google.cloud.gapic.datastore.v1 import datastore_client
129+
130+
patch1 = mock.patch.object(
131+
datastore_client.DatastoreClient, '__init__',
132+
return_value=None)
133+
patch2 = mock.patch.object(datastore_client.DatastoreClient, 'lookup')
134+
patch3 = mock.patch(
135+
'google.cloud.datastore._gax._grpc_catch_rendezvous')
136+
137+
with patch1 as mock_constructor:
138+
ds_api = self._make_one()
139+
mock_constructor.assert_called_once_with()
140+
with patch2 as mock_lookup:
141+
with patch3 as mock_catch_rendezvous:
142+
mock_catch_rendezvous.assert_not_called()
143+
ds_api.lookup(None, True, bb='cc')
144+
mock_lookup.assert_called_once_with(None, True, bb='cc')
145+
mock_catch_rendezvous.assert_called_once_with()
146+
147+
def test_run_query(self):
148+
from google.cloud.gapic.datastore.v1 import datastore_client
149+
150+
patch1 = mock.patch.object(
151+
datastore_client.DatastoreClient, '__init__',
152+
return_value=None)
153+
patch2 = mock.patch.object(
154+
datastore_client.DatastoreClient, 'run_query')
155+
patch3 = mock.patch(
156+
'google.cloud.datastore._gax._grpc_catch_rendezvous')
157+
158+
with patch1 as mock_constructor:
159+
ds_api = self._make_one()
160+
mock_constructor.assert_called_once_with()
161+
with patch2 as mock_run_query:
162+
with patch3 as mock_catch_rendezvous:
163+
mock_catch_rendezvous.assert_not_called()
164+
ds_api.run_query('47a', none=None)
165+
mock_run_query.assert_called_once_with('47a', none=None)
166+
mock_catch_rendezvous.assert_called_once_with()
167+
168+
def test_begin_transaction(self):
169+
from google.cloud.gapic.datastore.v1 import datastore_client
170+
171+
patch1 = mock.patch.object(
172+
datastore_client.DatastoreClient, '__init__',
173+
return_value=None)
174+
patch2 = mock.patch.object(
175+
datastore_client.DatastoreClient, 'begin_transaction')
176+
patch3 = mock.patch(
177+
'google.cloud.datastore._gax._grpc_catch_rendezvous')
178+
179+
with patch1 as mock_constructor:
180+
ds_api = self._make_one()
181+
mock_constructor.assert_called_once_with()
182+
with patch2 as mock_begin_transaction:
183+
with patch3 as mock_catch_rendezvous:
184+
mock_catch_rendezvous.assert_not_called()
185+
ds_api.begin_transaction('a', 'b', [], key='kei')
186+
mock_begin_transaction.assert_called_once_with(
187+
'a', 'b', [], key='kei')
188+
mock_catch_rendezvous.assert_called_once_with()
189+
127190
def test_commit(self):
128191
from google.cloud.gapic.datastore.v1 import datastore_client
129192

@@ -144,6 +207,50 @@ def test_commit(self):
144207
mock_commit.assert_called_once_with(1, 2, a=3)
145208
mock_catch_rendezvous.assert_called_once_with()
146209

210+
def test_rollback(self):
211+
from google.cloud.gapic.datastore.v1 import datastore_client
212+
213+
patch1 = mock.patch.object(
214+
datastore_client.DatastoreClient, '__init__',
215+
return_value=None)
216+
patch2 = mock.patch.object(
217+
datastore_client.DatastoreClient, 'rollback')
218+
patch3 = mock.patch(
219+
'google.cloud.datastore._gax._grpc_catch_rendezvous')
220+
221+
with patch1 as mock_constructor:
222+
ds_api = self._make_one()
223+
mock_constructor.assert_called_once_with()
224+
with patch2 as mock_rollback:
225+
with patch3 as mock_catch_rendezvous:
226+
mock_catch_rendezvous.assert_not_called()
227+
ds_api.rollback(11, 12, arp='marp')
228+
mock_rollback.assert_called_once_with(11, 12, arp='marp')
229+
mock_catch_rendezvous.assert_called_once_with()
230+
231+
def test_allocate_ids(self):
232+
from google.cloud.gapic.datastore.v1 import datastore_client
233+
234+
patch1 = mock.patch.object(
235+
datastore_client.DatastoreClient, '__init__',
236+
return_value=None)
237+
patch2 = mock.patch.object(
238+
datastore_client.DatastoreClient, 'allocate_ids')
239+
patch3 = mock.patch(
240+
'google.cloud.datastore._gax._grpc_catch_rendezvous')
241+
242+
with patch1 as mock_constructor:
243+
ds_api = self._make_one()
244+
mock_constructor.assert_called_once_with()
245+
with patch2 as mock_allocate_ids:
246+
with patch3 as mock_catch_rendezvous:
247+
mock_catch_rendezvous.assert_not_called()
248+
ds_api.allocate_ids(
249+
'hey', 'bai', bye=(47, 4), shy={'a': 4})
250+
mock_allocate_ids.assert_called_once_with(
251+
'hey', 'bai', bye=(47, 4), shy={'a': 4})
252+
mock_catch_rendezvous.assert_called_once_with()
253+
147254

148255
@unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
149256
class Test_make_datastore_api(unittest.TestCase):

0 commit comments

Comments
 (0)