Merge pull request #3163 from dhermes/wrap-up-2746

Restoring datastore Gax exception re-mapping.

Author: dhermes

datastore/google/cloud/datastore/_gax.py

@@ -103,6 +103,64 @@ class GAPICDatastoreAPI(datastore_client.DatastoreClient):
     :param kwargs: Keyword arguments to pass to constructor.
     """
 
+    def lookup(self, *args, **kwargs):
+        """Perform a ``lookup`` request.
+
+        A light wrapper around the the base method from the parent class.
+        Intended to provide exception re-mapping (from GaxError to our
+        native errors).
+
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
+
+        :rtype: :class:`.datastore_pb2.LookupResponse`
+        :returns: The returned protobuf response object.
+        """
+        with _grpc_catch_rendezvous():
+            return super(GAPICDatastoreAPI, self).lookup(*args, **kwargs)
+
+    def run_query(self, *args, **kwargs):
+        """Perform a ``runQuery`` request.
+
+        A light wrapper around the the base method from the parent class.
+        Intended to provide exception re-mapping (from GaxError to our
+        native errors).
+
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
+
+        :rtype: :class:`.datastore_pb2.RunQueryResponse`
+        :returns: The returned protobuf response object.
+        """
+        with _grpc_catch_rendezvous():
+            return super(GAPICDatastoreAPI, self).run_query(*args, **kwargs)
+
+    def begin_transaction(self, *args, **kwargs):
+        """Perform a ``beginTransaction`` request.
+
+        A light wrapper around the the base method from the parent class.
+        Intended to provide exception re-mapping (from GaxError to our
+        native errors).
+
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
+
+        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
+        :returns: The returned protobuf response object.
+        """
+        with _grpc_catch_rendezvous():
+            return super(GAPICDatastoreAPI, self).begin_transaction(
+                *args, **kwargs)
+
     def commit(self, *args, **kwargs):
         """Perform a ``commit`` request.
 
@@ -122,6 +180,45 @@ def commit(self, *args, **kwargs):
         with _grpc_catch_rendezvous():
             return super(GAPICDatastoreAPI, self).commit(*args, **kwargs)
 
