30 January 2013

ndb OR query with cursors

I have two annoyances with Appengine's NDB OR queries.  One is that the order must be by key, the other is that OR queries are done in series, so ndb.OR(a == 1, a == 2) are done as query(a == 1) then query(a == 2).  To work around this, I created a new cursor class that allows sorting by any attribute.

Using this cursor does come with it's own caveats, one is that the order between the queries is not maintained (fixable) and two the limit attribute is per query not for the results of all queries (fixable).

import base64
import json

from google.appengine.ext import ndb

class MultiCursor(object):
    @staticmethod
    def from_websafe_string(value):
        values = None
        try:
            value = base64.urlsafe_b64decode(str(value))
            values = json.loads(value)
        except Exception as e:
            logging.error('Invalid Cursor %s - %s', e, value)
        if not values:
            return MultiCursor()

        cursors = MultiCursor()
        for key, cursor in values.iteritems():
            try:
                if isinstance(cursor, basestring):
                    cursors.set(key, ndb.Cursor.from_websafe_string(cursor))
                else:
                    cursors.set(key, cursor)
            except:
                logging.error('MultiCursor.from_string - bad cursor (%s) for %s in %s',
                              cursor, key, values)
        return cursors

    def to_websafe_string(self):
        values = {}
        for key, value in self.values:
            if isinstance(value, ndb.Cursor):
                values[key] = value.to_websafe_string()
            else:
                values[key] = value
        return base64.urlsafe_b64encode(json.dumps(values))

    def __init__(self, values=None):
        self.values = values or {}

    def get(self, key):
        return self.values.get(key), self.values.get(key + _'done')

    def set(self, key, cursor, more):
        self.values[key] = cursor
        self.values[key + _'done'] = not more

    def __len__(self):
        return len(self.values)

To use the cursor the code should look like:
   @ndb.tasklet
    def _query_A_by_value_async(value, cursor, limit=100):
        start_cursor, done = cursor.get(value)
        if not done:
            query = Foo.query(Foo.A == value)
            query = Foo.order(-Foo.last_access_time)
            foos, end_cursor, more = yield query.fetch_page_async(
                limit, start_cursor=start_cursor)
            cursor.set(value, end_cursor, more)
            raise ndb.Return(foos, end_cursor, more)
        raise ndb.Return([], None, False)

No comments:

Post a Comment