fix: @non_transactional decorator was not working correctly with async (#554)

* fix: `@non_transactional` decorator was not working correctly with async

Fixes #552

* We don't really need `_get_transaction` at all.

In all cases, we know where to get the transaction from.

Author: Chris Rossi

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

@@ -125,8 +125,7 @@ def lookup(key, options):
     """
     context = context_module.get_context()
     use_datastore = context._use_datastore(key, options)
-    in_transaction = bool(_get_transaction(options))
-    if use_datastore and in_transaction:
+    if use_datastore and options.transaction:
         use_global_cache = False
     else:
         use_global_cache = context._use_global_cache(key, options)
@@ -316,8 +315,7 @@ def get_read_options(options, default_read_consistency=None):
         ValueError: When ``read_consistency`` is set to ``EVENTUAL`` and there
             is a transaction.
     """
-    transaction = _get_transaction(options)
-
+    transaction = options.transaction
     read_consistency = options.read_consistency
 
     if transaction is None:
@@ -332,27 +330,6 @@ def get_read_options(options, default_read_consistency=None):
     )
 
 
-def _get_transaction(options):
-    """Get the transaction for a request.
-
-    If specified, this will return the transaction from ``options``. Otherwise,
-    it will return the transaction for the current context.
-
-    Args:
-        options (_options.ReadOptions): The options for the request. Only
-            ``transaction`` will have any bearing here.
-
-    Returns:
-        Union[bytes, NoneType]: The transaction identifier, or :data:`None`.
-    """
-    transaction = getattr(options, "transaction", None)
-    if transaction is None:
-        context = context_module.get_context()
-        transaction = context.transaction
-
-    return transaction
-
-
 @tasklets.tasklet
 def put(entity, options):
     """Store an entity in datastore.
@@ -388,7 +365,7 @@ def put(entity, options):
             yield _cache.global_set(cache_key, cache_value, expires=expires)
 
     if use_datastore:
-        transaction = _get_transaction(options)
+        transaction = context.transaction
         if transaction:
             batch = _get_commit_batch(transaction, options)
         else:
@@ -432,7 +409,7 @@ def delete(key, options):
         if use_global_cache:
             yield _cache.global_lock(cache_key)
 
-        transaction = _get_transaction(options)
+        transaction = context.transaction
         if transaction:
             batch = _get_commit_batch(transaction, options)
         else:

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

@@ -399,3 +399,59 @@ def run(self):
         thread2.join()
 
         assert activity["calls"] == 2
+
+
+@pytest.mark.usefixtures("client_context")
+def test_non_transactional_means_no_transaction(dispose_of):
+    """Regression test for #552
+
+    https://github.com/googleapis/python-ndb/issues/552
+    """
+    N = 50
+
+    class SomeKind(ndb.Model):
+        pass
+
+    class OtherKind(ndb.Model):
+        pass
+
+    @ndb.tasklet
+    def create_entities():
+        parent_keys = yield [SomeKind().put_async() for _ in range(N)]
+
+        futures = []
+        for parent_key in parent_keys:
+            dispose_of(parent_key._key)
+            futures.append(OtherKind(parent=parent_key).put_async())
+            futures.append(OtherKind(parent=parent_key).put_async())
+
+        keys = yield futures
+        for key in keys:
+            dispose_of(key._key)
+
+        raise ndb.Return(keys)
+
+    @ndb.non_transactional()
+    @ndb.tasklet
+    def non_transactional_tasklet(keys):
+        entities = yield ndb.get_multi_async(keys)
+        raise ndb.Return(entities)
+
+    @ndb.non_transactional()
+    @ndb.tasklet
+    def also_a_non_transactional_tasklet():
+        entities = yield OtherKind.query().fetch_async()
+        raise ndb.Return(entities)
+
+    @ndb.transactional()
+    def test_lookup(keys):
+        entities = non_transactional_tasklet(keys).result()
+        assert len(entities) == N * 2
+
+    @ndb.transactional()
+    def test_query():
+        return also_a_non_transactional_tasklet().result()
+
+    keys = create_entities().result()
+    test_lookup(keys)
+    eventually(test_query, length_equals(N * 2))