Merge pull request #3127 from dhermes/ds-move-commit-to-GAPIC

Using GAPIC datastore object (and an HTTP equivalent) for commit.

Author: dhermes

datastore/google/cloud/datastore/_gax.py

@@ -16,9 +16,12 @@
 
 
 import contextlib
+import sys
 
 from google.cloud.gapic.datastore.v1 import datastore_client
 from google.cloud.proto.datastore.v1 import datastore_pb2_grpc
+from google.gax.errors import GaxError
+from google.gax.grpc import exc_to_code
 from google.gax.utils import metrics
 from grpc import StatusCode
 import six
@@ -71,13 +74,22 @@ def _grpc_catch_rendezvous():
     """
     try:
         yield
+    except GaxError as exc:
+        error_code = exc_to_code(exc.cause)
+        error_class = _GRPC_ERROR_MAPPING.get(error_code)
+        if error_class is None:
+            raise
+        else:
+            new_exc = error_class(exc.cause.details())
+            six.reraise(error_class, new_exc, sys.exc_info()[2])
     except exceptions.GrpcRendezvous as exc:
         error_code = exc.code()
         error_class = _GRPC_ERROR_MAPPING.get(error_code)
         if error_class is None:
             raise
         else:
-            raise error_class(exc.details())
+            new_exc = error_class(exc.details())
+            six.reraise(error_class, new_exc, sys.exc_info()[2])
 
 
 class _DatastoreAPIOverGRPC(object):
@@ -158,56 +170,38 @@ def begin_transaction(self, project, request_pb):
         with _grpc_catch_rendezvous():
             return self._stub.BeginTransaction(request_pb)
 
-    def commit(self, project, request_pb):
-        """Perform a ``commit`` request.
-
-        :type project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
-
-        :type request_pb: :class:`.datastore_pb2.CommitRequest`
-        :param request_pb: The request protobuf object.
 
-        :rtype: :class:`.datastore_pb2.CommitResponse`
-        :returns: The returned protobuf response object.
-        """
-        request_pb.project_id = project
-        with _grpc_catch_rendezvous():
-            return self._stub.Commit(request_pb)
+class GAPICDatastoreAPI(datastore_client.DatastoreClient):
+    """An API object that sends proto-over-gRPC requests.
 
-    def rollback(self, project, request_pb):
-        """Perform a ``rollback`` request.
+    A light wrapper around the parent class, with exception re-mapping
+    provided (from GaxError to our native errors).
 
-        :type project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
+    :type args: tuple
+    :param args: Positional arguments to pass to constructor.
 
-        :type request_pb: :class:`.datastore_pb2.RollbackRequest`
-        :param request_pb: The request protobuf object.
+    :type kwargs: dict
+    :param kwargs: Keyword arguments to pass to constructor.
+    """
 
-        :rtype: :class:`.datastore_pb2.RollbackResponse`
-        :returns: The returned protobuf response object.
-        """
-        request_pb.project_id = project
-        with _grpc_catch_rendezvous():
-            return self._stub.Rollback(request_pb)
+    def commit(self, *args, **kwargs):
+        """Perform a ``commit`` request.
 
-    def allocate_ids(self, project, request_pb):
-        """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 project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
+        :type args: tuple
+        :param args: Positional arguments to pass to base method.
 
-        :type request_pb: :class:`.datastore_pb2.AllocateIdsRequest`
-        :param request_pb: The request protobuf object.
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to pass to base method.
 
-        :rtype: :class:`.datastore_pb2.AllocateIdsResponse`
+        :rtype: :class:`.datastore_pb2.CommitResponse`
         :returns: The returned protobuf response object.
         """
-        request_pb.project_id = project
         with _grpc_catch_rendezvous():
-            return self._stub.AllocateIds(request_pb)
+            return super(GAPICDatastoreAPI, self).commit(*args, **kwargs)
 
 
 def make_datastore_api(client):
@@ -222,5 +216,5 @@ def make_datastore_api(client):
     channel = make_secure_channel(
         client._credentials, DEFAULT_USER_AGENT,
         datastore_client.DatastoreClient.SERVICE_ADDRESS)
-    return datastore_client.DatastoreClient(
+    return GAPICDatastoreAPI(
         channel=channel, lib_name='gccl', lib_version=__version__)

datastore/google/cloud/datastore/_http.py

@@ -216,23 +216,6 @@ def begin_transaction(self, project, request_pb):
                     self.connection.api_base_url,
                     request_pb, _datastore_pb2.BeginTransactionResponse)
 
