Skip to content

Commit a201894

Browse files
PrimevennPrimevenn
authored
fix: support ordering by key for multi queries (#630)
Fixes: #629
1 parent 49c1999 commit a201894

3 files changed

Lines changed: 156 additions & 5 deletions

File tree

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

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from google.cloud.datastore_v1.proto import datastore_pb2
2626
from google.cloud.datastore_v1.proto import entity_pb2
2727
from google.cloud.datastore_v1.proto import query_pb2
28-
from google.cloud.datastore import helpers
28+
from google.cloud.datastore import helpers, Key
2929

3030
from google.cloud.ndb import context as context_module
3131
from google.cloud.ndb import _datastore_api
@@ -801,10 +801,26 @@ def _compare(self, other):
801801
return NotImplemented
802802

803803
for order in self.order_by:
804-
this_value_pb = self.result_pb.entity.properties[order.name]
805-
this_value = helpers._get_value_from_value_pb(this_value_pb)
806-
other_value_pb = other.result_pb.entity.properties[order.name]
807-
other_value = helpers._get_value_from_value_pb(other_value_pb)
804+
805+
if order.name == "__key__":
806+
this_value = helpers.key_from_protobuf(
807+
self.result_pb.entity.key
808+
).flat_path
809+
other_value = helpers.key_from_protobuf(
810+
other.result_pb.entity.key
811+
).flat_path
812+
else:
813+
this_value_pb = self.result_pb.entity.properties[order.name]
814+
this_value = helpers._get_value_from_value_pb(this_value_pb)
815+
other_value_pb = other.result_pb.entity.properties[order.name]
816+
other_value = helpers._get_value_from_value_pb(other_value_pb)
817+
818+
# Compare key paths if ordering by key property
819+
if isinstance(this_value, Key):
820+
this_value = this_value.flat_path
821+
822+
if isinstance(other_value, Key):
823+
other_value = other_value.flat_path
808824

809825
direction = -1 if order.reverse else 1
810826

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -750,6 +750,64 @@ class SomeKind(ndb.Model):
750750
results[0].foo
751751

752752

753+
@pytest.mark.usefixtures("client_context")
754+
def test_multiquery_with_order_by_entity_key(ds_entity):
755+
"""Regression test for #629
756+
757+
https://github.com/googleapis/python-ndb/issues/629
758+
"""
759+
760+
for i in range(5):
761+
entity_id = test_utils.system.unique_resource_id()
762+
ds_entity(KIND, entity_id, foo=i)
763+
764+
class SomeKind(ndb.Model):
765+
foo = ndb.IntegerProperty()
766+
767+
query = (
768+
SomeKind.query()
769+
.order(SomeKind.key)
770+
.filter(ndb.OR(SomeKind.foo == 4, SomeKind.foo == 3, SomeKind.foo == 1))
771+
)
772+
773+
results = eventually(query.fetch, length_equals(3))
774+
assert [entity.foo for entity in results] == [1, 3, 4]
775+
776+
777+
@pytest.mark.usefixtures("client_context")
778+
def test_multiquery_with_order_key_property(ds_entity, client_context):
779+
"""Regression test for #629
780+
781+
https://github.com/googleapis/python-ndb/issues/629
782+
"""
783+
project = client_context.client.project
784+
namespace = client_context.get_namespace()
785+
786+
for i in range(5):
787+
entity_id = test_utils.system.unique_resource_id()
788+
ds_entity(
789+
KIND,
790+
entity_id,
791+
foo=i,
792+
bar=ds_key_module.Key(
793+
"test_key", i + 1, project=project, namespace=namespace
794+
),
795+
)
796+
797+
class SomeKind(ndb.Model):
798+
foo = ndb.IntegerProperty()
799+
bar = ndb.KeyProperty()
800+
801+
query = (
802+
SomeKind.query()
803+
.order(SomeKind.bar)
804+
.filter(ndb.OR(SomeKind.foo == 4, SomeKind.foo == 3, SomeKind.foo == 1))
805+
)
806+
807+
results = eventually(query.fetch, length_equals(3))
808+
assert [entity.foo for entity in results] == [1, 3, 4]
809+
810+
753811
@pytest.mark.usefixtures("client_context")
754812
def test_count_with_multi_query(ds_entity):
755813
for i in range(5):

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1550,6 +1550,83 @@ def test__compare_with_order_by():
15501550
)
15511551
assert result._compare("other") == NotImplemented
15521552

1553+
@staticmethod
1554+
def test__compare_with_order_by_entity_key():
1555+
def result(key_path):
1556+
key_pb = entity_pb2.Key(
1557+
partition_id=entity_pb2.PartitionId(project_id="testing"),
1558+
path=[key_path],
1559+
)
1560+
return _datastore_query._Result(
1561+
result_type=None,
1562+
result_pb=query_pb2.EntityResult(entity=entity_pb2.Entity(key=key_pb)),
1563+
order_by=[
1564+
query_module.PropertyOrder("__key__"),
1565+
],
1566+
)
1567+
1568+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) < result(
1569+
entity_pb2.Key.PathElement(kind="ThisKind", name="b")
1570+
)
1571+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="b")) > result(
1572+
entity_pb2.Key.PathElement(kind="ThisKind", name="a")
1573+
)
1574+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) != result(
1575+
entity_pb2.Key.PathElement(kind="ThisKind", name="b")
1576+
)
1577+
1578+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) < result(
1579+
entity_pb2.Key.PathElement(kind="ThisKind", id=2)
1580+
)
1581+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=2)) > result(
1582+
entity_pb2.Key.PathElement(kind="ThisKind", id=1)
1583+
)
1584+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) != result(
1585+
entity_pb2.Key.PathElement(kind="ThisKind", id=2)
1586+
)
1587+
1588+
@staticmethod
1589+
def test__compare_with_order_by_key_property():
1590+
def result(foo_key_path):
1591+
foo_key = entity_pb2.Key(
1592+
partition_id=entity_pb2.PartitionId(project_id="testing"),
1593+
path=[foo_key_path],
1594+
)
1595+
1596+
return _datastore_query._Result(
1597+
result_type=None,
1598+
result_pb=query_pb2.EntityResult(
1599+
entity=entity_pb2.Entity(
1600+
properties={
1601+
"foo": entity_pb2.Value(key_value=foo_key),
1602+
}
1603+
)
1604+
),
1605+
order_by=[
1606+
query_module.PropertyOrder("foo"),
1607+
],
1608+
)
1609+
1610+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) < result(
1611+
entity_pb2.Key.PathElement(kind="ThisKind", name="b")
1612+
)
1613+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="b")) > result(
1614+
entity_pb2.Key.PathElement(kind="ThisKind", name="a")
1615+
)
1616+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) != result(
1617+
entity_pb2.Key.PathElement(kind="ThisKind", name="b")
1618+
)
1619+
1620+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) < result(
1621+
entity_pb2.Key.PathElement(kind="ThisKind", id=2)
1622+
)
1623+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=2)) > result(
1624+
entity_pb2.Key.PathElement(kind="ThisKind", id=1)
1625+
)
1626+
assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) != result(
1627+
entity_pb2.Key.PathElement(kind="ThisKind", id=2)
1628+
)
1629+
15531630
@staticmethod
15541631
@mock.patch("google.cloud.ndb._datastore_query.model")
15551632
def test_entity_unsupported_result_type(model):

0 commit comments

Comments
 (0)