def _get_commit_batch(transaction, options):
"""Get the commit batch for the current context and transaction.
Args:
transaction (bytes): The transaction id. Different transactions will
have different batchs.
options (_options.Options): Options for the batch. Not supported at
this time.
Returns:
_TransactionalCommitBatch: The batch.
"""
# Support for different options will be tricky if we're in a transaction,
# since we can only do one commit, so any options that affect that gRPC
# call would all need to be identical. For now, no options are supported
# here.
for key, value in options.items():
if value:
raise NotImplementedError("Passed bad option: {!r}".format(key))
# Since we're in a transaction, we need to hang on to the batch until
# commit time, so we need to store it separately from other batches.
context = context_module.get_context()
batch = context.commit_batches.get(transaction)
if batch is None:
batch = _TransactionalCommitBatch(transaction, options)
context.commit_batches[transaction] = batch
return batch
def tasklet_wrapper(*args, **kwargs):
# The normal case is that the wrapped function is a generator function
# that returns a generator when called. We also support the case that
# the user has wrapped a regular function with the tasklet decorator.
# In this case, we fail to realize an actual tasklet, but we go ahead
# and create a future object and set the result to the function's
# return value so that from the user perspective there is no problem.
# This permissive behavior is inherited from legacy NDB.
context = context_module.get_context()
try:
returned = wrapped(*args, **kwargs)
except StopIteration as stop:
# If wrapped is a regular function and the function uses the
# deprecated "raise Return(result)" pattern rather than just
# returning the result, then we'll extract the result from the
# StopIteration exception.
returned = _get_return_value(stop)
if isinstance(returned, types.GeneratorType):
# We have a tasklet, start it
future = _TaskletFuture(returned, context, info=wrapped.__name__)
future._advance_tasklet()
else:
# We don't have a tasklet, but we fake it anyway
future = Future(info=wrapped.__name__)
future.set_result(returned)
return future
def _datastore_commit(mutations, transaction, retries=None, timeout=None):
"""Call Commit on Datastore.
Args:
mutations (List[datastore_pb2.Mutation]): The changes to persist to
Datastore.
transaction (Union[bytes, NoneType]): The identifier for the
transaction for this commit, or :data:`None` if no transaction is
being used.
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use :data:`_retry._DEFAULT_RETRIES`.
If :data:`0` is passed, the call is attempted only once.
timeout (float): Timeout, in seconds, to pass to gRPC call. If
:data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
Returns:
tasklets.Tasklet: A future for
:class:`google.cloud.datastore_v1.datastore_pb2.CommitResponse`
"""
if transaction is None:
mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
else:
mode = datastore_pb2.CommitRequest.TRANSACTIONAL
client = context_module.get_context().client
request = datastore_pb2.CommitRequest(
project_id=client.project,
mode=mode,
mutations=mutations,
transaction=transaction,
)
return make_call("Commit", request, retries=retries, timeout=timeout)
def _datastore_lookup(keys, read_options, retries=None, timeout=None):
"""Issue a Lookup call to Datastore using gRPC.
Args:
keys (Iterable[entity_pb2.Key]): The entity keys to
look up.
read_options (Union[datastore_pb2.ReadOptions, NoneType]): Options for
the request.
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use :data:`_retry._DEFAULT_RETRIES`.
If :data:`0` is passed, the call is attempted only once.
timeout (float): Timeout, in seconds, to pass to gRPC call. If
:data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
Returns:
tasklets.Future: Future object for eventual result of lookup.
"""
client = context_module.get_context().client
request = datastore_pb2.LookupRequest(
project_id=client.project,
keys=[key for key in keys],
read_options=read_options,
)
return make_call("Lookup", request, retries=retries, timeout=timeout)
def delete(key, options):
"""Delete an entity from Datastore.
Deleting an entity that doesn't exist does not result in an error. The
result is the same regardless.
Args:
key (datastore.Key): The key for the entity to be deleted.
options (_options.Options): Options for this request.
Returns:
tasklets.Future: Will be finished when entity is deleted. Result will
always be :data:`None`.
"""
context = context_module.get_context()
use_global_cache = context._use_global_cache(key, options)
use_datastore = context._use_datastore(key, options)
if use_global_cache:
cache_key = _cache.global_cache_key(key)
if use_datastore:
if use_global_cache:
yield _cache.global_lock(cache_key)
transaction = _get_transaction(options)
if transaction:
batch = _get_commit_batch(transaction, options)
else:
batch = _batch.get_batch(_NonTransactionalCommitBatch, options)
yield batch.delete(key)
if use_global_cache:
yield _cache.global_delete(cache_key)
def test_success(_datastore_api):
context_module.get_context().cache["foo"] = "bar"
on_commit_callback = mock.Mock()
def callback():
context = context_module.get_context()
assert not context.cache
context.call_on_commit(on_commit_callback)
return "I tried, momma."
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
commit_future = tasklets.Future("commit transaction")
_datastore_api.commit.return_value = commit_future
future = _transaction.transaction_async(callback)
_datastore_api.begin_transaction.assert_called_once_with(
False, retries=0
)
begin_future.set_result(b"tx123")
_datastore_api.commit.assert_called_once_with(b"tx123", retries=0)
commit_future.set_result(None)
assert future.result() == "I tried, momma."
on_commit_callback.assert_called_once_with()
def _datastore_begin_transaction(read_only, retries=None, timeout=None):
"""Calls ``BeginTransaction`` on Datastore.
Args:
read_only (bool): Whether to start a read-only or read-write
transaction.
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use :data:`_retry._DEFAULT_RETRIES`.
If :data:`0` is passed, the call is attempted only once.
timeout (float): Timeout, in seconds, to pass to gRPC call. If
:data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
Returns:
tasklets.Tasklet: A future for
:class:`google.cloud.datastore_v1.datastore_pb2.BeginTransactionResponse`
"""
client = context_module.get_context().client
if read_only:
options = datastore_pb2.TransactionOptions(
read_only=datastore_pb2.TransactionOptions.ReadOnly())
else:
options = datastore_pb2.TransactionOptions(
read_write=datastore_pb2.TransactionOptions.ReadWrite())
request = datastore_pb2.BeginTransactionRequest(
project_id=client.project, transaction_options=options)
return make_call("BeginTransaction",
request,
retries=retries,
timeout=timeout)
def test__context():
with patch_credentials("testing"):
client = client_module.Client()
with client.context():
context = context_module.get_context()
assert context.client is client
def put(entity, options):
"""Store an entity in datastore.
The entity can be a new entity to be saved for the first time or an
existing entity that has been updated.
Args:
entity_pb (datastore.Entity): The entity to be stored.
options (_options.Options): Options for this request.
Returns:
tasklets.Future: Result will be completed datastore key
(datastore.Key) for the entity.
"""
context = context_module.get_context()
use_global_cache = context._use_global_cache(entity.key, options)
use_datastore = context._use_datastore(entity.key, options)
if not (use_global_cache or use_datastore):
raise TypeError("use_global_cache and use_datastore can't both be False")
if not use_datastore and entity.key.is_partial:
raise TypeError("Can't store partial keys when use_datastore is False")
lock = None
entity_pb = helpers.entity_to_protobuf(entity)
cache_key = _cache.global_cache_key(entity.key)
if use_global_cache and not entity.key.is_partial:
if use_datastore:
lock = yield _cache.global_lock_for_write(cache_key)
else:
expires = context._global_cache_timeout(entity.key, options)
cache_value = entity_pb.SerializeToString()
yield _cache.global_set(cache_key, cache_value, expires=expires)
if use_datastore:
transaction = context.transaction
if transaction:
batch = _get_commit_batch(transaction, options)
else:
batch = _batch.get_batch(_NonTransactionalCommitBatch, options)
key_pb = yield batch.put(entity_pb)
if key_pb:
key = helpers.key_from_protobuf(key_pb)
else:
key = None
if lock:
if transaction:
def callback():
_cache.global_unlock_for_write(cache_key, lock).result()
context.call_on_transaction_complete(callback)
else:
yield _cache.global_unlock_for_write(cache_key, lock)
raise tasklets.Return(key)
def in_transaction():
"""Determine if there is a currently active transaction.
Returns:
bool: :data:`True` if there is a transaction for the current context,
otherwise :data:`False`.
"""
return context_module.get_context().transaction is not None
def lookup(key, options):
"""Look up a Datastore entity.
Gets an entity from Datastore, asynchronously. Checks the global cache,
first, if appropriate. Uses batching.
Args:
key (~datastore.Key): The key for the entity to retrieve.
options (_options.ReadOptions): The options for the request. For
example, ``{"read_consistency": EVENTUAL}``.
Returns:
:class:`~tasklets.Future`: If not an exception, future's result will be
either an entity protocol buffer or _NOT_FOUND.
"""
context = context_module.get_context()
use_datastore = context._use_datastore(key, options)
if use_datastore and options.transaction:
use_global_cache = False
else:
use_global_cache = context._use_global_cache(key, options)
if not (use_global_cache or use_datastore):
raise TypeError(
"use_global_cache and use_datastore can't both be False")
entity_pb = _NOT_FOUND
key_locked = False
if use_global_cache:
cache_key = _cache.global_cache_key(key)
result = yield _cache.global_get(cache_key)
key_locked = _cache.is_locked_value(result)
if not key_locked:
if result is not None:
entity_pb = entity_pb2.Entity()
entity_pb.MergeFromString(result)
elif use_datastore:
yield _cache.global_lock(cache_key, read=True)
yield _cache.global_watch(cache_key)
if entity_pb is _NOT_FOUND and use_datastore:
batch = _batch.get_batch(_LookupBatch, options)
entity_pb = yield batch.add(key)
# Do not cache misses
if use_global_cache and not key_locked:
if entity_pb is not _NOT_FOUND:
expires = context._global_cache_timeout(key, options)
serialized = entity_pb.SerializeToString()
yield _cache.global_compare_and_swap(cache_key,
serialized,
expires=expires)
else:
yield _cache.global_unwatch(cache_key)
raise tasklets.Return(entity_pb)
def _advance_tasklet(self, send_value=None, error=None):
"""Advance a tasklet one step by sending in a value or error."""
try:
with self.context.use():
# Send the next value or exception into the generator
if error:
self.generator.throw(type(error), error)
# send_value will be None if this is the first time
yielded = self.generator.send(send_value)
# Context may have changed in tasklet
self.context = context_module.get_context()
except StopIteration as stop:
# Generator has signalled exit, get the return value. This tasklet
# has finished.
self.set_result(_get_return_value(stop))
return
except Exception as error:
# An error has occurred in the tasklet. This tasklet has finished.
self.set_exception(error)
return
# This tasklet has yielded a value. We expect this to be a future
# object (either NDB or gRPC) or a sequence of futures, in the case of
# parallel yield.
def done_callback(yielded):
# To be called when a future dependency has completed. Advance the
# tasklet with the yielded value or error.
#
# It might be worth noting that legacy NDB added a callback to the
# event loop which, in turn, called _help_tasklet_along. I don't
# see a compelling reason not to go ahead and call _advance_tasklet
# immediately here, rather than queue it up to be called soon by
# the event loop. This is subject to change if the reason for the
# indirection in the original implementation becomes apparent.
error = yielded.exception()
if error:
self._advance_tasklet(error=error)
else:
self._advance_tasklet(yielded.result())
if isinstance(yielded, Future):
yielded.add_done_callback(done_callback)
elif isinstance(yielded, _remote.RemoteCall):
_eventloop.queue_rpc(yielded, done_callback)
elif isinstance(yielded, (list, tuple)):
future = _MultiFuture(yielded)
future.add_done_callback(done_callback)
else:
raise RuntimeError(
"A tasklet yielded an illegal value: {!r}".format(yielded))
def wrap_endpoint(*args, **kwargs):
"""Wraps the endpoint method in a NDB Context."""
context = context_module.get_context(raise_context_error=False)
if not context:
creds, project = auth.default()
with ndb.Client(project=project, credentials=creds).context():
return method(*args, **kwargs)
# If endpoint is inside a NDB context don't create a new context.
return method(*args, **kwargs)
def callback():
# The transaction uses its own in-memory cache, which should be empty in
# the transaction context and not include the key set above.
context = context_module.get_context()
assert not context.cache
context.call_on_commit(on_commit_callback)
context.call_on_transaction_complete(transaction_complete_callback)
return "I tried, momma."
def fetch(query):
"""Fetch query results.
Args:
query (query.QueryOptions): The query spec.
Returns:
tasklets.Future: Result is List[model.Model]: The query results.
"""
client = context_module.get_context().client
project_id = query.project
if not project_id:
project_id = client.project
namespace = query.namespace
if not namespace:
namespace = client.namespace
filter_pbs = (None, )
if query.filters:
filter_pbs = query.filters._to_filter()
if not isinstance(filter_pbs, (tuple, list)):
filter_pbs = (filter_pbs, )
multiple_queries = len(filter_pbs) > 1
if multiple_queries:
# If we're aggregating multiple queries, then limit and offset will be
# have to applied to the aggregate, not passed to Datastore to use on
# individual queries
offset = query.offset
limit = query.limit
query = query.copy(offset=0, limit=None)
else:
offset = limit = None
queries = [
_run_query(project_id, namespace, _query_to_protobuf(query, filter_pb))
for filter_pb in filter_pbs
]
result_sets = yield queries
result_sets = [(_Result(result_type, result_pb, query.order_by)
for result_type, result_pb in result_set)
for result_set in result_sets]
if len(result_sets) > 1:
sortable = bool(query.order_by)
results = _merge_results(result_sets, sortable)
else:
results = result_sets[0]
if offset or limit:
results = itertools.islice(results, offset, offset + limit)
return [result.entity(query.projection) for result in results]
def _advance_tasklet(self, send_value=None, error=None):
"""Advance a tasklet one step by sending in a value or error."""
try:
with self.context.use():
# Send the next value or exception into the generator
if error:
self.generator.throw(type(error), error)
# send_value will be None if this is the first time
yielded = self.generator.send(send_value)
# Context may have changed in tasklet
self.context = context_module.get_context()
except StopIteration as stop:
# Generator has signalled exit, get the return value. This tasklet
# has finished.
self.set_result(_get_return_value(stop))
return
except Exception as error:
# An error has occurred in the tasklet. This tasklet has finished.
self.set_exception(error)
return
# This tasklet has yielded a value. We expect this to be a future
# object (either NDB or gRPC) or a sequence of futures, in the case of
# parallel yield.
def done_callback(yielded):
# To be called when a future dependency has completed. Advance the
# tasklet with the yielded value or error.
#
# It was tempting to call `_advance_tasklet` (`_help_tasklet_along`
# in Legacy) directly. Doing so, it has been found, can lead to
# exceeding the maximum recursion depth. Queing it up to run on the
# event loop avoids this issue by keeping the call stack shallow.
error = yielded.exception()
if error:
_eventloop.call_soon(self._advance_tasklet, error=error)
else:
_eventloop.call_soon(self._advance_tasklet, yielded.result())
if isinstance(yielded, Future):
yielded.add_done_callback(done_callback)
elif isinstance(yielded, _remote.RemoteCall):
_eventloop.queue_rpc(yielded, done_callback)
elif isinstance(yielded, (list, tuple)):
future = _MultiFuture(yielded)
future.add_done_callback(done_callback)
else:
raise RuntimeError(
"A tasklet yielded an illegal value: {!r}".format(yielded))
def _key_function(values):
if not len(values) % 2:
context = context_module.get_context()
client = context.client
return key.Key(*values,
namespace=context.get_namespace(),
project=client.project)
_raise_cast_error(
"Key requires even number of operands or single string, {}".format(
values))
def stub():
"""Get the stub for the `Google Datastore` API.
Gets the stub from the current context.
Returns:
:class:`~google.cloud.datastore_v1.proto.datastore_pb2_grpc.DatastoreStub`:
The stub instance.
"""
context = context_module.get_context()
return context.stub
def get_event_loop():
"""Get the current event loop.
This function should be called within a context established by
:func:`~google.cloud.ndb.ndb_context`.
Returns:
EventLoop: The event loop for the current context.
"""
context = context_module.get_context()
return context.eventloop
def in_transaction():
"""Determine if there is a currently active transaction.
Returns:
bool: :data:`True` if there is a transaction for the current context,
otherwise :data:`False`.
"""
# Avoid circular import in Python 2.7
from google.cloud.ndb import context as context_module
return context_module.get_context().transaction is not None
def _handle_independent(self):
"""Always use a new transaction, "pausing" any existing transactions.
A function that uses this policy should not return any entities read in
the new transaction, as the entities are not transactionally consistent
with the caller's transaction.
"""
if in_transaction():
# Avoid circular import in Python 2.7
from google.cloud.ndb import context as context_module
context = context_module.get_context()
new_context = context.new(transaction=None)
return new_context
def retry_wrapper(*args, **kwargs):
from google.cloud.ndb import context as context_module
sleep_generator = core_retry.exponential_sleep_generator(
_DEFAULT_INITIAL_DELAY,
_DEFAULT_MAXIMUM_DELAY,
_DEFAULT_DELAY_MULTIPLIER,
)
for sleep_time in itertools.islice(sleep_generator, retries + 1):
context = context_module.get_context()
if not context.in_retry():
# We need to be able to identify if we are inside a nested
# retry. Here, we set the retry state in the context. This is
# used for deciding if an exception should be raised
# immediately or passed up to the outer retry block.
context.set_retry_state(repr(callback))
try:
result = callback(*args, **kwargs)
if isinstance(result, tasklets.Future):
result = yield result
except exceptions.NestedRetryException as e:
error = e
except Exception as e:
# `e` is removed from locals at end of block
error = e # See: https://goo.gl/5J8BMK
if not is_transient_error(error):
# If we are in an inner retry block, use special nested
# retry exception to bubble up to outer retry. Else, raise
# actual exception.
if context.get_retry_state() != repr(callback):
message = getattr(error, "message", str(error))
raise exceptions.NestedRetryException(message)
else:
raise error
else:
raise tasklets.Return(result)
finally:
# No matter what, if we are exiting the top level retry,
# clear the retry state in the context.
if context.get_retry_state() == repr(
callback): # pragma: NO BRANCH
context.clear_retry_state()
yield tasklets.sleep(sleep_time)
raise core_exceptions.RetryError(
"Maximum number of {} retries exceeded while calling {}".format(
retries, callback),
cause=error,
)
def non_transactional_inner_wrapper(*args, **kwargs):
# Avoid circular import in Python 2.7
from google.cloud.ndb import context as context_module
context = context_module.get_context()
if not context.in_transaction():
return wrapped(*args, **kwargs)
if not allow_existing:
raise exceptions.BadRequestError(
"{} cannot be called within a transaction".format(
wrapped.__name__))
new_context = context.new(transaction=None)
with new_context.use():
return wrapped(*args, **kwargs)
def get_event_loop():
"""Get the current event loop.
This function should be called within a context established by
:func:`~google.cloud.ndb.ndb_context`.
Returns:
EventLoop: The event loop for the current context.
"""
# Prevent circular import in Python 2.7
from google.cloud.ndb import context as context_module
context = context_module.get_context()
return context.eventloop
def _next_batch(self):
"""Get the next batch from Datastore.
If this batch isn't the last batch for the query, update the internal
query spec with a cursor pointing to the next batch.
"""
query = self._query
response = yield _datastore_run_query(query)
batch = response.batch
result_type = batch.entity_result_type
self._start_cursor = query.start_cursor
self._index = 0
self._batch = [
_Result(result_type, result_pb, query.order_by)
for result_pb in response.batch.entity_results
]
if result_type == RESULT_TYPE_FULL:
# If we cached a delete, remove it from the result set. This may come cause
# some queries to return less than their limit even if there are more
# results. As far as I can tell, that was also a possibility with the legacy
# version.
context = context_module.get_context()
self._batch = [
result
for result in self._batch
if result.check_cache(context) is not None
]
self._has_next_batch = more_results = batch.more_results == NOT_FINISHED
self._more_results_after_limit = batch.more_results == MORE_RESULTS_AFTER_LIMIT
if more_results:
# Fix up query for next batch
limit = self._query.limit
if limit is not None:
limit -= len(self._batch)
offset = self._query.offset
if offset:
offset -= response.batch.skipped_results
self._query = self._query.copy(
start_cursor=Cursor(batch.end_cursor),
offset=offset,
limit=limit,
)
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 (Dict[str, Any]): The options for the request. Only
``transaction`` will have any bearing here.
Returns:
Union[bytes, NoneType]: The transaction identifier, or :data:`None`.
"""
context = context_module.get_context()
return options.get("transaction", context.transaction)
def test_entity_full_entity_no_cache(model):
context = context_module.get_context()
with context.new(cache_policy=False).use():
key_pb = entity_pb2.Key(
partition_id=entity_pb2.PartitionId(project_id="testing"),
path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
)
entity = mock.Mock(key=key_pb)
model._entity_from_protobuf.return_value = entity
result = _datastore_query._Result(
_datastore_query.RESULT_TYPE_FULL,
mock.Mock(entity=entity,
cursor=b"123",
spec=("entity", "cursor")),
)
assert result.entity() is entity
def get_batch(batch_cls, options=None):
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current context. If there is
not already a batch started, a new structure is created and an idle
callback is added to the current event loop which will eventually perform
the batch look up.
Args:
batch_cls (type): Class representing the kind of operation being
batched.
options (_options.ReadOptions): The options for the request. Calls with
different options will be placed in different batches.
Returns:
batch_cls: An instance of the batch class.
"""
# prevent circular import in Python 2.7
from google.cloud.ndb import context as context_module
context = context_module.get_context()
batches = context.batches.get(batch_cls)
if batches is None:
context.batches[batch_cls] = batches = {}
if options is not None:
options_key = tuple(
sorted(((key, value) for key, value in options.items()
if value is not None)))
else:
options_key = ()
batch = batches.get(options_key)
if batch is not None and not batch.full():
return batch
def idler(batch):
def idle():
if batches.get(options_key) is batch:
del batches[options_key]
batch.idle_callback()
return idle
batches[options_key] = batch = batch_cls(options)
_eventloop.add_idle(idler(batch))
return batch
def _datastore_rollback(transaction, retries=None):
"""Calls Rollback in Datastore.
Args:
transaction (bytes): Transaction id.
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use :data:`_retry._DEFAULT_RETRIES`.
If :data:`0` is passed, the call is attempted only once.
Returns:
tasklets.Tasklet: Future for
:class:`google.cloud.datastore_v1.datastore_pb2.RollbackResponse`
"""
client = context_module.get_context().client
request = datastore_pb2.RollbackRequest(project_id=client.project,
transaction=transaction)
return make_call("Rollback", request, retries=retries)
def test_entity_full_entity_cached(model):
key = key_module.Key("ThisKind", 42)
key_pb = entity_pb2.Key(
partition_id=entity_pb2.PartitionId(project_id="testing"),
path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
)
entity = mock.Mock(key=key_pb)
cached_entity = mock.Mock(key=key_pb, _key=key)
context = context_module.get_context()
context.cache[key] = cached_entity
model._entity_from_protobuf.return_value = entity
result = _datastore_query._Result(
_datastore_query.RESULT_TYPE_FULL,
mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")),
)
assert result.entity() is not entity
assert result.entity() is cached_entity