Moving Transaction tombstone behavior onto Batch.

This is so that commit request protos are not re-used,
so we can re-use and modify them when passed into
`Connection.commit`.

gcloud/datastore/batch.py

@@ -63,10 +63,23 @@ class Batch(object):
 
     _id = None  # "protected" attribute, always None for non-transactions
 
+    _INITIAL = 0
+    """Enum value for _INITIAL status of batch/transaction."""
+
+    _IN_PROGRESS = 1
+    """Enum value for _IN_PROGRESS status of batch/transaction."""
+
+    _ABORTED = 2
+    """Enum value for _ABORTED status of batch/transaction."""
+
+    _FINISHED = 3
+    """Enum value for _FINISHED status of batch/transaction."""
+
     def __init__(self, client):
         self._client = client
         self._commit_request = _datastore_pb2.CommitRequest()
         self._partial_key_entities = []
+        self._status = self._INITIAL
 
     def current(self):
         """Return the topmost batch / transaction, or None."""
@@ -202,18 +215,26 @@ def delete(self, key):
         self._add_delete_key_pb().CopyFrom(key_pb)
 
     def begin(self):
-        """No-op
+        """Begins a batch.
+
+        This method is called automatically when entering a with
+        statement, however it can be called explicitly if you don't want
+        to use a context manager.
 
         Overridden by :class:`gcloud.datastore.transaction.Transaction`.
+
+        :raises: :class:`ValueError` if the batch has already begun.
         """
+        if self._status != self._INITIAL:
+            raise ValueError('Batch already started previously.')
+        self._status = self._IN_PROGRESS
 
-    def commit(self):
+    def _commit(self):
         """Commits the batch.
 
-        This is called automatically upon exiting a with statement,
-        however it can be called explicitly if you don't want to use a
-        context manager.
+        This is called by :meth:`commit`.
         """
+        # NOTE: ``self._commit_request`` will be modified.
         _, updated_keys = self.connection.commit(
             self.dataset_id, self._commit_request, self._id)
         # If the back-end returns without error, we are guaranteed that
@@ -224,12 +245,26 @@ def commit(self):
             new_id = new_key_pb.path_element[-1].id
             entity.key = entity.key.completed_key(new_id)
 
+    def commit(self):
+        """Commits the batch.
+
+        This is called automatically upon exiting a with statement,
+        however it can be called explicitly if you don't want to use a
+        context manager.
+        """
+        try:
+            self._commit()
+        finally:
+            self._status = self._FINISHED
+
     def rollback(self):
-        """No-op
+        """Rolls back the current batch.
+
+        Marks the batch as aborted (can't be used again).
 
         Overridden by :class:`gcloud.datastore.transaction.Transaction`.
         """
-        pass
+        self._status = self._ABORTED
 
     def __enter__(self):
         self._client._push_batch(self)

gcloud/datastore/connection.py

@@ -293,30 +293,31 @@ def begin_transaction(self, dataset_id):
                              _datastore_pb2.BeginTransactionResponse)
         return response.transaction
 
-    def commit(self, dataset_id, commit_request, transaction_id):
+    def commit(self, dataset_id, request, transaction_id):
         """Commit mutations in context of current transation (if any).
 
         Maps the ``DatastoreService.Commit`` protobuf RPC.
 
         :type dataset_id: string
         :param dataset_id: The ID dataset to which the transaction applies.
 
-        :type commit_request: :class:`._generated.datastore_pb2.CommitRequest`
-        :param commit_request: The protobuf with the mutations being committed.
+        :type request: :class:`._generated.datastore_pb2.CommitRequest`
+        :param request: The protobuf with the mutations being committed.
 
         :type transaction_id: string or None
         :param transaction_id: The transaction ID returned from
                                :meth:`begin_transaction`.  Non-transactional
                                batches must pass ``None``.
 
+        .. note::
+
+            This method will mutate ``request`` before using it.
+
         :rtype: tuple
         :returns': The pair of the number of index updates and a list of
                    :class:`._generated.entity_pb2.Key` for each incomplete key
                    that was completed in the commit.
         """
-        request = _datastore_pb2.CommitRequest()
-        request.CopyFrom(commit_request)
-
         if transaction_id:
             request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
             request.transaction = transaction_id

gcloud/datastore/transaction.py

@@ -90,22 +90,9 @@ class Transaction(Batch):
     :param client: the client used to connect to datastore.
     """
 
-    _INITIAL = 0
-    """Enum value for _INITIAL status of transaction."""
-
-    _IN_PROGRESS = 1
-    """Enum value for _IN_PROGRESS status of transaction."""
-
-    _ABORTED = 2
-    """Enum value for _ABORTED status of transaction."""
-
-    _FINISHED = 3
-    """Enum value for _FINISHED status of transaction."""
-
     def __init__(self, client):
         super(Transaction, self).__init__(client)
         self._id = None
-        self._status = self._INITIAL
 
     @property
     def id(self):
@@ -139,9 +126,7 @@ def begin(self):
 
         :raises: :class:`ValueError` if the transaction has already begun.
         """
-        if self._status != self._INITIAL:
-            raise ValueError('Transaction already started previously.')
-        self._status = self._IN_PROGRESS
+        super(Transaction, self).begin()
         self._id = self.connection.begin_transaction(self.dataset_id)
 
     def rollback(self):
@@ -155,7 +140,7 @@ def rollback(self):
         try:
             self.connection.rollback(self.dataset_id, self._id)
         finally:
-            self._status = self._ABORTED
+            super(Transaction, self).rollback()
             # Clear our own ID in case this gets accidentally reused.
             self._id = None
 
@@ -173,6 +158,5 @@ def commit(self):
         try:
             super(Transaction, self).commit()
         finally:
-            self._status = self._FINISHED
             # Clear our own ID in case this gets accidentally reused.
             self._id = None