Adding key_filter to datastore Query.

dhermes · dhermes · commit ddc8f8536a0e · 2016-01-30T10:28:11.000-08:00
Also updating system tests and system test helpers to avoid direct use of __key__. Fixes googleapis#1432.
diff --git a/gcloud/datastore/query.py b/gcloud/datastore/query.py
@@ -248,6 +248,18 @@ def keys_only(self):
         """Set the projection to include only keys."""
         self._projection[:] = ['__key__']
 
+    def key_filter(self, key, operator='='):
+        """Filter on a key.
+
+        :type key: :class:`gcloud.datastore.key.Key`
+        :param key: The key to filter on.
+
+        :type operator: string
+        :param operator: (Optional) One of ``=``, ``<``, ``<=``, ``>``, ``>=``.
+                         Defaults to ``=``.
+        """
+        self.add_filter('__key__', operator, key)
+
     @property
     def order(self):
         """Names of fields used to sort query results.
diff --git a/gcloud/datastore/test_query.py b/gcloud/datastore/test_query.py
@@ -244,6 +244,26 @@ def test_keys_only(self):
         query.keys_only()
         self.assertEqual(query.projection, ['__key__'])
 
+    def test_key_filter_defaults(self):
+        from gcloud.datastore.key import Key
+
+        client = self._makeClient()
+        query = self._makeOne(client)
+        self.assertEqual(query.filters, [])
+        key = Key('Kind', 1234, project='project')
+        query.key_filter(key)
+        self.assertEqual(query.filters, [('__key__', '=', key)])
+
+    def test_key_filter_explicit(self):
+        from gcloud.datastore.key import Key
+
+        client = self._makeClient()
+        query = self._makeOne(client)
+        self.assertEqual(query.filters, [])
+        key = Key('Kind', 1234, project='project')
+        query.key_filter(key, operator='>')
+        self.assertEqual(query.filters, [('__key__', '>', key)])
+
     def test_order_setter_empty(self):
         query = self._makeOne(self._makeClient(), order=['foo', '-bar'])
         query.order = []
diff --git a/system_tests/clear_datastore.py b/system_tests/clear_datastore.py
@@ -42,7 +42,8 @@ def print_func(message):
 
 def fetch_keys(kind, client, fetch_max=FETCH_MAX, query=None, cursor=None):
     if query is None:
-        query = client.query(kind=kind, projection=['__key__'])
+        query = client.query(kind=kind)
+        query.keys_only()
 
     iterator = query.fetch(limit=fetch_max, start_cursor=cursor)
 
diff --git a/system_tests/datastore.py b/system_tests/datastore.py
@@ -273,12 +273,12 @@ def test_ancestor_query(self):
         entities = list(filtered_query.fetch(limit=expected_matches + 1))
         self.assertEqual(len(entities), expected_matches)
 
-    def test_query___key___filter(self):
+    def test_query_key_filter(self):
         # Use the client for this test instead of the global.
         rickard_key = self.CLIENT.key(*populate_datastore.RICKARD)
 
         query = self._base_query()
-        query.add_filter('__key__', '=', rickard_key)
+        query.key_filter(rickard_key)
         expected_matches = 1
         # We expect 1, but allow the query to get 1 extra.
         entities = list(query.fetch(limit=expected_matches + 1))