-    def commit(self, project, request_pb):
-        """Perform a ``commit`` request.
-
-        :type project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
-
-        :type request_pb: :class:`.datastore_pb2.CommitRequest`
-        :param request_pb: The request protobuf object.
-
-        :rtype: :class:`.datastore_pb2.CommitResponse`
-        :returns: The returned protobuf response object.
-        """
-        return _rpc(self.connection.http, project, 'commit',
-                    self.connection.api_base_url,
-                    request_pb, _datastore_pb2.CommitResponse)
-
 
 class Connection(connection_module.Connection):
     """A connection to the Google Cloud Datastore via the Protobuf API.
@@ -366,37 +349,6 @@ def begin_transaction(self, project):
         request = _datastore_pb2.BeginTransactionRequest()
         return self._datastore_api.begin_transaction(project, request)
 
-    def commit(self, project, request, transaction_id):
-        """Commit mutations in context of current transaction (if any).
-
-        Maps the ``DatastoreService.Commit`` protobuf RPC.
-
-        :type project: str
-        :param project: The project to which the transaction applies.
-
-        :type request: :class:`.datastore_pb2.CommitRequest`
-        :param request: The protobuf with the mutations being committed.
-
-        :type transaction_id: str
-        :param transaction_id: (Optional) The transaction ID returned from
-                               :meth:`begin_transaction`.  Non-transactional
-                               batches must pass ``None``.
-
-        .. note::
-
-            This method will mutate ``request`` before using it.
-
-        :rtype: :class:`.datastore_pb2.CommitResponse`
-        :returns: The protobuf response from a commit request.
-        """
-        if transaction_id:
-            request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
-            request.transaction = transaction_id
-        else:
-            request.mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
-
-        return self._datastore_api.commit(project, request)
-
 
 class HTTPDatastoreAPI(object):
     """An API object that sends proto-over-HTTP requests.
@@ -410,6 +362,39 @@ class HTTPDatastoreAPI(object):
     def __init__(self, client):
         self.client = client
 
+    def commit(self, project, mode, mutations, transaction=None):
+        """Perform a ``commit`` request.
+
+        :type project: str
+        :param project: The project to connect to. This is
+                        usually your project name in the cloud console.
+
+        :type mode: :class:`.gapic.datastore.v1.enums.CommitRequest.Mode`
+        :param mode: The type of commit to perform. Expected to be one of
+                     ``TRANSACTIONAL`` or ``NON_TRANSACTIONAL``.
+
+        :type mutations: list
+        :param mutations: List of :class:`.datastore_pb2.Mutation`, the
+                          mutations to perform.
+
+        :type transaction: bytes
+        :param transaction: (Optional) The transaction ID returned from
+                            :meth:`begin_transaction`.  Non-transactional
+                            commits must pass :data:`None`.
+
+        :rtype: :class:`.datastore_pb2.CommitResponse`
+        :returns: The returned protobuf response object.
+        """
+        request_pb = _datastore_pb2.CommitRequest(
+            project_id=project,
+            mode=mode,
+            transaction=transaction,
+            mutations=mutations,
+        )
+        return _rpc(self.client._http, project, 'commit',
+                    self.client._base_url,
+                    request_pb, _datastore_pb2.CommitResponse)
+
     def rollback(self, project, transaction_id):
         """Perform a ``rollback`` request.
 

datastore/google/cloud/datastore/batch.py

@@ -77,7 +77,7 @@ class Batch(object):
 
     def __init__(self, client):
         self._client = client
-        self._commit_request = _datastore_pb2.CommitRequest()
+        self._mutations = []
         self._partial_key_entities = []
         self._status = self._INITIAL
 
@@ -110,7 +110,8 @@ def _add_partial_key_entity_pb(self):
         :returns: The newly created entity protobuf that will be
                   updated and sent with a commit.
         """
-        new_mutation = self.mutations.add()
+        new_mutation = _datastore_pb2.Mutation()
+        self._mutations.append(new_mutation)
         return new_mutation.insert
 
     def _add_complete_key_entity_pb(self):
@@ -123,7 +124,8 @@ def _add_complete_key_entity_pb(self):
         # We use ``upsert`` for entities with completed keys, rather than
         # ``insert`` or ``update``, in order not to create race conditions
         # based on prior existence / removal of the entity.
-        new_mutation = self.mutations.add()
+        new_mutation = _datastore_pb2.Mutation()
+        self._mutations.append(new_mutation)
         return new_mutation.upsert
 
     def _add_delete_key_pb(self):
@@ -133,7 +135,8 @@ def _add_delete_key_pb(self):
         :returns: The newly created key protobuf that will be
                   deleted when sent with a commit.
         """
-        new_mutation = self.mutations.add()
+        new_mutation = _datastore_pb2.Mutation()
+        self._mutations.append(new_mutation)
         return new_mutation.delete
 
     @property
@@ -150,7 +153,7 @@ def mutations(self):
         :returns: The list of :class:`.datastore_pb2.Mutation`
                   protobufs to be sent in the commit request.
         """
-        return self._commit_request.mutations
+        return self._mutations
 
     def put(self, entity):
         """Remember an entity's state to be saved during :meth:`commit`.
@@ -237,9 +240,13 @@ def _commit(self):
 
         This is called by :meth:`commit`.
         """
-        # NOTE: ``self._commit_request`` will be modified.
-        commit_response_pb = self._client._connection.commit(
-            self.project, self._commit_request, self._id)
+        if self._id is None:
+            mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        else:
+            mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
+
+        commit_response_pb = self._client._datastore_api.commit(
+            self.project, mode, self._mutations, transaction=self._id)
         _, updated_keys = _parse_commit_response(commit_response_pb)
         # If the back-end returns without error, we are guaranteed that
         # :meth:`Connection.commit` will return keys that match (length and