fix: support ordering by key for multi queries

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,26 @@ 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
+                ).flat_path
+                other_value = helpers.key_from_protobuf(
+                    other.result_pb.entity.key
+                ).flat_path
+            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 property
+                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,64 @@ class SomeKind(ndb.Model):
         results[0].foo
 
 
+@pytest.mark.usefixtures("client_context")
+def test_multiquery_with_order_by_entity_key(ds_entity):
+    """Regression test for #629
+
+    https://github.com/googleapis/python-ndb/issues/629
+    """
+
+    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_key_property(ds_entity, client_context):
+    """Regression test for #629
+
+    https://github.com/googleapis/python-ndb/issues/629
+    """
+    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):

tests/unit/test__datastore_query.py

@@ -1550,6 +1550,83 @@ def test__compare_with_order_by():
         )
         assert result._compare("other") == NotImplemented
 
+    @staticmethod
+    def test__compare_with_order_by_entity_key():
+        def result(key_path):
+            key_pb = entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[key_path],
+            )
+            return _datastore_query._Result(
+                result_type=None,
+                result_pb=query_pb2.EntityResult(entity=entity_pb2.Entity(key=key_pb)),
+                order_by=[
+                    query_module.PropertyOrder("__key__"),
+                ],
+            )
+
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) < result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="b")
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="b")) > result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="a")
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) != result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="b")
+        )
+
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) < result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=2)
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=2)) > result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=1)
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) != result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=2)
+        )
+
+    @staticmethod
+    def test__compare_with_order_by_key_property():
+        def result(foo_key_path):
+            foo_key = entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[foo_key_path],
+            )
+
+            return _datastore_query._Result(
+                result_type=None,
+                result_pb=query_pb2.EntityResult(
+                    entity=entity_pb2.Entity(
+                        properties={
+                            "foo": entity_pb2.Value(key_value=foo_key),
+                        }
+                    )
+                ),
+                order_by=[
+                    query_module.PropertyOrder("foo"),
+                ],
+            )
+
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) < result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="b")
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="b")) > result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="a")
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", name="a")) != result(
+            entity_pb2.Key.PathElement(kind="ThisKind", name="b")
+        )
+
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) < result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=2)
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=2)) > result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=1)
+        )
+        assert result(entity_pb2.Key.PathElement(kind="ThisKind", id=1)) != result(
+            entity_pb2.Key.PathElement(kind="ThisKind", id=2)
+        )
+
     @staticmethod
     @mock.patch("google.cloud.ndb._datastore_query.model")
     def test_entity_unsupported_result_type(model):