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 _call_fut(self, entity):
from google.cloud.datastore.helpers import entity_to_protobuf
return entity_to_protobuf(entity)
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()
def ConvertFromEntity(entity):
return helpers.entity_to_protobuf(entity)
