Merge pull request #204 from tseaver/feature-datastore_query_cursor

Add cursor handling to datastore.query.Query.

gcloud/datastore/connection.py

@@ -186,7 +186,7 @@ def run_query(self, dataset_id, query_pb, namespace=None):
  Under the hood this is doing...
 
  >>> connection.run_query('dataset-id', query.to_protobuf())
- [<list of Entity Protobufs>]
+ [<list of Entity Protobufs>], cursor, more_results, skipped_results
 
  :type dataset_id: string
  :param dataset_id: The ID of the dataset over which to run the query.
@@ -205,7 +205,11 @@ def run_query(self, dataset_id, query_pb, namespace=None):
  request.query.CopyFrom(query_pb)
  response = self._rpc(dataset_id, 'runQuery', request,
  datastore_pb.RunQueryResponse)
- return [e.entity for e in response.batch.entity_result]
+ return ([e.entity for e in response.batch.entity_result],
+ response.batch.end_cursor,
+ response.batch.more_results,
+ response.batch.skipped_results,
+ )
 
  def lookup(self, dataset_id, key_pbs):
  """Lookup keys from a dataset in the Cloud Datastore.

gcloud/datastore/query.py

@@ -4,6 +4,7 @@
 from gcloud.datastore import helpers
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
+import base64
 
 
 class Query(object):
@@ -53,6 +54,7 @@ class Query(object):
  def __init__(self, kind=None, dataset=None):
  self._dataset = dataset
  self._pb = datastore_pb.Query()
+ self._cursor = None
 
  if kind:
  self._pb.kind.add().name = kind
@@ -303,12 +305,55 @@ def fetch(self, limit=None):
  if limit:
  clone = self.limit(limit)
 
- entity_pbs = self.dataset().connection().run_query(
+ (entity_pbs,
+ end_cursor,
+ more_results,
+ skipped_results) = self.dataset().connection().run_query(
  query_pb=clone.to_protobuf(), dataset_id=self.dataset().id())
 
+ self._cursor = end_cursor
  return [Entity.from_protobuf(entity, dataset=self.dataset())
  for entity in entity_pbs]
 
+ def cursor(self):
+ """Returns cursor ID
+
+ .. Caution:: Invoking this method on a query that has not yet been
+ executed will raise a RuntimeError.
+
+ :rtype: string
+ :returns: base64-encoded cursor ID string denoting the last position
+ consumed in the query's result set.
+ """
+ if not self._cursor:
+ raise RuntimeError('No cursor')
+ return base64.b64encode(self._cursor)
+
+ def with_cursor(self, start_cursor, end_cursor=None):
+ """Specifies the starting / ending positions in a query's result set.
+
+ :type start_cursor: bytes
+ :param start_cursor: Base64-encoded cursor string specifying where to
+ start reading query results.
+
+ :type end_cursor: bytes
+ :param end_cursor: Base64-encoded cursor string specifying where to stop
+ reading query results.
+
+ :rtype: :class:`Query`
+ :returns: If neither cursor is passed, returns self; else, returns a
+ clone of the :class:`Query`, with cursors updated.
+
+ """
+ clone = self
+ if start_cursor or end_cursor:
+ clone = self._clone()
+ if start_cursor:
+ clone._pb.start_cursor = base64.b64decode(start_cursor)
+ if end_cursor:
+ clone._pb.end_cursor = base64.b64decode(end_cursor)
+ return clone
+
  def order(self, *properties):
  """Adds a sort order to the query.
  

gcloud/datastore/test_connection.py

@@ -322,7 +322,11 @@ def test_run_query_wo_namespace_empty_result(self):
  'runQuery',
  ])
  http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
- self.assertEqual(conn.run_query(DATASET_ID, q_pb), [])
+ pbs, end, more, skipped = conn.run_query(DATASET_ID, q_pb)
+ self.assertEqual(pbs, [])
+ self.assertEqual(end, '')
+ self.assertTrue(more)
+ self.assertEqual(skipped, 0)
  cw = http._called_with
  self.assertEqual(cw['uri'], URI)
  self.assertEqual(cw['method'], 'POST')