+    def rollback(self, *args, **kwargs):
+        """Perform a ``rollback`` request.
+
+        A light wrapper around the the base method from the parent class.
+        Intended to provide exception re-mapping (from GaxError to our
+        native errors).
+
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
+
+        :rtype: :class:`.datastore_pb2.RollbackResponse`
+        :returns: The returned protobuf response object.
+        """
+        with _grpc_catch_rendezvous():
+            return super(GAPICDatastoreAPI, self).rollback(*args, **kwargs)
+
+    def allocate_ids(self, *args, **kwargs):
+        """Perform an ``allocateIds`` request.
+
+        A light wrapper around the the base method from the parent class.
+        Intended to provide exception re-mapping (from GaxError to our
+        native errors).
+
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
+
+        :rtype: :class:`.datastore_pb2.AllocateIdsResponse`
+        :returns: The returned protobuf response object.
+        """
+        with _grpc_catch_rendezvous():
+            return super(GAPICDatastoreAPI, self).allocate_ids(
+                *args, **kwargs)
+
 
 def make_datastore_api(client):
     """Create an instance of the GAPIC Datastore API.

datastore/unit_tests/test__gax.py

@@ -124,6 +124,69 @@ def _get_target_class():
     def _make_one(self, *args, **kwargs):
         return self._get_target_class()(*args, **kwargs)
 
+    def test_lookup(self):
+        from google.cloud.gapic.datastore.v1 import datastore_client
+
+        patch1 = mock.patch.object(
+            datastore_client.DatastoreClient, '__init__',
+            return_value=None)
+        patch2 = mock.patch.object(datastore_client.DatastoreClient, 'lookup')
+        patch3 = mock.patch(
+            'google.cloud.datastore._gax._grpc_catch_rendezvous')
+
+        with patch1 as mock_constructor:
+            ds_api = self._make_one()
+            mock_constructor.assert_called_once_with()
+            with patch2 as mock_lookup:
+                with patch3 as mock_catch_rendezvous:
+                    mock_catch_rendezvous.assert_not_called()
+                    ds_api.lookup(None, True, bb='cc')
+                    mock_lookup.assert_called_once_with(None, True, bb='cc')
+                    mock_catch_rendezvous.assert_called_once_with()
+
+    def test_run_query(self):
+        from google.cloud.gapic.datastore.v1 import datastore_client
+
+        patch1 = mock.patch.object(
+            datastore_client.DatastoreClient, '__init__',
+            return_value=None)
+        patch2 = mock.patch.object(
+            datastore_client.DatastoreClient, 'run_query')
+        patch3 = mock.patch(
+            'google.cloud.datastore._gax._grpc_catch_rendezvous')
+
+        with patch1 as mock_constructor:
+            ds_api = self._make_one()
+            mock_constructor.assert_called_once_with()
+            with patch2 as mock_run_query:
+                with patch3 as mock_catch_rendezvous:
+                    mock_catch_rendezvous.assert_not_called()
+                    ds_api.run_query('47a', none=None)
+                    mock_run_query.assert_called_once_with('47a', none=None)
+                    mock_catch_rendezvous.assert_called_once_with()
+
+    def test_begin_transaction(self):
+        from google.cloud.gapic.datastore.v1 import datastore_client
+
+        patch1 = mock.patch.object(
+            datastore_client.DatastoreClient, '__init__',
+            return_value=None)
+        patch2 = mock.patch.object(
+            datastore_client.DatastoreClient, 'begin_transaction')
+        patch3 = mock.patch(
+            'google.cloud.datastore._gax._grpc_catch_rendezvous')
+
+        with patch1 as mock_constructor:
+            ds_api = self._make_one()
+            mock_constructor.assert_called_once_with()
+            with patch2 as mock_begin_transaction:
+                with patch3 as mock_catch_rendezvous:
+                    mock_catch_rendezvous.assert_not_called()
+                    ds_api.begin_transaction('a', 'b', [], key='kei')
+                    mock_begin_transaction.assert_called_once_with(
+                        'a', 'b', [], key='kei')
+                    mock_catch_rendezvous.assert_called_once_with()
+
     def test_commit(self):
         from google.cloud.gapic.datastore.v1 import datastore_client
 
@@ -144,6 +207,50 @@ def test_commit(self):
                     mock_commit.assert_called_once_with(1, 2, a=3)
                     mock_catch_rendezvous.assert_called_once_with()
 
+    def test_rollback(self):
+        from google.cloud.gapic.datastore.v1 import datastore_client
+
+        patch1 = mock.patch.object(
+            datastore_client.DatastoreClient, '__init__',
+            return_value=None)
+        patch2 = mock.patch.object(
+            datastore_client.DatastoreClient, 'rollback')
+        patch3 = mock.patch(
+            'google.cloud.datastore._gax._grpc_catch_rendezvous')
+
+        with patch1 as mock_constructor:
+            ds_api = self._make_one()
+            mock_constructor.assert_called_once_with()
+            with patch2 as mock_rollback:
+                with patch3 as mock_catch_rendezvous:
+                    mock_catch_rendezvous.assert_not_called()
+                    ds_api.rollback(11, 12, arp='marp')
+                    mock_rollback.assert_called_once_with(11, 12, arp='marp')
+                    mock_catch_rendezvous.assert_called_once_with()
+
+    def test_allocate_ids(self):
+        from google.cloud.gapic.datastore.v1 import datastore_client
+
+        patch1 = mock.patch.object(
+            datastore_client.DatastoreClient, '__init__',
+            return_value=None)
+        patch2 = mock.patch.object(
+            datastore_client.DatastoreClient, 'allocate_ids')
+        patch3 = mock.patch(
+            'google.cloud.datastore._gax._grpc_catch_rendezvous')
+
+        with patch1 as mock_constructor:
+            ds_api = self._make_one()
+            mock_constructor.assert_called_once_with()
+            with patch2 as mock_allocate_ids:
+                with patch3 as mock_catch_rendezvous:
+                    mock_catch_rendezvous.assert_not_called()
+                    ds_api.allocate_ids(
+                        'hey', 'bai', bye=(47, 4), shy={'a': 4})
+                    mock_allocate_ids.assert_called_once_with(
+                        'hey', 'bai', bye=(47, 4), shy={'a': 4})
+                    mock_catch_rendezvous.assert_called_once_with()
+
 
 @unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
 class Test_make_datastore_api(unittest.TestCase):