Implements group_by, projection and offset on datastore Query.

gcloud/datastore/query.py

@@ -59,6 +59,7 @@ def __init__(self, kind=None, dataset=None, namespace=None):
         self._namespace = namespace
         self._pb = datastore_pb.Query()
         self._cursor = None
+        self._offset = 0
 
         if kind:
             self._pb.kind.add().name = kind
@@ -414,3 +415,103 @@ def order(self, *properties):
                 property_order.direction = property_order.ASCENDING
 
         return clone
+
+    def projection(self, projection=None):
+        """Adds a projection to the query.
+
+        This is a hybrid getter / setter, used as::
+
+          >>> query = Query('Person')
+          >>> query.projection()  # Get the projection for this query.
+          []
+          >>> query = query.projection(['name'])
+          >>> query.projection()  # Get the projection for this query.
+          ['name']
+
+        :type projection: sequence of strings
+        :param projection: Each value is a string giving the name of a
+                           property to be included in the projection query.
+
+        :rtype: :class:`Query` or `list` of strings.
+        :returns: If no arguments, returns the current projection.
+                  If a projection is provided, returns a clone of the
+                  :class:`Query` with that projection set.
+        """
+        if projection is None:
+            return [prop_expr.property.name
+                    for prop_expr in self._pb.projection]
+
+        clone = self._clone()
+
+        # Reset projection values to empty.
+        clone._pb.ClearField('projection')
+
+        # Add each name to list of projections.
+        for projection_name in projection:
+            clone._pb.projection.add().property.name = projection_name
+        return clone
+
+    def offset(self, offset=None):
+        """Adds offset to the query to allow pagination.
+
+        NOTE: Paging with cursors should be preferred to using an offset.
+
+        This is a hybrid getter / setter, used as::
+
+          >>> query = Query('Person')
+          >>> query.offset()  # Get the offset for this query.
+          0
+          >>> query = query.offset(10)
+          >>> query.offset()  # Get the offset for this query.
+          10
+
+        :type offset: non-negative integer.
+        :param offset: Value representing where to start a query for
+                       a given kind.
+
+        :rtype: :class:`Query` or `int`.
+        :returns: If no arguments, returns the current offset.
+                  If an offset is provided, returns a clone of the
+                  :class:`Query` with that offset set.
+        """
+        if offset is None:
+            return self._offset
+
+        clone = self._clone()
+        clone._offset = offset
+        clone._pb.offset = offset
+        return clone
+
+    def group_by(self, group_by=None):
+        """Adds a group_by to the query.
+
+        This is a hybrid getter / setter, used as::
+
+          >>> query = Query('Person')
+          >>> query.group_by()  # Get the group_by for this query.
+          []
+          >>> query = query.group_by(['name'])
+          >>> query.group_by()  # Get the group_by for this query.
+          ['name']
+
+        :type group_by: sequence of strings
+        :param group_by: Each value is a string giving the name of a
+                         property to use to group results together.
+
+        :rtype: :class:`Query` or `list` of strings.
+        :returns: If no arguments, returns the current group_by.
+                  If a list of group by properties is provided, returns a clone
+                  of the :class:`Query` with that list of values set.
+        """
+        if group_by is None:
+            return [prop_ref.name for prop_ref in self._pb.group_by]
+
+        clone = self._clone()
+
+        # Reset group_by values to empty.
+        clone._pb.ClearField('group_by')
+
+        # Add each name to list of group_bys.
+        for group_by_name in group_by:
+            clone._pb.group_by.add().name = group_by_name
+        return clone

gcloud/datastore/test_query.py

@@ -451,6 +451,88 @@ def test_order_multiple(self):
         self.assertEqual(prop_pb.property.name, 'bar')
         self.assertEqual(prop_pb.direction, prop_pb.DESCENDING)
 
+    def test_projection_empty(self):
+        _KIND = 'KIND'
+        before = self._makeOne(_KIND)
+        after = before.projection([])
+        self.assertFalse(after is before)
+        self.assertTrue(isinstance(after, self._getTargetClass()))
+        self.assertEqual(before.to_protobuf(), after.to_protobuf())
+
+    def test_projection_non_empty(self):
+        _KIND = 'KIND'
+        before = self._makeOne(_KIND)
+        after = before.projection(['field1', 'field2'])
+        projection_pb = list(after.to_protobuf().projection)
+        self.assertEqual(len(projection_pb), 2)
+        prop_pb1 = projection_pb[0]
+        self.assertEqual(prop_pb1.property.name, 'field1')
+        prop_pb2 = projection_pb[1]
+        self.assertEqual(prop_pb2.property.name, 'field2')
+
+    def test_get_projection_non_empty(self):
+        _KIND = 'KIND'
+        _PROJECTION = ['field1', 'field2']
+        after = self._makeOne(_KIND).projection(_PROJECTION)
+        self.assertEqual(after.projection(), _PROJECTION)
+
+    def test_projection_multiple_calls(self):
+        _KIND = 'KIND'
+        _PROJECTION1 = ['field1', 'field2']
+        _PROJECTION2 = ['field3']
+        before = self._makeOne(_KIND).projection(_PROJECTION1)
+        self.assertEqual(before.projection(), _PROJECTION1)
+        after = before.projection(_PROJECTION2)
+        self.assertEqual(after.projection(), _PROJECTION2)
+
+    def test_set_offset(self):
+        _KIND = 'KIND'
+        _OFFSET = 42
+        before = self._makeOne(_KIND)
+        after = before.offset(_OFFSET)
+        offset_pb = after.to_protobuf().offset
+        self.assertEqual(offset_pb, _OFFSET)
+
+    def test_get_offset(self):
+        _KIND = 'KIND'
+        _OFFSET = 10
+        after = self._makeOne(_KIND).offset(_OFFSET)
+        self.assertEqual(after.offset(), _OFFSET)
+
+    def test_group_by_empty(self):
+        _KIND = 'KIND'
+        before = self._makeOne(_KIND)
+        after = before.group_by([])
+        self.assertFalse(after is before)
+        self.assertTrue(isinstance(after, self._getTargetClass()))
+        self.assertEqual(before.to_protobuf(), after.to_protobuf())
+
+    def test_group_by_non_empty(self):
+        _KIND = 'KIND'
+        before = self._makeOne(_KIND)
+        after = before.group_by(['field1', 'field2'])
+        group_by_pb = list(after.to_protobuf().group_by)
+        self.assertEqual(len(group_by_pb), 2)
+        prop_pb1 = group_by_pb[0]
+        self.assertEqual(prop_pb1.name, 'field1')
+        prop_pb2 = group_by_pb[1]
+        self.assertEqual(prop_pb2.name, 'field2')
+
+    def test_get_group_by_non_empty(self):
+        _KIND = 'KIND'
+        _GROUP_BY = ['field1', 'field2']
+        after = self._makeOne(_KIND).group_by(_GROUP_BY)
+        self.assertEqual(after.group_by(), _GROUP_BY)
+
+    def test_group_by_multiple_calls(self):
+        _KIND = 'KIND'
+        _GROUP_BY1 = ['field1', 'field2']
+        _GROUP_BY2 = ['field3']
+        before = self._makeOne(_KIND).group_by(_GROUP_BY1)
+        self.assertEqual(before.group_by(), _GROUP_BY1)
+        after = before.group_by(_GROUP_BY2)
+        self.assertEqual(after.group_by(), _GROUP_BY2)
+
 
 class _Dataset(object):