fix: fix delete in transaction (#333)

This fixes a bug that caused deleting an entity inside of a transaction
to hang.

Fixes #271

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/_datastore_api.py

@@ -647,7 +647,8 @@ def put(self, entity_pb):
         self.mutations.append(mutation)
 
         # If we have an incomplete key, add the incomplete key to a batch for a
-        # call to AllocateIds
+        # call to AllocateIds, since the call to actually store the entity
+        # won't happen until the end of the transaction.
         if not _complete(entity_pb.key):
             # If this is the first key in the batch, we also need to
             # schedule our idle handler to get called
@@ -657,12 +658,28 @@ def put(self, entity_pb):
             self.incomplete_mutations.append(mutation)
             self.incomplete_futures.append(future)
 
-        # Complete keys get passed back None
+        # Can't wait for result, since batch won't be sent until transaction
+        # has ended. Complete keys get passed back None.
         else:
             future.set_result(None)
 
         return future
 
+    def delete(self, key):
+        """Add a key to batch to be deleted.
+
+        Args:
+            entity_pb (datastore.Key): The entity's key to be deleted.
+
+        Returns:
+            tasklets.Future: Result will be :data:`None`, always.
+        """
+        # Can't wait for result, since batch won't be sent until transaction
+        # has ended.
+        future = super(_TransactionalCommitBatch, self).delete(key)
+        future.set_result(None)
+        return future
+
     def idle_callback(self):
         """Call AllocateIds on any incomplete keys in the batch."""
         if not self.incomplete_mutations:

packages/google-cloud-ndb/google/cloud/ndb/_transaction.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 
 import functools
+import logging
 
 from google.cloud.ndb import exceptions
 from google.cloud.ndb import _retry
 from google.cloud.ndb import tasklets
 
+log = logging.getLogger(__name__)
+
 
 def in_transaction():
     """Determine if there is a currently active transaction.
@@ -102,9 +105,11 @@ def _transaction_async(context, callback, read_only=False):
     from google.cloud.ndb import _datastore_api
 
     # Start the transaction
+    log.debug("Start transaction")
     transaction_id = yield _datastore_api.begin_transaction(
         read_only, retries=0
     )
+    log.debug("Transaction Id: {}".format(transaction_id))
 
     on_commit_callbacks = []
     tx_context = context.new(

packages/google-cloud-ndb/tests/system/test_crud.py

@@ -1167,3 +1167,31 @@ class SomeKind(ndb.Model):
 
     keys = SomeKind.allocate_ids(N)
     assert len(keys) == N
+
+
+@pytest.mark.usefixtures("client_context")
+def test_delete_multi_with_transactional(dispose_of):
+    """Regression test for issue #271
+
+    https://github.com/googleapis/python-ndb/issues/271
+    """
+    N = 10
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    @ndb.transactional()
+    def delete_them(entities):
+        ndb.delete_multi([entity.key for entity in entities])
+
+    foos = list(range(N))
+    entities = [SomeKind(foo=foo) for foo in foos]
+    keys = ndb.put_multi(entities)
+    dispose_of(*(key._key for key in keys))
+
+    entities = ndb.get_multi(keys)
+    assert [entity.foo for entity in entities] == foos
+
+    assert delete_them(entities) is None
+    entities = ndb.get_multi(keys)
+    assert entities == [None] * N

packages/google-cloud-ndb/tests/unit/test__datastore_api.py

@@ -777,18 +777,15 @@ def __init__(self, delete=None):
             def __eq__(self, other):
                 return self.delete == other.delete
 
-        eventloop = mock.Mock(spec=("add_idle", "run"))
-        with in_context.new(
-            eventloop=eventloop, transaction=b"tx123"
-        ).use() as context:
+        with in_context.new(transaction=b"tx123").use() as context:
             datastore_pb2.Mutation = Mutation
 
             key1 = key_module.Key("SomeKind", 1)._key
             key2 = key_module.Key("SomeKind", 2)._key
             key3 = key_module.Key("SomeKind", 3)._key
-            _api.delete(key1, _options.Options())
-            _api.delete(key2, _options.Options())
-            _api.delete(key3, _options.Options())
+            assert _api.delete(key1, _options.Options()).result() is None
+            assert _api.delete(key2, _options.Options()).result() is None
+            assert _api.delete(key3, _options.Options()).result() is None
 
             batch = context.commit_batches[b"tx123"]
             assert batch.mutations == [