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
+
+