Using GAE's search API with ndb's tasklets

Question

The Google AppEngine search API can return asynchronous results. The docs say very little about these futures, but they do have a .get_result() method which looks a lot like an ndb.Future. I thought it would be fun to try to use one in a tasklet:

@ndb.tasklet
def async_query(index):
    results = yield [index.search_async('foo'), index.search_async('bar')]
    raise ndb.Return(results)

Unfortunately, this doesn't work. ndb doesn't like this because the future returned by the search API doesn't seem to be compatible with ndb.Future. However, the tasklet documentation also specifically mentions that they have been made to work with urlfetch futures. Is there a way to get similar behavior with the search api?

Answer 1

It turns out that there is a way to make this work (at least, my tests on the dev_appserver seem to be passing ;-).

@ndb.tasklet
def async_query(index, query):
    fut = index.search_async(query)
    yield fut._rpc
    raise ndb.Return(fut._get_result_hook())

Now if I want to do multiple queries and intermix some datastore queries (i.e. maybe I use search API to get IDs for model entities),

@ndb.tasklet
def get_entities(queries):
    search_results = yield [async_query(q) for q in queries]
    ids = set()
    for search_result in search_results:
        for doc in search_results.results:
            ids.add(doc.doc_id)
    entities = yield [FooModel.get_by_id_async(id_) for id_ in ids_]
    raise ndb.Return(entities)

This is super hacky -- and I doubt that it is officially supported since I am using internal members of the async search result class ... Use at your own risk :-).