Visitar URL original
fix: transparently add sort properties to projection for multiqueries… · googleapis/google-cloud-python@491bfe1 · GitHub
Skip to content

Commit 491bfe1

Browse files
author
Chris Rossi
authored
fix: transparently add sort properties to projection for multiqueries (#511)
Since the result sets for multiqueries have to be merged, if a projection is specified, make sure the projection used for the underlying Datstore queries includes any properties we need for a sorted merge. Fixes #509.
1 parent 0148a44 commit 491bfe1

4 files changed

Lines changed: 153 additions & 1 deletion

File tree

‎packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,9 +518,29 @@ class _MultiQueryIteratorImpl(QueryIterator):
518518
internal use only.
519519
"""
520520

521+
_extra_projections = None
522+
521523
def __init__(self, query, raw=False):
524+
projection = query.projection
525+
if query.order_by and projection:
526+
# In an ordered multiquery, result sets have to be merged in order
527+
# by this iterator, so if there's a projection we may need to add a
528+
# property or two to underlying Datastore queries to make sure we
529+
# have the data needed for sorting.
530+
projection = list(projection)
531+
extra_projections = []
532+
for order in query.order_by:
533+
if order.name not in projection:
534+
projection.append(order.name)
535+
extra_projections.append(order.name)
536+
537+
if extra_projections:
538+
self._extra_projections = extra_projections
539+
522540
queries = [
523-
query.copy(filters=node, offset=None, limit=None)
541+
query.copy(
542+
filters=node, projection=projection, offset=None, limit=None
543+
)
524544
for node in query.filters._nodes
525545
]
526546
self._result_sets = [iterate(_query, raw=True) for _query in queries]
@@ -620,6 +640,14 @@ def next(self):
620640
# Won't block
621641
next_result = self._next_result
622642
self._next_result = None
643+
644+
# If we had to set extra properties in the projection, elide them now
645+
if self._extra_projections:
646+
properties = next_result.result_pb.entity.properties
647+
for name in self._extra_projections:
648+
if name in properties:
649+
del properties[name]
650+
623651
if self._raw:
624652
return next_result
625653
else:

‎packages/google-cloud-ndb/tests/system/index.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,11 @@ indexes:
55
- name: bar
66
- name: foo
77

8+
- kind: SomeKind
9+
properties:
10+
- name: foo
11+
- name: bar
12+
813
- kind: SomeKind
914
properties:
1015
- name: bar.one

‎packages/google-cloud-ndb/tests/system/test_query.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
"""
1818

1919
import datetime
20+
import functools
2021
import operator
2122

2223
import pytest
@@ -651,6 +652,56 @@ class SomeKind(ndb.Model):
651652
eventually(query.count, equals(5))
652653

653654

655+
@pytest.mark.usefixtures("client_context")
656+
def test_keys_only_multiquery_with_order(ds_entity):
657+
"""Regression test for #509
658+
659+
https://github.com/googleapis/python-ndb/issues/509
660+
"""
661+
keys = []
662+
for i in range(5):
663+
entity_id = test_utils.system.unique_resource_id()
664+
ds_entity(KIND, entity_id, foo=i)
665+
keys.append(ndb.Key(KIND, entity_id))
666+
667+
class SomeKind(ndb.Model):
668+
foo = ndb.IntegerProperty()
669+
670+
query = (
671+
SomeKind.query()
672+
.order(SomeKind.foo)
673+
.filter(ndb.OR(SomeKind.foo < 100, SomeKind.foo > -1))
674+
)
675+
results = eventually(
676+
functools.partial(query.fetch, keys_only=True), length_equals(5)
677+
)
678+
assert keys == [entity.key for entity in results]
679+
680+
681+
@pytest.mark.usefixtures("client_context")
682+
def test_multiquery_with_projection_and_order(ds_entity):
683+
"""Regression test for #509
684+
685+
https://github.com/googleapis/python-ndb/issues/509
686+
"""
687+
for i in range(5):
688+
entity_id = test_utils.system.unique_resource_id()
689+
ds_entity(KIND, entity_id, foo=i, bar="bar " + str(i))
690+
691+
class SomeKind(ndb.Model):
692+
foo = ndb.IntegerProperty()
693+
bar = ndb.StringProperty()
694+
695+
query = (
696+
SomeKind.query(projection=[SomeKind.bar])
697+
.order(SomeKind.foo)
698+
.filter(ndb.OR(SomeKind.foo < 100, SomeKind.foo > -1))
699+
)
700+
results = eventually(query.fetch, length_equals(5))
701+
with pytest.raises(ndb.UnprojectedPropertyError):
702+
results[0].foo
703+
704+
654705
@pytest.mark.usefixtures("client_context")
655706
def test_count_with_multi_query(ds_entity):
656707
for i in range(5):

‎packages/google-cloud-ndb/tests/unit/test__datastore_query.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -740,6 +740,46 @@ def test_constructor_sortable():
740740
)
741741
assert iterator._sortable
742742

