def _Dynamic_Put(self, put_request, put_response):
if put_request.has_transaction():
entities = put_request.entity_list()
requires_id = lambda x: x.id() == 0 and not x.has_name()
new_ents = [e for e in entities
if requires_id(e.key().path().element_list()[-1])]
id_request = remote_api_pb.PutRequest()
if new_ents:
for ent in new_ents:
e = id_request.add_entity()
e.mutable_key().CopyFrom(ent.key())
e.mutable_entity_group()
id_response = datastore_pb.PutResponse()
super(RemoteDatastoreStub, self).MakeSyncCall(
'remote_datastore', 'GetIDs', id_request, id_response)
assert id_request.entity_size() == id_response.key_size()
for key, ent in zip(id_response.key_list(), new_ents):
ent.mutable_key().CopyFrom(key)
ent.mutable_entity_group().add_element().CopyFrom(
key.path().element(0))
txid = put_request.transaction().handle()
txdata = self.__transactions[txid]
assert (txdata.thread_id ==
thread.get_ident()), "Transactions are single-threaded."
for entity in entities:
txdata.entities[entity.key().Encode()] = (entity.key(), entity)
put_response.add_key().CopyFrom(entity.key())
else:
super(RemoteDatastoreStub, self).MakeSyncCall(
'datastore_v3', 'Put', put_request, put_response)
def test_dynamic_put(self):
PREFIX = "x\x01"
db_batch = flexmock()
db_batch.should_receive('valid_data_version').and_return(True)
zookeeper = flexmock()
zookeeper.should_receive("acquire_lock").and_return(True)
zookeeper.should_receive("release_lock").and_return(True)
zookeeper.should_receive("get_transaction_id").and_return(1)
entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name",
"prop1val", ns="blah")
entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
entity_proto2 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name",
"prop2val", ns="blah")
entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
db_batch.should_receive('batch_get_entity').and_return(
{entity_key1: {}, entity_key2: {}})
db_batch.should_receive('batch_mutate')
dd = DatastoreDistributed(db_batch, zookeeper)
putreq_pb = datastore_pb.PutRequest()
putreq_pb.add_entity()
putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
putreq_pb.add_entity()
putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
putresp_pb = datastore_pb.PutResponse()
dd.dynamic_put('test', putreq_pb, putresp_pb)
self.assertEquals(len(putresp_pb.key_list()), 2)
def test_dynamic_put(self):
PREFIX = "x\x01"
db_batch = flexmock(session=flexmock())
db_batch.should_receive('valid_data_version').and_return(True)
entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name",
"prop1val", ns="blah")
entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
entity_proto2 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name",
"prop2val", ns="blah")
entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
db_batch.should_receive('batch_get_entity').and_return(
{entity_key1: {}, entity_key2: {}})
db_batch.should_receive('batch_mutate')
dd = DatastoreDistributed(db_batch, self.get_zookeeper())
putreq_pb = datastore_pb.PutRequest()
putreq_pb.add_entity()
putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
putreq_pb.add_entity()
putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
putresp_pb = datastore_pb.PutResponse()
entity_lock = flexmock(EntityLock)
entity_lock.should_receive('acquire')
entity_lock.should_receive('release')
flexmock(ScatteredAllocator).should_receive('next').\
and_return(random.randint(1, 500))
dd.dynamic_put('test', putreq_pb, putresp_pb)
self.assertEquals(len(putresp_pb.key_list()), 2)
def _Dynamic_Commit(self, transaction, transaction_response):
txid = transaction.handle()
if txid not in self.__transactions:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST,
'Transaction %d not found.' % (txid, ))
txdata = self.__transactions[txid]
assert (txdata.thread_id == thread.get_ident()
), "Transactions are single-threaded."
del self.__transactions[txid]
tx = remote_api_pb.TransactionRequest()
for key, hash in txdata.preconditions.values():
precond = tx.add_precondition()
precond.mutable_key().CopyFrom(key)
if hash:
precond.set_hash(hash)
puts = tx.mutable_puts()
deletes = tx.mutable_deletes()
for key, entity in txdata.entities.values():
if entity:
puts.add_entity().CopyFrom(entity)
else:
deletes.add_key().CopyFrom(key)
super(RemoteDatastoreStub,
self).MakeSyncCall('remote_datastore', 'Transaction', tx,
datastore_pb.PutResponse())
def test_dynamic_put(self):
PREFIX = "x!"
db_batch = flexmock()
db_batch.should_receive("batch_put_entity").and_return(None)
db_batch.should_receive("batch_get_entity").and_return({PREFIX: {}})
db_batch.should_receive("batch_delete").and_return(None)
zookeeper = flexmock()
zookeeper.should_receive("acquire_lock").and_return(True)
zookeeper.should_receive("release_lock").and_return(True)
zookeeper.should_receive("get_transaction_id").and_return(1)
entity_proto1 = self.get_new_entity_proto("test",
"test_kind",
"bob",
"prop1name",
"prop1val",
ns="blah")
entity_proto2 = self.get_new_entity_proto("test",
"test_kind",
"nancy",
"prop1name",
"prop2val",
ns="blah")
dd = DatastoreDistributed(db_batch, zookeeper)
putreq_pb = datastore_pb.PutRequest()
putreq_pb.add_entity()
putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
putreq_pb.add_entity()
putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
putresp_pb = datastore_pb.PutResponse()
dd.dynamic_put('test', putreq_pb, putresp_pb)
self.assertEquals(len(putresp_pb.key_list()), 2)
def __put_v1_entity(self, v1_entity, v1_txn):
"""Writes a v1 entity to the datastore in a transaction and return its key.
Args:
v1_entity: the entity to write
v1_txn: the transaction in which to write the entity.
Returns:
the key of the entity, which may have been allocated.
"""
v3_entity = entity_pb.EntityProto()
self.__entity_converter.v1_to_v3_entity(v1_entity, v3_entity)
v3_txn = datastore_pb.Transaction()
self.__service_converter.v1_to_v3_txn(v1_txn, v3_txn)
v3_put_req = datastore_pb.PutRequest()
v3_put_req.mutable_transaction().CopyFrom(v3_txn)
v3_put_req.entity_list().append(v3_entity)
v3_put_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_put_req, v3_put_resp)
v3_key = v3_put_resp.key(0)
v1_key = googledatastore.Key()
self.__entity_converter.v3_to_v1_key(v3_key, v1_key)
return v1_key
def put_request(self, app_id, http_request_data):
""" High level function for doing puts.
Args:
app_id: Name of the application.
http_request_data: Stores the protocol buffer request from the AppServer.
Returns:
Returns an encoded put response.
"""
global datastore_access
putreq_pb = datastore_pb.PutRequest(http_request_data)
putresp_pb = datastore_pb.PutResponse()
if READ_ONLY:
logger.warning(
'Unable to put in read-only mode: {}'.format(putreq_pb))
raise gen.Return(
('', datastore_pb.Error.CAPABILITY_DISABLED,
'Datastore is in read-only mode.'))
try:
yield datastore_access.dynamic_put(app_id, putreq_pb, putresp_pb)
raise gen.Return((putresp_pb.Encode(), 0, ''))
except (dbconstants.InternalError, zktransaction.ZKInternalException,
dbconstants.AppScaleDBConnectionError) as error:
raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
except dbconstants.Timeout as error:
raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(error)))
except (dbconstants.BadRequest, zktransaction.ZKBadRequest) as error:
raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
except zktransaction.ZKTransactionException as error:
logger.exception('Concurrent transaction during {}'.format(putreq_pb))
raise gen.Return(
('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
def testLocalPutRequest(self):
"""A request for Put should go local or remote depending on the key."""
put_request = self.CreatePutRequest('localapp')
put_response = datastore_pb.PutResponse()
self.mock_stub.MakeSyncCall('datastore_v3', 'Put', put_request,
put_response)
self.mox.ReplayAll()
self.stub.MakeSyncCall('datastore_v3', 'Put', put_request, put_response)
def testRemotePutTransactionRequest(self):
"""A remote transactional PUT should fail."""
put_request = self.CreatePutRequest('remoteapp')
put_request.mutable_transaction().set_app('remoteapp')
put_request.mutable_transaction().set_handle(123)
put_response = datastore_pb.PutResponse()
self.assertRaises(remote_api_put_stub.RemoteTransactionsUnimplemented,
self.stub.MakeSyncCall, 'datastore_v3', 'Put',
put_request, put_response)
def testRemotePutRequest(self):
"""A request for Put should go local or remote depending on the key."""
put_request = self.CreatePutRequest('remoteapp')
put_response = datastore_pb.PutResponse()
key = db.Key.from_path('MyModel', 1, _app='remoteapp')
expected_put_response = datastore_pb.PutResponse()
expected_put_response.add_key().CopyFrom(key._ToPb())
expected_post = self.RemoteApiRequest('Put', put_request).Encode()
expected_response = self.RemoteApiResponse(
expected_put_response).Encode()
remote_api_put_stub.urlfetch.fetch(
self.remote_url, expected_post, urlfetch.POST,
{'X-appcfg-api-version': '1', 'auth': 'good'}, follow_redirects=False
).AndReturn(MockUrlfetchResult(200, expected_response))
self.mox.ReplayAll()
self.stub.MakeSyncCall('datastore_v3', 'Put', put_request, put_response)
self.assertEqual(put_response, expected_put_response)
def _put_entity(ds_access, entity):
""" Updates or creates an entity.
Args:
ds_access: A DatastoreDistributed client.
entity: A datastore.Entity object.
Raises:
ApplicationError if unable to put entity.
"""
request = datastore_pb.PutRequest()
new_entity = request.add_entity()
new_entity.CopyFrom(entity.ToPb())
response = datastore_pb.PutResponse()
ds_access._Dynamic_Put(request, response)
def testRemotePutRequestUnauthorized(self):
"""A remote put with a 'bad' urlfetch response."""
put_request = self.CreatePutRequest('remoteapp')
put_response = datastore_pb.PutResponse()
expected_post = self.RemoteApiRequest('Put', put_request).Encode()
remote_api_put_stub.urlfetch.fetch(
self.remote_url, expected_post, urlfetch.POST,
{'X-appcfg-api-version': '1', 'auth': 'good'}, follow_redirects=False
).AndReturn(MockUrlfetchResult(403, 'not authorized'))
self.mox.ReplayAll()
self.assertRaises(remote_api_put_stub.FetchFailed, self.stub.MakeSyncCall,
'datastore_v3', 'Put', put_request, put_response)
def __upsert_v3_entity(self, v3_entity, v3_txn):
"""Upsert a v3 entity.
Args:
v3_entity: an entity_pb.EntityProto
v3_txn: a datastore_pb.Transaction or None
Returns:
a tuple (the number of index writes that occurred,
the key of the entity)
"""
v3_put_req = datastore_pb.PutRequest()
if v3_txn:
v3_put_req.mutable_transaction().CopyFrom(v3_txn)
v3_put_req.entity_list().append(v3_entity)
v3_put_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_put_req, v3_put_resp)
return (v3_put_resp.cost().index_writes(), v3_put_resp.key(0))
def put_request(self, app_id, http_request_data):
""" High level function for doing puts.
Args:
app_id: Name of the application.
http_request_data: Stores the protocol buffer request from the AppServer.
Returns:
Returns an encoded put response.
"""
global datastore_access
putreq_pb = datastore_pb.PutRequest(http_request_data)
putresp_pb = datastore_pb.PutResponse()
if READ_ONLY:
logger.warning('Unable to put in read-only mode: {}'.
format(putreq_pb))
return (putresp_pb.Encode(), datastore_pb.Error.CAPABILITY_DISABLED,
'Datastore is in read-only mode.')
try:
datastore_access.dynamic_put(app_id, putreq_pb, putresp_pb)
return (putresp_pb.Encode(), 0, "")
except zktransaction.ZKBadRequest as zkie:
logger.exception('Illegal argument during {}'.format(putreq_pb))
return (putresp_pb.Encode(),
datastore_pb.Error.BAD_REQUEST,
"Illegal arguments for transaction. {0}".format(str(zkie)))
except zktransaction.ZKInternalException:
logger.exception('ZKInternalException during {}'.format(putreq_pb))
return (putresp_pb.Encode(),
datastore_pb.Error.INTERNAL_ERROR,
"Internal error with ZooKeeper connection.")
except zktransaction.ZKTransactionException:
logger.exception('Concurrent transaction during {}'.
format(putreq_pb))
return (putresp_pb.Encode(),
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Concurrent transaction exception on put.")
except dbconstants.AppScaleDBConnectionError:
logger.exception('DB connection error during {}'.format(putreq_pb))
return (putresp_pb.Encode(),
datastore_pb.Error.INTERNAL_ERROR,
"Datastore connection error on put.")
def test_dynamic_put(self):
db_batch = flexmock(session=flexmock())
db_batch.should_receive('valid_data_version_sync').and_return(True)
entity_proto1 = self.get_new_entity_proto(
"test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
entity_proto2 = self.get_new_entity_proto(
"test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
async_result = gen.Future()
async_result.set_result({entity_key1: {}, entity_key2: {}})
db_batch.should_receive('batch_get_entity').and_return(async_result)
db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
transaction_manager = flexmock(
create_transaction_id=lambda project, xg: 1,
delete_transaction_id=lambda project, txid: None,
set_groups=lambda project, txid, groups: None)
dd = DatastoreDistributed(db_batch, transaction_manager,
self.get_zookeeper())
dd.index_manager = flexmock(
projects={'test': flexmock(indexes_pb=[])})
putreq_pb = datastore_pb.PutRequest()
putreq_pb.add_entity()
putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
putreq_pb.add_entity()
putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
putresp_pb = datastore_pb.PutResponse()
async_true = gen.Future()
async_true.set_result(True)
entity_lock = flexmock(EntityLock)
entity_lock.should_receive('acquire').and_return(async_true)
entity_lock.should_receive('release')
flexmock(ScatteredAllocator).should_receive('next').\
and_return(random.randint(1, 500))
yield dd.dynamic_put('test', putreq_pb, putresp_pb)
self.assertEquals(len(putresp_pb.key_list()), 2)
def testRemotePutRemoteException(self):
"""Test that a remote exception is bubbled back up."""
put_request = self.CreatePutRequest('remoteapp')
put_response = datastore_pb.PutResponse()
expected_post = self.RemoteApiRequest('Put', put_request).Encode()
expected_exception = db.Timeout('too slow')
remote_response = remote_api_pb.Response()
remote_response.mutable_exception().set_contents(
pickle.dumps(expected_exception))
expected_response = remote_response.Encode()
remote_api_put_stub.urlfetch.fetch(
self.remote_url, expected_post, urlfetch.POST,
{'X-appcfg-api-version': '1', 'auth': 'good'}, follow_redirects=False
).AndReturn(MockUrlfetchResult(200, expected_response))
self.mox.ReplayAll()
self.assertRaises(db.Timeout, self.stub.MakeSyncCall,
'datastore_v3', 'Put', put_request, put_response)
def __insert_v3_entity(self, v3_entity, v3_txn):
"""Inserts a v3 entity.
Args:
v3_entity: an entity_pb.EntityProto
v3_txn: a datastore_pb.Transaction or None
Returns:
a tuple (the number of index writes that occurred,
the entity key)
Raises:
ApplicationError: if the entity already exists
"""
if not v3_txn:
v3_txn = datastore_pb.Transaction()
v3_begin_txn_req = datastore_pb.BeginTransactionRequest()
v3_begin_txn_req.set_app(v3_entity.key().app())
self.__make_v3_call('BeginTransaction', v3_begin_txn_req, v3_txn)
_, key = self.__insert_v3_entity(v3_entity, v3_txn)
v3_resp = datastore_pb.CommitResponse()
self.__make_v3_call('Commit', v3_txn, v3_resp)
return (v3_resp.cost().index_writes(), key)
if datastore_pbs.is_complete_v3_key(v3_entity.key()):
v3_get_req = datastore_pb.GetRequest()
v3_get_req.mutable_transaction().CopyFrom(v3_txn)
v3_get_req.key_list().append(v3_entity.key())
v3_get_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_get_req, v3_get_resp)
if v3_get_resp.entity(0).has_entity():
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, 'Entity already exists.')
v3_put_req = datastore_pb.PutRequest()
v3_put_req.mutable_transaction().CopyFrom(v3_txn)
v3_put_req.entity_list().append(v3_entity)
v3_put_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_put_req, v3_put_resp)
return (v3_put_resp.cost().index_writes(), v3_put_resp.key(0))
def store_entity_batch(self, entity_batch):
""" Stores the given entity batch.
Args:
entity_batch: A list of entities to store.
Returns:
True on success, False otherwise.
"""
logging.debug("Entity batch to process: {0}".format(entity_batch))
# Convert encoded entities to EntityProto objects, change the app ID if
# it's different than the original and encode again.
new_entities_encoded = []
ent_protos = []
for entity in entity_batch:
ent_proto = entity_pb.EntityProto()
ent_proto.ParseFromString(entity)
ent_proto.key().set_app(self.app_id)
ent_protos.append(ent_proto)
new_entities_encoded.append(ent_proto.Encode())
logging.debug("Entities encoded: {0}".format(new_entities_encoded))
# Create a PutRequest with the entities to be stored.
put_request = datastore_pb.PutRequest()
put_response = datastore_pb.PutResponse()
for entity in new_entities_encoded:
new_entity = put_request.add_entity()
new_entity.MergeFromString(entity)
logging.debug("Put request: {0}".format(put_request))
try:
self.ds_distributed.dynamic_put(self.app_id, put_request,
put_response)
self.entities_restored += len(ent_protos)
except zk.ZKInternalException, zkie:
logging.error("ZK internal exception for app id {0}, " \
"info {1}".format(self.app_id, str(zkie)))
return False
def __update_v3_entity(self, v3_entity, v3_txn):
"""Updates a v3 entity.
Args:
v3_entity: a datastore_v4_pb.Entity
v3_txn: a datastore_pb.Transaction or None
Returns:
the number of index writes that occurred
Raises:
ApplicationError: if the entity does not exist
"""
if not v3_txn:
v3_txn = datastore_pb.Transaction()
v3_begin_txn_req = datastore_pb.BeginTransactionRequest()
v3_begin_txn_req.set_app(v3_entity.key().app())
self.__make_v3_call('BeginTransaction', v3_begin_txn_req, v3_txn)
self.__update_v3_entity(v3_entity, v3_txn)
v3_resp = datastore_pb.CommitResponse()
self.__make_v3_call('Commit', v3_txn, v3_resp)
return v3_resp.cost().index_writes()
v3_get_req = datastore_pb.GetRequest()
v3_get_req.mutable_transaction().CopyFrom(v3_txn)
v3_get_req.key_list().append(v3_entity.key())
v3_get_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_get_req, v3_get_resp)
if not v3_get_resp.entity(0).has_entity():
raise apiproxy_errors.ApplicationError(
datastore_v4_pb.Error.BAD_REQUEST, 'Entity does not exist.')
v3_put_req = datastore_pb.PutRequest()
v3_put_req.mutable_transaction().CopyFrom(v3_txn)
v3_put_req.entity_list().append(v3_entity)
v3_put_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_put_req, v3_put_resp)
return v3_put_resp.cost().index_writes()
def __apply_v4_deprecated_mutation(self, v4_deprecated_mutation, v4_txn):
"""Applies a v4 DeprecatedMutation.
Args:
v4_deprecated_mutation: a datastore_v4_pb.DeprecatedMutation
v4_txn: an optional v4 transaction handle or None
Returns:
a datastore_v4_pb.DeprecatedMutationResult
"""
index_writes = 0
v3_txn = None
if v4_txn:
v3_txn = datastore_pb.Transaction()
self.__service_converter.v4_to_v3_txn(v4_txn, v3_txn)
for v4_entity in v4_deprecated_mutation.insert_list():
v3_entity = entity_pb.EntityProto()
self.__entity_converter.v4_to_v3_entity(v4_entity, v3_entity)
index_writes += self.__insert_v3_entity(v3_entity, v3_txn)
for v4_entity in v4_deprecated_mutation.update_list():
v3_entity = entity_pb.EntityProto()
self.__entity_converter.v4_to_v3_entity(v4_entity, v3_entity)
index_writes += self.__update_v3_entity(v3_entity, v3_txn)
v3_insert_auto_req = datastore_pb.PutRequest()
if v3_txn:
v3_insert_auto_req.mutable_transaction().CopyFrom(v3_txn)
for v4_entity in v4_deprecated_mutation.insert_auto_id_list():
v3_entity = entity_pb.EntityProto()
self.__entity_converter.v4_to_v3_entity(v4_entity, v3_entity)
v3_insert_auto_req.entity_list().append(v3_entity)
v3_insert_auto_id_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_insert_auto_req, v3_insert_auto_id_resp)
index_writes += v3_insert_auto_id_resp.cost().index_writes()
v3_upsert_req = datastore_pb.PutRequest()
if v3_txn:
v3_upsert_req.mutable_transaction().CopyFrom(v3_txn)
for v4_entity in v4_deprecated_mutation.upsert_list():
v3_entity = entity_pb.EntityProto()
self.__entity_converter.v4_to_v3_entity(v4_entity, v3_entity)
v3_upsert_req.entity_list().append(v3_entity)
v3_upsert_resp = datastore_pb.PutResponse()
self.__make_v3_call('Put', v3_upsert_req, v3_upsert_resp)
index_writes += v3_upsert_resp.cost().index_writes()
v3_delete_req = datastore_pb.DeleteRequest()
if v3_txn:
v3_delete_req.mutable_transaction().CopyFrom(v3_txn)
for v4_key in v4_deprecated_mutation.delete_list():
self.__entity_converter.v4_to_v3_reference(v4_key,
v3_delete_req.add_key())
v3_delete_resp = datastore_pb.DeleteResponse()
self.__make_v3_call('Delete', v3_delete_req, v3_delete_resp)
index_writes += v3_delete_resp.cost().index_writes()
v4_deprecated_mutation_result = datastore_v4_pb.DeprecatedMutationResult(
)
for v3_ref in v3_insert_auto_id_resp.key_list():
self.__entity_converter.v3_to_v4_key(
v3_ref, v4_deprecated_mutation_result.add_insert_auto_id_key())
v4_deprecated_mutation_result.set_index_updates(index_writes)
return v4_deprecated_mutation_result
