def _transaction_async(context, callback, read_only=False):
# Avoid circular import in Python 2.7
from google.cloud.ndb import _datastore_api
# Start the transaction
utils.logging_debug(log, "Start transaction")
transaction_id = yield _datastore_api.begin_transaction(read_only,
retries=0)
utils.logging_debug(log, "Transaction Id: {}", transaction_id)
on_commit_callbacks = []
tx_context = context.new(
transaction=transaction_id,
on_commit_callbacks=on_commit_callbacks,
batches=None,
commit_batches=None,
cache=None,
# We could just pass `None` here and let the `Context` constructor
# instantiate a new event loop, but our unit tests inject a subclass of
# `EventLoop` that makes testing a little easier. This makes sure the
# new event loop is of the same type as the current one, to propagate
# the event loop class used for testing.
eventloop=type(context.eventloop)(),
retry=context.get_retry_state(),
)
# The outer loop is dependent on the inner loop
def run_inner_loop(inner_context):
with inner_context.use():
if inner_context.eventloop.run1():
return True # schedule again
context.eventloop.add_idle(run_inner_loop, tx_context)
with tx_context.use():
try:
# Run the callback
result = callback()
if isinstance(result, tasklets.Future):
result = yield result
# Make sure we've run everything we can run before calling commit
_datastore_api.prepare_to_commit(transaction_id)
tx_context.eventloop.run()
# Commit the transaction
yield _datastore_api.commit(transaction_id, retries=0)
# Rollback if there is an error
except Exception as e: # noqa: E722
tx_context.cache.clear()
yield _datastore_api.rollback(transaction_id)
raise e
tx_context._clear_global_cache()
for callback in on_commit_callbacks:
callback()
raise tasklets.Return(result)
def test_prepare_to_commit(get_commit_batch):
_api.prepare_to_commit(b"123")
get_commit_batch.assert_called_once_with(b"123", _options.Options())
batch = get_commit_batch.return_value
assert batch.preparing_to_commit is True