Add Support for Query Cursor & Order

gcloud/datastore/connection.py

@@ -208,7 +208,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.
@@ -226,7 +226,7 @@ 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/helpers.py

@@ -60,7 +60,8 @@ def get_protobuf_attribute_and_value(val):
   return name + '_value', value
 
 
-def get_value_from_protobuf(pb):
+def get_value_from_value_pb(value_pb):
+  from entity import Entity
   """Given a protobuf for a Property, get the correct value.
 
   The Cloud Datastore Protobuf API returns a Property Protobuf
@@ -75,27 +76,55 @@ def get_value_from_protobuf(pb):
 
   :returns: The value provided by the Protobuf.
   """
-
-  if pb.value.HasField('timestamp_microseconds_value'):
-    microseconds = pb.value.timestamp_microseconds_value
+  if value_pb.HasField('timestamp_microseconds_value'):
+    microseconds = value_pb.timestamp_microseconds_value
     return (datetime.utcfromtimestamp(0) +
             timedelta(microseconds=microseconds))
 
-  elif pb.value.HasField('key_value'):
-    return Key.from_protobuf(pb.value.key_value)
+  elif value_pb.HasField('key_value'):
+    return Key.from_protobuf(value_pb.key_value)
+
+  elif value_pb.HasField('boolean_value'):
+    return value_pb.boolean_value
+
+  elif value_pb.HasField('double_value'):
+    return value_pb.double_value
+
+  elif value_pb.HasField('integer_value'):
+    return value_pb.integer_value
+
+  elif value_pb.HasField('string_value'):
+    return value_pb.string_value
 
-  elif pb.value.HasField('boolean_value'):
-    return pb.value.boolean_value
+  elif value_pb.HasField('blob_key_value'):
+    return value_pb.blob_key_value
 
-  elif pb.value.HasField('double_value'):
-    return pb.value.double_value
+  elif value_pb.HasField('blob_value'):
+    return value_pb.blob_value
 
-  elif pb.value.HasField('integer_value'):
-    return pb.value.integer_value
+  elif value_pb.HasField('entity_value'):
+    return Entity.from_protobuf(value_pb.entity_value)
 
-  elif pb.value.HasField('string_value'):
-    return pb.value.string_value
+  elif value_pb.list_value:
+    return [get_value_from_value_pb(k) for k in value_pb.list_value]
 
   else:
     # TODO(jjg): Should we raise a ValueError here?
     return None
+
+def get_value_from_protobuf(pb):
+  """Given a protobuf for a Property, get the correct value.
+
+  The Cloud Datastore Protobuf API returns a Property Protobuf
+  which has one value set and the rest blank.
+  This method retrieves the the one value provided.
+
+  Some work is done to coerce the return value into a more useful type
+  (particularly in the case of a timestamp value, or a key value).
+
+  :type pb: :class:`gcloud.datastore.datastore_v1_pb2.Property`
+  :param pb: The Property Protobuf.
+
+  :returns: The value provided by the Protobuf.
+  """
+  return get_value_from_value_pb(pb.value)

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
 
 
 # TODO: Figure out how to properly handle namespaces.
@@ -55,6 +56,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
@@ -308,8 +310,66 @@ 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 a base64-encoded cursor string denoting the position in the query's result
+    set following the last result retrieved.
+
+    .. Caution:: Invoking this method on a query that has not yet has been
+      executed will raise an AssertionError exception.
+
+    :rtype: string
+    :returns: The lastest end_cursor for query
+    """
+    assert self._cursor
+    return base64.b64encode(self._cursor)
+
+  def with_cursor(self, start_cursor, end_cursor=None):
+    """Specifies the starting and (optionally) ending positions within a query's
+    result set from which to retrieve results.
+
+    :type start_cursor: bytes
+    :param start_cursor: Base64-encoded cursor string specifying where to start the query.
+
+    :type end_cursor: bytes
+    :param end_cursor: Base64-encoded cursor string specifying where to end the query.
+
+    """
+    if start_cursor:
+      self._pb.start_cursor = base64.b64decode(start_cursor)
+    if end_cursor:
+      self._pb.end_cursor = base64.b64decode(end_cursor)
+
+  def order(self, *properties):
+    """Adds a sort order to the query. If more than one sort order is added,
+    they will be applied in the order specified.
+
+    :type properties: string
+    :param properties: String giving the name of the property on which to sort,
+    optionally preceded by a hyphen (-) to specify descending order.
+    Omitting the hyphen specifies ascending order by default.
+
+    :rtype: :class:`Query`
+    :returns: A Query order by properties.
+    """
+    clone = self._clone()
+
+    for p in properties:
+      property_order = clone._pb.order.add()
+
+      if p.startswith('-'):
+        property_order.property.name = p[1:]
+        property_order.direct = property_order.DESCENDING
+      else:
+        property_order.property.name = p
+        property_order.direction = property_order.ASCENDING
+
+    return clone
+
+