def test_ndb_context():
assert _runstate.states.current() is None
with _runstate.ndb_context():
one = _runstate.current()
with _runstate.ndb_context():
two = _runstate.current()
assert one is not two
two.eventloop = unittest.mock.Mock(spec=("run", ))
two.eventloop.run.assert_not_called()
assert _runstate.current() is one
two.eventloop.run.assert_called_once_with()
assert _runstate.states.current() is None
def _get_read_options(options):
"""Get the read options for a request.
Args:
options (Dict[str, Any]): The options for the request. For example,
``{"read_consistency": EVENTUAL}``. May contain options unrelated
to creating a :class:`datastore_pb2.ReadOptions` instance, which
will be ignored.
Returns:
datastore_pb2.ReadOptions: The options instance for passing to the
Datastore gRPC API.
Raises:
ValueError: When ``read_consistency`` is set to ``EVENTUAL`` and there
is a transaction.
"""
state = _runstate.current()
transaction = options.get("transaction", state.transaction)
read_consistency = options.get("read_consistency")
if read_consistency is None:
read_consistency = options.get("read_policy") # Legacy NDB
if transaction is not None and read_consistency is EVENTUAL:
raise ValueError(
"read_consistency must be EVENTUAL when in transaction")
return datastore_pb2.ReadOptions(read_consistency=read_consistency,
transaction=transaction)
def _get_batch(batch_cls, options):
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current run state. 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 (Dict[str, Any]): The options for the request. For example,
``{"read_consistency": EVENTUAL}``. Calls with different options
will be placed in different batches.
Returns:
batch_cls: An instance of the batch class.
"""
state = _runstate.current()
batches = state.batches.get(batch_cls)
if batches is None:
state.batches[batch_cls] = batches = {}
options_key = tuple(sorted(options.items()))
batch = batches.get(options_key)
if batch is not None:
return batch
def idle():
batch = batches.pop(options_key)
batch.idle_callback()
batches[options_key] = batch = _LookupBatch(options)
_eventloop.add_idle(idle)
return batch
def _datastore_commit(mutations, transaction):
"""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.
Returns:
grpc.Future: 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 = _runstate.current().client
request = datastore_pb2.CommitRequest(
project_id=client.project,
mode=mode,
mutations=mutations,
transaction=transaction,
)
api = stub()
return api.Commit.future(request)
def test_ndb_context():
assert _runstate.states.current() is None
with _runstate.ndb_context():
one = _runstate.current()
with _runstate.ndb_context():
two = _runstate.current()
assert one is not two
two.eventloop = unittest.mock.Mock(spec=("run",))
two.eventloop.run.assert_not_called()
assert _runstate.current() is one
two.eventloop.run.assert_called_once_with()
assert _runstate.states.current() is None
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.
"""
state = _runstate.current()
return state.stub
def test_state_context():
assert _runstate.states.current() is None
client1 = object()
client2 = object()
with _runstate.state_context(client1):
one = _runstate.current()
assert one.client is client1
with _runstate.state_context(client2):
two = _runstate.current()
assert two.client is client2
assert one is not two
two.eventloop = unittest.mock.Mock(spec=("run", ))
two.eventloop.run.assert_not_called()
assert _runstate.current() is one
two.eventloop.run.assert_called_once_with()
assert _runstate.states.current() is None
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 = _runstate.current()
return context.eventloop
def _perform_batch_lookup():
"""Perform a Datastore Lookup on all batched Lookup requests.
Meant to be used as an idle callback, so that calls to lookup entities can
be batched into a single request to the back end service as soon as running
code has need of one of the results.
"""
state = _runstate.current()
batch = state.batches.pop(_BATCH_LOOKUP, None)
if batch is None:
return
rpc = _datastore_lookup(batch.keys())
_eventloop.queue_rpc(rpc, BatchLookupCallback(batch))
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`.
"""
state = _runstate.current()
return options.get("transaction", state.transaction)
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.
"""
state = _runstate.current()
# Be lazy and avoid circular dependency with _runstate
if state.eventloop is None:
state.eventloop = EventLoop()
return state.eventloop
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.
"""
state = _runstate.current()
# Be lazy and avoid circular dependency with _runstate
if state.eventloop is None:
state.eventloop = EventLoop()
return state.eventloop
def _datastore_lookup(keys):
"""Issue a Lookup call to Datastore using gRPC.
Args:
keys (Iterable[datastore_v1.proto.entity_pb2.Key]): The entity keys to look up.
Returns:
:class:`grpc.Future`: Future object for eventual result of lookup.
"""
client = _runstate.current().client
request = datastore_pb2.LookupRequest(
project_id=client.project, keys=[key for key in keys]
)
api = stub()
return api.Lookup.future(request)
def _get_lookup_batch():
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current run state. 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.
Returns:
Dict[~datastore_v1.proto.entity_pb2.Key, List[~tasklets.Future]]
"""
state = _runstate.current()
batch = state.batches.get(_BATCH_LOOKUP)
if batch is not None:
return batch
state.batches[_BATCH_LOOKUP] = batch = {}
_eventloop.add_idle(_perform_batch_lookup)
return batch
def _datastore_lookup(keys, read_options):
"""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.
Returns:
:class:`grpc.Future`: Future object for eventual result of lookup.
"""
client = _runstate.current().client
request = datastore_pb2.LookupRequest(
project_id=client.project,
keys=[key for key in keys],
read_options=read_options,
)
api = stub()
return api.Lookup.future(request)
def _project_from_app(app, allow_empty=False):
"""Convert a legacy Google App Engine app string to a project.
Args:
app (str): The application value to be used. If the caller passes
:data:`None` and ``allow_empty`` is :data:`False`, then this will
use the project set by the current client context. (See
:meth:`~client.Client.context`.)
allow_empty (bool): Flag determining if an empty (i.e. :data:`None`)
project is allowed. Defaults to :data:`False`.
Returns:
str: The cleaned project.
"""
if app is None:
if allow_empty:
return None
client = _runstate.current().client
app = client.project
# NOTE: This is the same behavior as in the helper
# ``google.cloud.datastore.key._clean_app()``.
parts = app.split("~", 1)
return parts[-1]
def stub():
"""Get the stub for the `Google Datastore` API.
Gets the stub from the current context, creating one if there isn't one
already.
Returns:
:class:`~google.cloud.datastore_v1.proto.datastore_pb2_grpc.DatastoreStub`:
The stub instance.
"""
state = _runstate.current()
if state.stub is None:
client = state.client
if client.secure:
channel = _helpers.make_secure_channel(client._credentials,
_http.DEFAULT_USER_AGENT,
client.host)
else:
channel = grpc.insecure_channel(client.host)
state.stub = datastore_pb2_grpc.DatastoreStub(channel)
return state.stub
def test_assert_as_context_manager(self):
context = self._make_one()
with context:
assert _runstate.current() is context
with pytest.raises(exceptions.ContextError):
_runstate.current()
