def test_dynamic_get(self):
entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name",
"prop2val", ns="blah")
zookeeper = flexmock()
zookeeper.should_receive("get_valid_transaction_id").and_return(1)
zookeeper.should_receive("register_updated_key").and_return(1)
zookeeper.should_receive("acquire_lock").and_return(True)
db_batch = flexmock()
db_batch.should_receive("batch_put_entity").and_return(None)
db_batch.should_receive("batch_get_entity").and_return(
{"test\x00blah\x00test_kind:nancy\x01":
{
APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
APP_ENTITY_SCHEMA[1]: 1
}
})
dd = DatastoreDistributed(db_batch, zookeeper)
entity_key = entity_proto1.key()
get_req = datastore_pb.GetRequest()
key = get_req.add_key()
key.MergeFrom(entity_key)
get_resp = datastore_pb.GetResponse()
dd.dynamic_get("test", get_req, get_resp)
self.assertEquals(get_resp.entity_size(), 1)
# Now test while in a transaction
get_resp = datastore_pb.GetResponse()
get_req.mutable_transaction().set_handle(1)
dd.dynamic_get("test", get_req, get_resp)
self.assertEquals(get_resp.entity_size(), 1)
def get_request(self, app_id, http_request_data):
""" High level function for doing gets.
Args:
app_id: Name of the application.
http_request_data: Stores the protocol buffer request from the AppServer.
Returns:
An encoded get response.
"""
global datastore_access
getreq_pb = datastore_pb.GetRequest(http_request_data)
getresp_pb = datastore_pb.GetResponse()
try:
datastore_access.dynamic_get(app_id, getreq_pb, getresp_pb)
except zktransaction.ZKBadRequest as zkie:
logger.exception('Illegal argument during {}'.format(getreq_pb))
return (getresp_pb.Encode(), datastore_pb.Error.BAD_REQUEST,
"Illegal arguments for transaction. {0}".format(str(zkie)))
except zktransaction.ZKInternalException:
logger.exception('ZKInternalException during {}'.format(getreq_pb))
return (getresp_pb.Encode(), datastore_pb.Error.INTERNAL_ERROR,
"Internal error with ZooKeeper connection.")
except zktransaction.ZKTransactionException:
logger.exception(
'Concurrent transaction during {}'.format(getreq_pb))
return (getresp_pb.Encode(),
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Concurrent transaction exception on get.")
except dbconstants.AppScaleDBConnectionError:
logger.exception('DB connection error during {}'.format(getreq_pb))
return (getresp_pb.Encode(), datastore_pb.Error.INTERNAL_ERROR,
"Datastore connection error on get.")
return (getresp_pb.Encode(), 0, "")
def _Dynamic_Get(self, request, response):
"""Intercepts get requests and returns them from cache if available."""
logging.info("Tx: %s, Keys: %s", request.has_transaction(),
[str(x) for x in request.key_list()])
if request.has_transaction():
self.CallWrappedStub('Get', request, response)
return
new_request = datastore_pb.GetRequest()
new_response = datastore_pb.GetResponse()
encoded_keys = [k.Encode() for k in request.key_list()]
cached = memcache.get_multi(encoded_keys)
for key, encoded_key in itertools.izip(request.key_list(),
encoded_keys):
if encoded_key not in cached:
new_request.add_key().CopyFrom(key)
if new_request.key_size() > 0:
self.CallWrappedStub('Get', new_request, new_response)
entity_iter = iter(new_response.entity_list())
to_put = dict()
for encoded_key in encoded_keys:
entity = cached.get(encoded_key, None)
if entity:
response.add_entity().mutable_entity().CopyFrom(entity)
else:
entity = entity_iter.next()
if entity.entity().IsInitialized():
# self.entity_cache[encoded_key] = entity.entity()
to_put[encoded_key] = entity.entity()
response.add_entity().CopyFrom(entity)
if to_put:
memcache.set_multi(to_put)
def _Dynamic_Lookup(self, req, resp):
new_txn = None
try:
try:
self.__service_validator.validate_lookup_req(req)
if req.read_options.WhichOneof(
'consistency_type') == 'new_transaction':
new_txn = self.__begin_adhoc_txn(req)
v3_req = self.__service_converter.v1_to_v3_get_req(
req, new_txn=new_txn)
except (cloud_datastore_validator.ValidationError,
datastore_pbs.InvalidConversionError), e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, str(e))
v3_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_req, v3_resp)
try:
v1_resp = self.__service_converter.v3_to_v1_lookup_resp(
v3_resp, new_txn=new_txn)
except datastore_pbs.InvalidConversionError, e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.INTERNAL_ERROR, str(e))
def get_request(self, app_id, http_request_data):
""" High level function for doing gets.
Args:
app_id: Name of the application.
http_request_data: Stores the protocol buffer request from the AppServer.
Returns:
An encoded get response.
"""
global datastore_access
getreq_pb = datastore_pb.GetRequest(http_request_data)
getresp_pb = datastore_pb.GetResponse()
try:
yield datastore_access.dynamic_get(app_id, getreq_pb, getresp_pb)
except zktransaction.ZKBadRequest as error:
logger.exception('Illegal argument during {}'.format(getreq_pb))
raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
except zktransaction.ZKInternalException as error:
logger.exception('ZKInternalException during {}'.format(getreq_pb))
raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
except zktransaction.ZKTransactionException as error:
logger.exception('Concurrent transaction during {}'.format(getreq_pb))
raise gen.Return(
('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
except dbconstants.AppScaleDBConnectionError:
logger.exception('DB connection error during get')
raise gen.Return(
('', datastore_pb.Error.INTERNAL_ERROR,
'Datastore connection error on get.'))
raise gen.Return((getresp_pb.Encode(), 0, ''))
def _Dynamic_Lookup(self, req, resp):
try:
self.__service_validator.validate_lookup_req(req)
v3_req = self.__service_converter.v4_to_v3_get_req(req)
v3_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_req, v3_resp)
except datastore_pbs.InvalidConversionError, e:
raise apiproxy_errors.ApplicationError(
datastore_v4_pb.Error.BAD_REQUEST, str(e))
def _Dynamic_Get(self, get_request, get_response):
txid = None
if get_request.has_transaction():
txid = get_request.transaction().handle()
txdata = self.__transactions[txid]
assert (txdata.thread_id ==
thread.get_ident()), "Transactions are single-threaded."
keys = [(k, k.Encode()) for k in get_request.key_list()]
new_request = datastore_pb.GetRequest()
for key, enckey in keys:
if enckey not in txdata.entities:
new_request.add_key().CopyFrom(key)
else:
new_request = get_request
if new_request.key_size() > 0:
super(RemoteDatastoreStub, self).MakeSyncCall(
'datastore_v3', 'Get', new_request, get_response)
if txid is not None:
newkeys = new_request.key_list()
entities = get_response.entity_list()
for key, entity in zip(newkeys, entities):
entity_hash = None
if entity.has_entity():
entity_hash = sha.new(entity.entity().Encode()).digest()
txdata.preconditions[key.Encode()] = (key, entity_hash)
new_response = datastore_pb.GetResponse()
it = iter(get_response.entity_list())
for key, enckey in keys:
if enckey in txdata.entities:
cached_entity = txdata.entities[enckey][1]
if cached_entity:
new_response.add_entity().mutable_entity().CopyFrom(cached_entity)
else:
new_response.add_entity()
else:
new_entity = it.next()
if new_entity.has_entity():
assert new_entity.entity().key() == key
new_response.add_entity().CopyFrom(new_entity)
else:
new_response.add_entity()
get_response.CopyFrom(new_response)
def _Dynamic_Transaction(self, request, response):
"""Handle a Transaction request.
We handle transactions by accumulating Put and Delete requests on the client
end, as well as recording the key and hash of Get requests. When Commit is
called, Transaction is invoked, which verifies that all the entities in the
precondition list still exist and their hashes match, then performs a
transaction of its own to make the updates.
"""
begin_request = datastore_pb.BeginTransactionRequest()
begin_request.set_app(os.environ['APPLICATION_ID'])
begin_request.set_allow_multiple_eg(request.allow_multiple_eg())
tx = datastore_pb.Transaction()
self.__call('datastore_v3', 'BeginTransaction', begin_request, tx)
preconditions = request.precondition_list()
if preconditions:
get_request = datastore_pb.GetRequest()
get_request.mutable_transaction().CopyFrom(tx)
for precondition in preconditions:
key = get_request.add_key()
key.CopyFrom(precondition.key())
get_response = datastore_pb.GetResponse()
self.__call('datastore_v3', 'Get', get_request, get_response)
entities = get_response.entity_list()
assert len(entities) == request.precondition_size()
for precondition, entity in zip(preconditions, entities):
if precondition.has_hash() != entity.has_entity():
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Transaction precondition failed.")
elif entity.has_entity():
entity_hash = hashlib.sha1(entity.entity().Encode()).digest()
if precondition.hash() != entity_hash:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Transaction precondition failed.")
if request.has_puts():
put_request = request.puts()
put_request.mutable_transaction().CopyFrom(tx)
self.__call('datastore_v3', 'Put', put_request, response)
if request.has_deletes():
delete_request = request.deletes()
delete_request.mutable_transaction().CopyFrom(tx)
self.__call('datastore_v3', 'Delete', delete_request,
datastore_pb.DeleteResponse())
self.__call('datastore_v3', 'Commit', tx, datastore_pb.CommitResponse())
def testRemoteGetRequest(self):
"""A request for Get should always pass through."""
remote_key = db.Key.from_path('MyModel', 'Name1', _app='remoteapp')
get_request = datastore_pb.GetRequest()
get_request.add_key().CopyFrom(remote_key._ToPb())
get_response = datastore_pb.GetResponse()
self.mock_stub.MakeSyncCall('datastore_v3', 'Get', get_request,
get_response)
self.mox.ReplayAll()
self.stub.MakeSyncCall('datastore_v3', 'Get', get_request, get_response)
def test_dynamic_get(self):
entity_proto1 = self.get_new_entity_proto(
"test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
zk_client = flexmock()
zk_client.should_receive('add_listener')
async_result = gen.Future()
async_result.set_result({
"test\x00blah\x00test_kind:nancy\x01": {
APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
APP_ENTITY_SCHEMA[1]: 1
}
})
zookeeper = flexmock(handle=zk_client)
zookeeper.should_receive("get_valid_transaction_id").and_return(1)
zookeeper.should_receive("acquire_lock").and_return(True)
db_batch = flexmock()
db_batch.should_receive('valid_data_version_sync').and_return(True)
db_batch.should_receive("batch_get_entity").and_return(async_result)
db_batch.should_receive('record_reads').and_return(ASYNC_NONE)
transaction_manager = flexmock()
dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
entity_key = entity_proto1.key()
get_req = datastore_pb.GetRequest()
key = get_req.add_key()
key.MergeFrom(entity_key)
get_resp = datastore_pb.GetResponse()
yield dd.dynamic_get("test", get_req, get_resp)
self.assertEquals(get_resp.entity_size(), 1)
# Now test while in a transaction
get_resp = datastore_pb.GetResponse()
get_req.mutable_transaction().set_handle(1)
yield dd.dynamic_get("test", get_req, get_resp)
self.assertEquals(get_resp.entity_size(), 1)
def _Dynamic_Transaction(self, request, response):
"""Handle a Transaction request.
We handle transactions by accumulating Put requests on the client end, as
well as recording the key and hash of Get requests. When Commit is called,
Transaction is invoked, which verifies that all the entities in the
precondition list still exist and their hashes match, then performs a
transaction of its own to make the updates.
"""
tx = datastore_pb.Transaction()
apiproxy_stub_map.MakeSyncCall('datastore_v3', 'BeginTransaction',
api_base_pb.VoidProto(), tx)
preconditions = request.precondition_list()
if preconditions:
get_request = datastore_pb.GetRequest()
get_request.mutable_transaction().CopyFrom(tx)
for precondition in preconditions:
key = get_request.add_key()
key.CopyFrom(precondition.key())
get_response = datastore_pb.GetResponse()
apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Get', get_request,
get_response)
entities = get_response.entity_list()
assert len(entities) == request.precondition_size()
for precondition, entity in zip(preconditions, entities):
if precondition.has_hash() != entity.has_entity():
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Transaction precondition failed.")
elif entity.has_entity():
entity_hash = sha.new(entity.entity().Encode()).digest()
if precondition.hash() != entity_hash:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.CONCURRENT_TRANSACTION,
"Transaction precondition failed.")
if request.has_puts():
put_request = request.puts()
put_request.mutable_transaction().CopyFrom(tx)
apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Put',
put_request, response)
if request.has_deletes():
delete_request = request.deletes()
delete_request.mutable_transaction().CopyFrom(tx)
apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Delete',
delete_request, api_base_pb.VoidProto())
apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Commit', tx,
api_base_pb.VoidProto())
def _Dynamic_TransactionQuery(self, request, response):
if not request.has_ancestor():
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST,
'No ancestor in transactional query.')
app_id = datastore_types.ResolveAppId(None)
if (datastore_rpc._GetDatastoreType(app_id) !=
datastore_rpc.BaseConnection.HIGH_REPLICATION_DATASTORE):
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST,
'remote_api supports transactional queries only in the '
'high-replication datastore.')
entity_group_key = entity_pb.Reference()
entity_group_key.CopyFrom(request.ancestor())
group_path = entity_group_key.mutable_path()
root = entity_pb.Path_Element()
root.MergeFrom(group_path.element(0))
group_path.clear_element()
group_path.add_element().CopyFrom(root)
eg_element = group_path.add_element()
eg_element.set_type(metadata.EntityGroup.KIND_NAME)
eg_element.set_id(metadata.EntityGroup.ID)
begin_request = datastore_pb.BeginTransactionRequest()
begin_request.set_app(app_id)
tx = datastore_pb.Transaction()
self.__call('datastore_v3', 'BeginTransaction', begin_request, tx)
request.mutable_transaction().CopyFrom(tx)
self.__call('datastore_v3', 'RunQuery', request, response.mutable_result())
get_request = datastore_pb.GetRequest()
get_request.mutable_transaction().CopyFrom(tx)
get_request.add_key().CopyFrom(entity_group_key)
get_response = datastore_pb.GetResponse()
self.__call('datastore_v3', 'Get', get_request, get_response)
entity_group = get_response.entity(0)
response.mutable_entity_group_key().CopyFrom(entity_group_key)
if entity_group.has_entity():
response.mutable_entity_group().CopyFrom(entity_group.entity())
self.__call('datastore_v3', 'Commit', tx, datastore_pb.CommitResponse())
def get(self, key, txid=None):
request = datastore_pb.GetRequest()
req_key = request.add_key()
req_key.MergeFrom(key._ToPb())
if txid is not None:
req_tx = request.mutable_transaction()
req_tx.set_app(self.project_id)
req_tx.set_handle(txid)
encoded_response = yield self._make_request('Get', request.Encode())
get_response = datastore_pb.GetResponse(encoded_response)
response_entity = get_response.entity(0).entity()
if not response_entity.has_key():
return
raise gen.Return(Entity.FromPb(response_entity))
def _get_entity_by_key(ds_access, entity_key):
""" Fetches an entity.
Args:
ds_access: A DatastoreDistributed client.
entity_key: A datastore.Key object.
Returns:
A datastore.Entity object.
Raises:
ApplicationError if unable to fetch entity.
"""
reference = entity_key._ToPb()
request = datastore_pb.GetRequest()
key = request.add_key()
key.CopyFrom(reference)
response = datastore_pb.GetResponse()
ds_access._Dynamic_Get(request, response)
entity = datastore.Entity.FromPb(response.entity(0).entity())
return entity
def _Dynamic_Lookup(self, req, resp):
try:
self.__service_validator.validate_lookup_req(req)
v3_req = self.__service_converter.v1_to_v3_get_req(req)
except cloud_datastore_validator.ValidationError as e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, str(e))
except datastore_pbs.InvalidConversionError as e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, str(e))
v3_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_req, v3_resp)
try:
v1_resp = self.__service_converter.v3_to_v1_lookup_resp(v3_resp)
except datastore_pbs.InvalidConversionError as e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.INTERNAL_ERROR, str(e))
resp.CopyFrom(v1_resp)
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 __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 __get_v1_entity_version(self, v1_key, v1_txn):
"""Returns the version of an entity, or _NO_VERSION if it does not exist.
Args:
v1_key: the key of the entity to lookup.
v1_txn: the transaction to use when retrieving the entity.
Returns:
the version number of the entity if it was found, or _NO_VERSION
otherwise.
"""
v3_key = entity_pb.Reference()
self.__entity_converter.v1_to_v3_reference(v1_key, v3_key)
v3_txn = datastore_pb.Transaction()
self.__service_converter.v1_to_v3_txn(v1_txn, v3_txn)
v3_get_req = datastore_pb.GetRequest()
v3_get_req.mutable_transaction().CopyFrom(v3_txn)
v3_get_req.key_list().append(v3_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():
return v3_get_resp.entity(0).version()
return _NO_VERSION
datastore_pb.Error.INTERNAL_ERROR, str(e))
resp.CopyFrom(v1_resp)
def _Dynamic_Lookup(self, req, resp):
try:
self.__service_validator.validate_lookup_req(req)
v3_req = self.__service_converter.v1_to_v3_get_req(req)
except cloud_datastore_validator.ValidationError, e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, str(e))
except datastore_pbs.InvalidConversionError, e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.BAD_REQUEST, str(e))
v3_resp = datastore_pb.GetResponse()
self.__make_v3_call('Get', v3_req, v3_resp)
try:
v1_resp = self.__service_converter.v3_to_v1_lookup_resp(v3_resp)
except datastore_pbs.InvalidConversionError, e:
raise apiproxy_errors.ApplicationError(
datastore_pb.Error.INTERNAL_ERROR, str(e))
resp.CopyFrom(v1_resp)
def _Dynamic_AllocateIds(self, req, resp):
v3_stub = apiproxy_stub_map.apiproxy.GetStub(V3_SERVICE_NAME)
v3_refs = None
try:
self.__service_validator.validate_allocate_ids_req(req)