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