def test_w_transaction(datastore_pb2, in_context): class Mutation: def __init__(self, upsert=None): self.upsert = upsert def __eq__(self, other): return self.upsert == other.upsert def MockEntity(*path): key = ds_key_module.Key(*path, project="testing") return entity.Entity(key=key) with in_context.new(transaction=b"123").use() as context: datastore_pb2.Mutation = Mutation entity1 = MockEntity("a", "1") _api.put(entity1, _options.Options()) entity2 = MockEntity("a") _api.put(entity2, _options.Options()) entity3 = MockEntity("b") _api.put(entity3, _options.Options()) batch = context.commit_batches[b"123"] assert batch.mutations == [ Mutation(upsert=helpers.entity_to_protobuf(entity1)), Mutation(upsert=helpers.entity_to_protobuf(entity2)), Mutation(upsert=helpers.entity_to_protobuf(entity3)), ] assert batch.transaction == b"123" assert batch.incomplete_mutations == [ Mutation(upsert=helpers.entity_to_protobuf(entity2)), Mutation(upsert=helpers.entity_to_protobuf(entity3)), ]
def test_no_transaction(datastore_pb2, in_context): class Mutation: def __init__(self, upsert=None): self.upsert = upsert def __eq__(self, other): return self.upsert == other.upsert def MockEntity(*path): key = ds_key_module.Key(*path, project="testing") return entity.Entity(key=key) datastore_pb2.Mutation = Mutation entity1 = MockEntity("a", "1") _api.put(entity1, _options.Options()) entity2 = MockEntity("a") _api.put(entity2, _options.Options()) entity3 = MockEntity("b") _api.put(entity3, _options.Options()) batch = in_context.batches[_api._NonTransactionalCommitBatch][()] assert batch.mutations == [ Mutation(upsert=helpers.entity_to_protobuf(entity1)), Mutation(upsert=helpers.entity_to_protobuf(entity2)), Mutation(upsert=helpers.entity_to_protobuf(entity3)), ]
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 test_entity_from_protobuf(): class Animal(polymodel.PolyModel): pass class Cat(Animal): pass key = datastore.Key("Cat", 123, project="testing") datastore_entity = datastore.Entity(key=key) protobuf = helpers.entity_to_protobuf(datastore_entity) entity = model._entity_from_protobuf(protobuf) assert isinstance(entity, Cat)
def _assign_entity_to_pb(entity_pb, entity): """Copy ``entity`` into ``entity_pb``. Helper method for ``Batch.put``. :type entity_pb: :class:`.entity_pb2.Entity` :param entity_pb: The entity owned by a mutation. :type entity: :class:`google.cloud.datastore.entity.Entity` :param entity: The entity being updated within the batch / transaction. """ bare_entity_pb = helpers.entity_to_protobuf(entity) bare_entity_pb._pb.key.CopyFrom(bare_entity_pb._pb.key) entity_pb._pb.CopyFrom(bare_entity_pb._pb)
def _assign_entity_to_pb(entity_pb, entity): """Copy ``entity`` into ``entity_pb``. Helper method for ``Batch.put``. :type entity_pb: :class:`.entity_pb2.Entity` :param entity_pb: The entity owned by a mutation. :type entity: :class:`google.cloud.datastore.entity.Entity` :param entity: The entity being updated within the batch / transaction. """ bare_entity_pb = helpers.entity_to_protobuf(entity) bare_entity_pb.key.CopyFrom(bare_entity_pb.key) entity_pb.CopyFrom(bare_entity_pb)
def _call_fut(self, entity): from google.cloud.datastore.helpers import entity_to_protobuf return entity_to_protobuf(entity)
def ByteSize(self): if self.entity is not None: return helpers.entity_to_protobuf(self.entity).ByteSize() else: return self.key.to_protobuf().ByteSize()
def ConvertFromEntity(entity): return helpers.entity_to_protobuf(entity)
def _serialize_entity(entity): """Converts an Entity object to a serialized proto string.""" if entity is None: return None return helpers.entity_to_protobuf(entity)._pb.SerializeToString()