fix: make sure keys_only ordered multiquery returns keys not entities (#527)

Chris Rossi · web-flow · commit 0803d1350d09 · 2020-09-02T08:47:24.000-04:00
Fixes #526
diff --git a/packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py b/packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py
@@ -580,6 +580,7 @@ class _MultiQueryIteratorImpl(QueryIterator):
     """
 
     _extra_projections = None
+    _coerce_keys_only = False
 
     def __init__(self, query, raw=False):
         projection = query.projection
@@ -592,10 +593,12 @@ def __init__(self, query, raw=False):
             extra_projections = []
             for order in query.order_by:
                 if order.name not in projection:
-                    projection.append(order.name)
                     extra_projections.append(order.name)
 
             if extra_projections:
+                if projection == ["__key__"]:
+                    self._coerce_keys_only = True
+                projection.extend(extra_projections)
                 self._extra_projections = extra_projections
 
         queries = [
@@ -707,7 +710,10 @@ def next(self):
         if self._raw:
             return next_result
         else:
-            return next_result.entity()
+            entity = next_result.entity()
+            if self._coerce_keys_only:
+                return entity._key
+            return entity
 
     __next__ = next
 
diff --git a/packages/google-cloud-ndb/tests/system/test_query.py b/packages/google-cloud-ndb/tests/system/test_query.py
@@ -673,7 +673,7 @@ class SomeKind(ndb.Model):
     results = eventually(
         functools.partial(query.fetch, keys_only=True), length_equals(5)
     )
-    assert keys == [entity.key for entity in results]
+    assert keys == results
 
 
 @pytest.mark.usefixtures("client_context")
diff --git a/packages/google-cloud-ndb/tests/unit/test__datastore_query.py b/packages/google-cloud-ndb/tests/unit/test__datastore_query.py
@@ -910,6 +910,30 @@ def test_constructor_sortable_with_projection_needs_extra():
             projection=["bar", "foo"],
         )
         assert iterator._sortable
+        assert not iterator._coerce_keys_only
+
+    @staticmethod
+    def test_constructor_sortable_with_projection_needs_extra_keys_only():
+        foo = model.StringProperty("foo")
+        order_by = [query_module.PropertyOrder("foo")]
+        query = query_module.QueryOptions(
+            filters=query_module.OR(foo == "this", foo == "that"),
+            order_by=order_by,
+            projection=("__key__",),
+        )
+        iterator = _datastore_query._MultiQueryIteratorImpl(query)
+        assert iterator._result_sets[0]._query == query_module.QueryOptions(
+            filters=foo == "this",
+            order_by=order_by,
+            projection=["__key__", "foo"],
+        )
+        assert iterator._result_sets[1]._query == query_module.QueryOptions(
+            filters=foo == "that",
+            order_by=order_by,
+            projection=["__key__", "foo"],
+        )
+        assert iterator._sortable
+        assert iterator._coerce_keys_only
 
     @staticmethod
     def test_iter():
@@ -981,6 +1005,39 @@ def test_next_with_extra_projections():
         assert iterator.next() is next_result
         assert "foo" not in next_result.result_pb.entity.properties
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_next_coerce_keys_only():
+        foo = model.StringProperty("foo")
+        order_by = [
+            query_module.PropertyOrder("foo"),
+            query_module.PropertyOrder("food"),
+        ]
+        query = query_module.QueryOptions(
+            filters=query_module.OR(foo == "this", foo == "that"),
+            order_by=order_by,
+            projection=["__key__"],
+        )
+        iterator = _datastore_query._MultiQueryIteratorImpl(query)
+        iterator._next_result = next_result = mock.Mock(
+            result_pb=mock.Mock(
+                entity=mock.Mock(
+                    properties={"foo": 1, "bar": "two"},
+                    spec=("properties",),
+                ),
+                spec=("entity",),
+            ),
+            entity=mock.Mock(
+                return_value=mock.Mock(
+                    _key="thekey",
+                )
+            ),
+            spec=("result_pb", "entity"),
+        )
+
+        assert iterator.next() == "thekey"
+        assert "foo" not in next_result.result_pb.entity.properties
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     def test_iterate_async():

Original file line number	Diff line number	Diff line change
`@@ -673,7 +673,7 @@ class SomeKind(ndb.Model):`
`673`	`673`	`results = eventually(`
`674`	`674`	`functools.partial(query.fetch, keys_only=True), length_equals(5)`
`675`	`675`	`)`
`676`		`- assert keys == [entity.key for entity in results]`
	`676`	`+ assert keys == results`
`677`	`677`
`678`	`678`
`679`	`679`	`@pytest.mark.usefixtures("client_context")`