Implements group_by, projection and offset on datastore Query.

In addition, __eq__ was implemented on datastore.key.Key
to allow for easy comparison.

Author: dhermes

gcloud/datastore/key.py

@@ -251,3 +251,14 @@ def parent(self):
 
     def __repr__(self):
         return '<Key%s>' % self.path()
+
+    def __eq__(self, other):
+        if self is other:
+            return True
+
+        return (self._dataset_id == other._dataset_id and
+                self.namespace() == other.namespace() and
+                self.path() == other.path())
+
+    def __ne__(self, other):
+        return not self.__eq__(other)

gcloud/datastore/query.py

@@ -60,6 +60,9 @@ def __init__(self, kind=None, dataset=None, namespace=None):
         self._namespace = namespace
         self._pb = datastore_pb.Query()
         self._cursor = None
+        self._projection = []
+        self._offset = 0
+        self._group_by = []
 
         if kind:
             self._pb.kind.add().name = kind
@@ -404,3 +407,104 @@ 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 self._projection
+
+        clone = self._clone()
+        clone._projection = projection
+
+        # Reset projection values to empty.
+        clone._pb.projection._values = []
+
+        # 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 self._group_by
+
+        clone = self._clone()
+        clone._group_by = group_by
+
+        # Reset group_by values to empty.
+        clone._pb.group_by._values = []
+
+        # 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_key.py

@@ -322,3 +322,13 @@ def test_parent_explicit_top_level(self):
     def test_parent_explicit_nested(self):
         key = self._getTargetClass().from_path('abc', 'def', 'ghi', 123)
         self.assertEqual(key.parent().path(), [{'kind': 'abc', 'name': 'def'}])
+
+    def test_key___eq__(self):
+        key1 = self._getTargetClass().from_path('abc', 'def')
+        key2 = self._getTargetClass().from_path('abc', 'def')
+        self.assertFalse(key1 is key2)
+        self.assertEqual(key1, key2)
+
+        self.assertEqual(key1, key1)
+        key3 = self._getTargetClass().from_path('abc', 'ghi')
+        self.assertNotEqual(key1, key3)

gcloud/datastore/test_query.py

@@ -415,6 +415,70 @@ 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_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)
+
 
 class _Dataset(object):