MultiQuery Support to Ease Migration

[speedplane]

This library should support the MultiQuery class (docs, source) found in the original datastore API.

Background

According to the App Engine migration guide, the original API is being phased out and developers are being asked to port their code to the GCL libraries. Unfortunately, these newer libraries are missing some of the features found in the older libraries.

The older datastore API supported the class google.appengine.api.datastore.MultiQuery. This class could run multiple queries at once and then merge sort the results in memory. There were plenty of limitations with the MultiQuery class, including large result sets and cursors. However, there are several main advantages:

1. Supporting MultiQuery will ease transition from the old API to the new one
2. It is a helpful tool for quickly being able to perform many queries at once.
3. It can assist middleware developers when implementing database IN operators.

Accordingly, I propose a similar class be built for the new datastore libraries.

If MultiQuery is not supported, some migration solution should be outlined in the migration guide.

[speedplane]

I went ahead and built a small class to handle this. The DatastoreMultiQuery class takes as a constructor an array of query objects, and then emulates a unified Query class, and specifically the fetch function. Cursors and ordering on multiple attributes don't work, but with some work this functionality could be added.

class DatastoreMultiQuery(object):
    '''Query that emulates the cloud.datastore.query with multiple queries.'''
    def __init__(self, queries, order = None):
        if isinstance(order, list):
            if len(order) > 1:
                raise NotImplementedError("Can't handle ordering: %s"%order)
        self.order = order
        self.queries = queries
        
        # Make sure the attributes of the multiquery are compatible.
        self.all_attrs = {}
        for a in ["kind", "namespace", "project", "order"]:
            attrs = [getattr(q, a, None) or None for q in queries]
            def make_hashable(x):
                if not x:
                    return None
                if isinstance(x, list):
                    return tuple(x)
                return x
            attrs = set(map(make_hashable, attrs))
            if len(attrs) > 1:
                raise NotImplementedError("Cannot perform MultiQuery on "
                    "multiple %s: %s"%(a, attrs))
            attrs = filter(None, attrs)
            if len(attrs) == 1:
                self.all_attrs[a] = attrs[0]
        logging.info("DatastoreMultiQuery: %s"%self.all_attrs)
    def __repr__(self):
        return "<DatastoreMultiQuery %s>"%(", ".join(["%s=%s"%(k, v)
            for k, v in self.all_attrs.items()]))
    def fetch(self, limit, offset = 0, start_cursor = None,  end_cursor = None):
        if start_cursor or end_cursor:
            raise NotImplementedError("Cannot handle cursors.")
        return MultiIterator([q.fetch(limit=limit, offset = offset).__iter__()
            for q in self.queries], order = self.order)
    def add_filter(*args, **kwargs):
        for q in self.queries:
            q.add_filter(*args, **kwargs)
    def keys_only(self):
        for q in self.queries:
            q.keys_only()
            
class MultiIterator(object):
    def __init__(self, its, order):
        if not isinstance(order, list):
            raise NotImplementedError("Expecting a list: %s"%order)
        if len(order) > 1:
            raise NotImplementedError("Can't handle ordering: %s"%order)
        self.order = order
        self.its = its
        self.cur_items = [None]*len(its)
        self.active = [True]*len(its)
        self.started = False
        self.last_iter = None
    
    def __iter__(self):
        while True:
            # Advance the iterator.
            if not self.started:
                # The first time we start, we have to increment all items.
                for it_i, it in enumerate(self.its):
                    try:
                        self.cur_items[it_i] = next(it)
                    except StopIteration:
                        self.active[it_i] = False
                self.started = True
            elif self.active[self.last_iter]:
                # After, we need to increment just the last iterator we used.
                try:
                    self.cur_items[self.last_iter] = next(it)
                except StopIteration:
                    self.active[self.last_iter] = False
            
            # Get all active items.
            cur_items = [(it_i, item) for it_i, item in
                enumerate(self.cur_items) if self.active[it_i]]
            
            # Order the items within the array.
            if self.order and self.order[0]:
                o = self.order[0]
                reverse = False
                if o.startswith("-"):
                    o = o[1:]
                    reverse = True
                # Sort them according to sort order.
                cur_items = sorted(cur_items,
                    key = lambda it_i_i: it_i_i[1].get(o), reverse = reverse)
            
            # Get the first item in the sorted array.
            if cur_items:
                self.last_iter = cur_items[0][0]
                yield cur_items[0][1]
            else:
                self.last_iter = None
                raise StopIteration