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