Fix #170: handle namespaces in queries

gcloud/datastore/query.py

@@ -41,6 +41,9 @@ class Query(object):
 
     :type dataset: :class:`gcloud.datastore.dataset.Dataset`
     :param dataset: The dataset to query.
+
+    :type namespace: string or None
+    :param dataset: The namespace to which to restrict results.
     """
 
     OPERATORS = {
@@ -52,8 +55,9 @@ class Query(object):
     }
     """Mapping of operator strings and their protobuf equivalents."""
 
-    def __init__(self, kind=None, dataset=None):
+    def __init__(self, kind=None, dataset=None, namespace=None):
         self._dataset = dataset
+        self._namespace = namespace
         self._pb = datastore_pb.Query()
         self._cursor = None
 
@@ -66,11 +70,20 @@ def _clone(self):
         :rtype: :class:`gcloud.datastore.query.Query`
         :returns: a copy of 'self'.
         """
-        clone = self.__class__(dataset=self._dataset)
+        clone = self.__class__(dataset=self._dataset,
+                               namespace=self._namespace)
         clone._pb.CopyFrom(self._pb)
         clone._cursor = self._cursor
         return clone
 
+    def namespace(self):
+        """This query's namespace
+
+        :rtype: string or None
+        :returns: the namespace assigned to this query
+        """
+        return self._namespace
+
     def to_protobuf(self):
         """Convert :class:`Query` instance to :class:`.datastore_v1_pb2.Query`.
 
@@ -316,7 +329,10 @@ def fetch(self, limit=None):
             clone = self.limit(limit)
 
         query_results = self.dataset().connection().run_query(
-            query_pb=clone.to_protobuf(), dataset_id=self.dataset().id())
+            query_pb=clone.to_protobuf(),
+            dataset_id=self.dataset().id(),
+            namespace=self._namespace,
+            )
         entity_pbs, end_cursor = query_results[:2]
 
         self._cursor = end_cursor

gcloud/datastore/test_query.py

@@ -8,39 +8,44 @@ def _getTargetClass(self):
 
         return Query
 
-    def _makeOne(self, kind=None, dataset=None):
-        return self._getTargetClass()(kind, dataset)
+    def _makeOne(self, kind=None, dataset=None, namespace=None):
+        return self._getTargetClass()(kind, dataset, namespace)
 
     def test_ctor_defaults(self):
-        query = self._makeOne()
+        query = self._getTargetClass()()
         self.assertEqual(query.dataset(), None)
         self.assertEqual(list(query.kind()), [])
         self.assertEqual(query.limit(), 0)
+        self.assertEqual(query.namespace(), None)
 
     def test_ctor_explicit(self):
         from gcloud.datastore.dataset import Dataset
 
         _DATASET = 'DATASET'
         _KIND = 'KIND'
+        _NAMESPACE = 'NAMESPACE'
         dataset = Dataset(_DATASET)
-        query = self._makeOne(_KIND, dataset)
+        query = self._makeOne(_KIND, dataset, _NAMESPACE)
         self.assertTrue(query.dataset() is dataset)
         kq_pb, = list(query.kind())
         self.assertEqual(kq_pb.name, _KIND)
+        self.assertEqual(query.namespace(), _NAMESPACE)
 
     def test__clone(self):
         from gcloud.datastore.dataset import Dataset
 
         _DATASET = 'DATASET'
         _KIND = 'KIND'
         _CURSOR = 'DEADBEEF'
+        _NAMESPACE = 'NAMESPACE'
         dataset = Dataset(_DATASET)
-        query = self._makeOne(_KIND, dataset)
+        query = self._makeOne(_KIND, dataset, _NAMESPACE)
         query._cursor = _CURSOR
         clone = query._clone()
         self.assertFalse(clone is query)
         self.assertTrue(isinstance(clone, self._getTargetClass()))
         self.assertTrue(clone.dataset() is dataset)
+        self.assertEqual(clone.namespace(), _NAMESPACE)
         kq_pb, = list(clone.kind())
         self.assertEqual(kq_pb.name, _KIND)
         self.assertEqual(clone._cursor, _CURSOR)
@@ -101,7 +106,7 @@ def test_ancestor_w_non_key_non_list(self):
         query = self._makeOne()
         self.assertRaises(TypeError, query.ancestor, object())
 
-    def test_ancester_wo_existing_ancestor_query_w_key_and_propfilter(self):
+    def test_ancestor_wo_existing_ancestor_query_w_key_and_propfilter(self):
         from gcloud.datastore.key import Key
         _KIND = 'KIND'
         _ID = 123
@@ -121,7 +126,7 @@ def test_ancester_wo_existing_ancestor_query_w_key_and_propfilter(self):
         self.assertEqual(p_pb.property.name, '__key__')
         self.assertEqual(p_pb.value.key_value, key.to_protobuf())
 
-    def test_ancester_wo_existing_ancestor_query_w_key(self):
+    def test_ancestor_wo_existing_ancestor_query_w_key(self):
         from gcloud.datastore.key import Key
         _KIND = 'KIND'
         _ID = 123
@@ -137,7 +142,7 @@ def test_ancester_wo_existing_ancestor_query_w_key(self):
         self.assertEqual(p_pb.property.name, '__key__')
         self.assertEqual(p_pb.value.key_value, key.to_protobuf())
 
-    def test_ancester_wo_existing_ancestor_query_w_list(self):
+    def test_ancestor_wo_existing_ancestor_query_w_list(self):
         from gcloud.datastore.key import Key
         _KIND = 'KIND'
         _ID = 123
@@ -153,7 +158,7 @@ def test_ancester_wo_existing_ancestor_query_w_list(self):
         self.assertEqual(p_pb.property.name, '__key__')
         self.assertEqual(p_pb.value.key_value, key.to_protobuf())
 
-    def test_ancester_clears_existing_ancestor_query_w_only(self):
+    def test_ancestor_clears_existing_ancestor_query_w_only(self):
         _KIND = 'KIND'
         _ID = 123
         query = self._makeOne()
@@ -164,7 +169,7 @@ def test_ancester_clears_existing_ancestor_query_w_only(self):
         q_pb = after.to_protobuf()
         self.assertEqual(list(q_pb.filter.composite_filter.filter), [])
 
-    def test_ancester_clears_existing_ancestor_query_w_others(self):
+    def test_ancestor_clears_existing_ancestor_query_w_others(self):
         _KIND = 'KIND'
         _ID = 123
         _NAME = 'NAME'
@@ -257,6 +262,7 @@ def test_fetch_default_limit(self):
         expected_called_with = {
             'dataset_id': _DATASET,
             'query_pb': query.to_protobuf(),
+            'namespace': None,
         }
         self.assertEqual(connection._called_with, expected_called_with)
 
@@ -266,6 +272,7 @@ def test_fetch_explicit_limit(self):
         _DATASET = 'DATASET'
         _KIND = 'KIND'
         _ID = 123
+        _NAMESPACE = 'NAMESPACE'
         entity_pb = Entity()
         path_element = entity_pb.key.path_element.add()
         path_element.kind = _KIND
@@ -276,7 +283,7 @@ def test_fetch_explicit_limit(self):
         connection = _Connection(entity_pb)
         connection._cursor = _CURSOR
         dataset = _Dataset(_DATASET, connection)
-        query = self._makeOne(_KIND, dataset)
+        query = self._makeOne(_KIND, dataset, _NAMESPACE)
         limited = query.limit(13)
         entities = query.fetch(13)
         self.assertEqual(query._cursor, _CURSOR)
@@ -286,6 +293,7 @@ def test_fetch_explicit_limit(self):
         expected_called_with = {
             'dataset_id': _DATASET,
             'query_pb': limited.to_protobuf(),
+            'namespace': _NAMESPACE,
         }
         self.assertEqual(connection._called_with, expected_called_with)