From ddc8f8536a0e6edd78f4e7e4f23b4604afd7d7dc Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Sat, 30 Jan 2016 10:28:11 -0800
Subject: [PATCH 1/2] Adding key_filter to datastore Query.

Also updating system tests and system test helpers to
avoid direct use of __key__.

Fixes #1432.
---
 gcloud/datastore/query.py       | 12 ++++++++++++
 gcloud/datastore/test_query.py  | 20 ++++++++++++++++++++
 system_tests/clear_datastore.py |  3 ++-
 system_tests/datastore.py       |  4 ++--
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/gcloud/datastore/query.py b/gcloud/datastore/query.py
index f51f08070c50..a1d012b65061 100644
--- 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
index 83b84a8f4364..a16140fefd12 100644
--- 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
index 8c57022a26fc..85a7b08c1914 100644
--- 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
index 44c2b06d0477..d995426d9983 100644
--- 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))

From 2b0d2453c2dc384725dd286763726805a282fa30 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Sat, 30 Jan 2016 10:37:29 -0800
Subject: [PATCH 2/2] Fixing broken system test helpers.

They were using the dataset ID environment variable key rather
than the value.
---
 system_tests/clear_datastore.py    | 2 +-
 system_tests/populate_datastore.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/system_tests/clear_datastore.py b/system_tests/clear_datastore.py
index 85a7b08c1914..f544dd484cbb 100644
--- a/system_tests/clear_datastore.py
+++ b/system_tests/clear_datastore.py
@@ -90,7 +90,7 @@ def remove_kind(kind, client):
 def remove_all_entities(client=None):
     if client is None:
         # Get a client that uses the test dataset.
-        client = datastore.Client(project=TESTS_DATASET)
+        client = datastore.Client(project=os.getenv(TESTS_DATASET))
     for kind in ALL_KINDS:
         remove_kind(kind, client)
 
diff --git a/system_tests/populate_datastore.py b/system_tests/populate_datastore.py
index 3c349fdff074..e4be6cfbadad 100644
--- a/system_tests/populate_datastore.py
+++ b/system_tests/populate_datastore.py
@@ -91,7 +91,7 @@ def print_func(message):
 def add_characters(client=None):
     if client is None:
         # Get a client that uses the test dataset.
-        client = datastore.Client(project=TESTS_DATASET)
+        client = datastore.Client(project=os.getenv(TESTS_DATASET))
     with client.transaction() as xact:
         for key_path, character in zip(KEY_PATHS, CHARACTERS):
             if key_path[-1] != character['name']: