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)) )