def test_it(self):
from google.cloud.datastore_v1.proto import datastore_pb2
from google.cloud.datastore_v1.proto import entity_pb2
index_updates = 1337
keys = [
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(
kind='Foo',
id=1234,
),
], ),
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(
kind='Bar',
name='baz',
),
], ),
]
response = datastore_pb2.CommitResponse(
mutation_results=[
datastore_pb2.MutationResult(key=key) for key in keys
],
index_updates=index_updates,
)
result = self._call_fut(response)
self.assertEqual(result, (index_updates, keys))
def test_idle_callback_success(datastore_allocate_ids, in_context):
def Mutation():
path = [entity_pb2.Key.PathElement(kind="SomeKind")]
return datastore_pb2.Mutation(upsert=entity_pb2.Entity(
key=entity_pb2.Key(path=path)))
mutation1, mutation2 = Mutation(), Mutation()
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.incomplete_mutations = [mutation1, mutation2]
future1, future2 = tasklets.Future(), tasklets.Future()
batch.incomplete_futures = [future1, future2]
rpc = tasklets.Future("_datastore_allocate_ids")
datastore_allocate_ids.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with in_context.new(eventloop=eventloop).use():
batch.idle_callback()
rpc.set_result(
mock.Mock(keys=[
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(kind="SomeKind", id=1)
]),
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(kind="SomeKind", id=2)
]),
]))
allocating_ids = batch.allocating_ids[0]
assert future1.result().path[0].id == 1
assert mutation1.upsert.key.path[0].id == 1
assert future2.result().path[0].id == 2
assert mutation2.upsert.key.path[0].id == 2
assert allocating_ids.result() is None
def _make_key(id_):
from google.cloud.datastore_v1.proto import entity_pb2
key = entity_pb2.Key()
elem = key.path.add()
elem.id = id_
return key
def _make_entity(kind, id_, project):
from google.cloud.datastore_v1.proto import entity_pb2
key = entity_pb2.Key()
key.partition_id.project_id = project
elem = key.path.add()
elem.kind = kind
elem.id = id_
return entity_pb2.Entity(key=key)
def _make_mutation(id_):
from google.cloud.datastore_v1.proto import datastore_pb2
from google.cloud.datastore_v1.proto import entity_pb2
key = entity_pb2.Key()
key.partition_id.project_id = 'PROJECT'
elem = key.path.add()
elem.kind = 'Kind'
elem.id = id_
return datastore_pb2.MutationResult(key=key)
def idle_callback(self):
"""Perform a Datastore Lookup on all batched Lookup requests."""
keys = []
for todo_key in self.todo.keys():
key_pb = entity_pb2.Key()
key_pb.ParseFromString(todo_key)
keys.append(key_pb)
read_options = _get_read_options(self.options)
rpc = _datastore_lookup(keys, read_options)
_eventloop.queue_rpc(rpc, self.lookup_callback)
def idle_callback(self):
"""Perform a Datastore Lookup on all batched Lookup requests."""
keys = []
for todo_key in self.todo.keys():
key_pb = entity_pb2.Key()
key_pb.ParseFromString(todo_key)
keys.append(key_pb)
read_options = _get_read_options(self.options)
retries = self.options.get("retries")
rpc = _datastore_lookup(keys, read_options, retries=retries)
rpc.add_done_callback(self.lookup_callback)
def to_protobuf(self):
from google.cloud.datastore_v1.proto import entity_pb2
key = self._key = entity_pb2.Key()
# Don't assign it, because it will just get ripped out
# key.partition_id.project_id = self.project
element = key.path.add()
element.kind = self._kind
if self._id is not None:
element.id = self._id
return key
def test_entity_key_only():
key_pb = entity_pb2.Key(
partition_id=entity_pb2.PartitionId(project_id="testing"),
path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
)
result = _datastore_query._Result(
_datastore_query.RESULT_TYPE_KEY_ONLY,
mock.Mock(
entity=mock.Mock(key=key_pb, spec=("key",)),
cursor=b"123",
spec=("entity", "cursor"),
),
)
assert result.entity() == key_module.Key("ThisKind", 42)
def test_entity_full_entity(model):
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
model._entity_from_protobuf.assert_called_once_with(entity)
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 _makePB(self, project=None, namespace=None, path=()):
from google.cloud.datastore_v1.proto import entity_pb2
pb = entity_pb2.Key()
if project is not None:
pb.partition_id.project_id = project
if namespace is not None:
pb.partition_id.namespace_id = namespace
for elem in path:
added = pb.path.add()
added.kind = elem['kind']
if 'id' in elem:
added.id = elem['id']
if 'name' in elem:
added.name = elem['name']
return pb
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
def to_protobuf(self):
"""Return a protobuf corresponding to the key.
:rtype: :class:`.entity_pb2.Key`
:returns: The protobuf representing the key.
"""
key = _entity_pb2.Key()
key.partition_id.project_id = self.project
if self.namespace:
key.partition_id.namespace_id = self.namespace
for item in self.path:
element = key.path.add()
if 'kind' in item:
element.kind = item['kind']
if 'id' in item:
element.id = item['id']
if 'name' in item:
element.name = item['name']
return key
def Mutation():
path = [entity_pb2.Key.PathElement(kind="SomeKind")]
return datastore_pb2.Mutation(
upsert=entity_pb2.Entity(key=entity_pb2.Key(path=path))
)