743+
@staticmethod
744+
def test_constructor_sortable_with_projection():
745+
foo = model.StringProperty("foo")
746+
order_by = [query_module.PropertyOrder("foo")]
747+
query = query_module.QueryOptions(
748+
filters=query_module.OR(foo == "this", foo == "that"),
749+
order_by=order_by,
750+
projection=["foo"],
751+
)
752+
iterator = _datastore_query._MultiQueryIteratorImpl(query)
753+
assert iterator._result_sets[0]._query == query_module.QueryOptions(
754+
filters=foo == "this", order_by=order_by, projection=["foo"],
755+
)
756+
assert iterator._result_sets[1]._query == query_module.QueryOptions(
757+
filters=foo == "that", order_by=order_by, projection=["foo"],
758+
)
759+
assert iterator._sortable
760+
761+
@staticmethod
762+
def test_constructor_sortable_with_projection_needs_extra():
763+
foo = model.StringProperty("foo")
764+
order_by = [query_module.PropertyOrder("foo")]
765+
query = query_module.QueryOptions(
766+
filters=query_module.OR(foo == "this", foo == "that"),
767+
order_by=order_by,
768+
projection=["bar"],
769+
)
770+
iterator = _datastore_query._MultiQueryIteratorImpl(query)
771+
assert iterator._result_sets[0]._query == query_module.QueryOptions(
772+
filters=foo == "this",
773+
order_by=order_by,
774+
projection=["bar", "foo"],
775+
)
776+
assert iterator._result_sets[1]._query == query_module.QueryOptions(
777+
filters=foo == "that",
778+
order_by=order_by,
779+
projection=["bar", "foo"],
780+
)
781+
assert iterator._sortable
782+
743783
@staticmethod
744784
def test_iter():
745785
foo = model.StringProperty("foo")
@@ -783,6 +823,34 @@ def test_has_next_async_exhausted():
783823
iterator._result_sets = []
784824
assert not iterator.has_next_async().result()
785825

826+
@staticmethod
827+
@pytest.mark.usefixtures("in_context")
828+
def test_next_with_extra_projections():
829+
foo = model.StringProperty("foo")
830+
order_by = [
831+
query_module.PropertyOrder("foo"),
832+
query_module.PropertyOrder("food"),
833+
]
834+
query = query_module.QueryOptions(
835+
filters=query_module.OR(foo == "this", foo == "that"),
836+
order_by=order_by,
837+
projection=["bar"],
838+
)
839+
iterator = _datastore_query._MultiQueryIteratorImpl(query)
840+
iterator._next_result = next_result = mock.Mock(
841+
result_pb=mock.Mock(
842+
entity=mock.Mock(
843+
properties={"foo": 1, "bar": "two"}, spec=("properties",),
844+
),
845+
spec=("entity",),
846+
),
847+
spec=("result_pb",),
848+
)
849+
iterator._raw = True
850+
851+
assert iterator.next() is next_result
852+
assert "foo" not in next_result.result_pb.entity.properties
853+
786854
@staticmethod
787855
@pytest.mark.usefixtures("in_context")
788856
def test_iterate_async():

0 commit comments

Comments
 (0)