@@ -357,8 +361,8 @@ def test_run_query_w_namespace_nonempty_result(self):
  'runQuery',
  ])
  http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
- result = conn.run_query(DATASET_ID, q_pb, 'NS')
- returned, = result # One entity.
+ pbs, end, more, skipped = conn.run_query(DATASET_ID, q_pb, 'NS')
+ returned, = pbs, # One entity.
  cw = http._called_with
  self.assertEqual(cw['uri'], URI)
  self.assertEqual(cw['method'], 'POST')

gcloud/datastore/test_query.py

@@ -202,6 +202,7 @@ def test_fetch_default_limit(self):
 
  def test_fetch_explicit_limit(self):
  from gcloud.datastore.datastore_v1_pb2 import Entity
+ _CURSOR = 'CURSOR'
  _DATASET = 'DATASET'
  _KIND = 'KIND'
  _ID = 123
@@ -213,10 +214,12 @@ def test_fetch_explicit_limit(self):
  prop.name = 'foo'
  prop.value.string_value = 'Foo'
  connection = _Connection(entity_pb)
+ connection._cursor = _CURSOR
  dataset = _Dataset(_DATASET, connection)
  query = self._makeOne(_KIND, dataset)
  limited = query.limit(13)
  entities = query.fetch(13)
+ self.assertEqual(query._cursor, _CURSOR)
  self.assertEqual(len(entities), 1)
  self.assertEqual(entities[0].key().path(),
  [{'kind': _KIND, 'id': _ID}])
@@ -225,6 +228,80 @@ def test_fetch_explicit_limit(self):
  'query_pb': limited.to_protobuf(),
  })
 
+ def test_cursor_not_fetched(self):
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ self.assertRaises(RuntimeError, query.cursor)
+
+ def test_cursor_fetched(self):
+ import base64
+ _CURSOR = 'CURSOR'
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ query._cursor = _CURSOR
+ self.assertEqual(query.cursor(), base64.b64encode(_CURSOR))
+
+ def test_with_cursor_neither(self):
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ self.assertTrue(query.with_cursor(None) is query)
+
+ def test_with_cursor_w_start(self):
+ import base64
+ _CURSOR = 'CURSOR'
+ _CURSOR_B64 = base64.b64encode(_CURSOR)
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ after = query.with_cursor(_CURSOR_B64)
+ self.assertFalse(after is query)
+ q_pb = after.to_protobuf()
+ self.assertEqual(q_pb.start_cursor, _CURSOR)
+ self.assertEqual(q_pb.end_cursor, '')
+
+ def test_with_cursor_w_end(self):
+ import base64
+ _CURSOR = 'CURSOR'
+ _CURSOR_B64 = base64.b64encode(_CURSOR)
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ after = query.with_cursor(None, _CURSOR_B64)
+ self.assertFalse(after is query)
+ q_pb = after.to_protobuf()
+ self.assertEqual(q_pb.start_cursor, '')
+ self.assertEqual(q_pb.end_cursor, _CURSOR)
+
+ def test_with_cursor_w_both(self):
+ import base64
+ _START = 'START'
+ _START_B64 = base64.b64encode(_START)
+ _END = 'CURSOR'
+ _END_B64 = base64.b64encode(_END)
+ _DATASET = 'DATASET'
+ _KIND = 'KIND'
+ connection = _Connection()
+ dataset = _Dataset(_DATASET, connection)
+ query = self._makeOne(_KIND, dataset)
+ after = query.with_cursor(_START_B64, _END_B64)
+ self.assertFalse(after is query)
+ q_pb = after.to_protobuf()
+ self.assertEqual(q_pb.start_cursor, _START)
+ self.assertEqual(q_pb.end_cursor, _END)
+
  def test_order_empty(self):
  _KIND = 'KIND'
  before = self._makeOne(_KIND)
@@ -285,10 +362,13 @@ def connection(self):
 
 class _Connection(object):
  _called_with = None
+ _cursor = ''
+ _more = True
+ _skipped = 0
 
  def __init__(self, *result):
  self._result = list(result)
 
  def run_query(self, **kw):
  self._called_with = kw
- return self._result
+ return self._result, self._cursor, self._more, self._skipped