fix - support ordering by key for multi queries

Author: Primevenn

google/cloud/ndb/_datastore_query.py

@@ -25,7 +25,7 @@
 from google.cloud.datastore_v1.proto import datastore_pb2
 from google.cloud.datastore_v1.proto import entity_pb2
 from google.cloud.datastore_v1.proto import query_pb2
-from google.cloud.datastore import helpers
+from google.cloud.datastore import helpers, Key
 
 from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _datastore_api
@@ -801,10 +801,22 @@ def _compare(self, other):
             return NotImplemented
 
         for order in self.order_by:
-            this_value_pb = self.result_pb.entity.properties[order.name]
-            this_value = helpers._get_value_from_value_pb(this_value_pb)
-            other_value_pb = other.result_pb.entity.properties[order.name]
-            other_value = helpers._get_value_from_value_pb(other_value_pb)
+
+            if order.name == '__key__':
+                this_value = helpers.key_from_protobuf(self.result_pb.entity.key)
+                other_value = helpers.key_from_protobuf(other.result_pb.entity.key)
+            else:
+                this_value_pb = self.result_pb.entity.properties[order.name]
+                this_value = helpers._get_value_from_value_pb(this_value_pb)
+                other_value_pb = other.result_pb.entity.properties[order.name]
+                other_value = helpers._get_value_from_value_pb(other_value_pb)
+
+            # Compare key paths if ordering by key
+            if isinstance(this_value, Key):
+                this_value = this_value.flat_path
+
+            if isinstance(other_value, Key):
+                other_value = other_value.flat_path
 
             direction = -1 if order.reverse else 1
 

tests/system/test_query.py

@@ -750,6 +750,49 @@ class SomeKind(ndb.Model):
         results[0].foo
 
 
+@pytest.mark.usefixtures("client_context")
+def test_multiquery_with_order_by_key(ds_entity):
+    """Verify multiqueries work when ordering by key"""
+
+    for i in range(5):
+        entity_id = test_utils.system.unique_resource_id()
+        ds_entity(KIND, entity_id, foo=i)
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    query = (
+        SomeKind.query()
+        .order(SomeKind.key)
+        .filter(ndb.OR(SomeKind.foo == 4, SomeKind.foo == 3, SomeKind.foo == 1))
+    )
+
+    results = eventually(query.fetch, length_equals(3))
+    assert [entity.foo for entity in results] == [1, 3, 4]
+
+@pytest.mark.usefixtures("client_context")
+def test_multiquery_with_order_value_by_key(ds_entity, client_context):
+    """Check ordering by key values an Entity references"""
+    project = client_context.client.project
+    namespace = client_context.get_namespace()
+
+    for i in range(5):
+        entity_id = test_utils.system.unique_resource_id()
+        ds_entity(KIND, entity_id, foo=i, bar=ds_key_module.Key("test_key", i + 1, project=project, namespace=namespace))
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.KeyProperty()
+
+    query = (
+        SomeKind.query()
+            .order(SomeKind.bar)
+            .filter(ndb.OR(SomeKind.foo == 4, SomeKind.foo == 3, SomeKind.foo == 1))
+    )
+
+    results = eventually(query.fetch, length_equals(3))
+    assert [entity.foo for entity in results] == [1, 3, 4]
+
 @pytest.mark.usefixtures("client_context")
 def test_count_with_multi_query(ds_entity):
     for i in range(5):