Moving Transaction tombstone behavior onto Batch.

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.project, 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, project):
                              _datastore_pb2.BeginTransactionResponse)
         return response.transaction
 
-    def commit(self, project, commit_request, transaction_id):
+    def commit(self, project, request, transaction_id):
         """Commit mutations in context of current transation (if any).
 
         Maps the ``DatastoreService.Commit`` protobuf RPC.
 
         :type project: string
         :param project: The project 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/test_batch.py

@@ -37,6 +37,7 @@ def test_ctor(self):
         self.assertEqual(batch.connection, connection)
         self.assertEqual(batch.namespace, _NAMESPACE)
         self.assertTrue(batch._id is None)
+        self.assertEqual(batch._status, batch._INITIAL)
         self.assertTrue(isinstance(batch.mutations, datastore_pb2.Mutation))
         self.assertEqual(batch._partial_key_entities, [])
 
@@ -200,13 +201,39 @@ def test_delete_w_completed_key_w_prefixed_project(self):
         mutated_key = _mutated_pb(self, batch.mutations, 'delete')
         self.assertEqual(mutated_key, key._key)
 
+    def test_begin(self):
+        _PROJECT = 'PROJECT'
+        client = _Client(_PROJECT, None)
+        batch = self._makeOne(client)
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+
+    def test_begin_fail(self):
+        _PROJECT = 'PROJECT'
+        client = _Client(_PROJECT, None)
+        batch = self._makeOne(client)
+        batch._status = batch._IN_PROGRESS
+        with self.assertRaises(ValueError):
+            batch.begin()
+
+    def test_rollback(self):
+        _PROJECT = 'PROJECT'
+        client = _Client(_PROJECT, None)
+        batch = self._makeOne(client)
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.rollback()
+        self.assertEqual(batch._status, batch._ABORTED)
+
     def test_commit(self):
         _PROJECT = 'PROJECT'
         connection = _Connection()
         client = _Client(_PROJECT, connection)
         batch = self._makeOne(client)
 
+        self.assertEqual(batch._status, batch._INITIAL)
         batch.commit()
+        self.assertEqual(batch._status, batch._FINISHED)
 
         self.assertEqual(connection._committed,
                          [(_PROJECT, batch._commit_request, None)])
@@ -222,7 +249,9 @@ def test_commit_w_partial_key_entities(self):
         key._id = None
         batch._partial_key_entities.append(entity)
 
+        self.assertEqual(batch._status, batch._INITIAL)
         batch.commit()
+        self.assertEqual(batch._status, batch._FINISHED)
 
         self.assertEqual(connection._committed,
                          [(_PROJECT, batch._commit_request, None)])

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.project)
 
     def rollback(self):
@@ -155,7 +140,7 @@ def rollback(self):
         try:
             self.connection.rollback(self.project, 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