Fix handling of repeated properties with projection queries. (#257)

This fixes issues, generally, with the handling of repeated properties
with projection queries, but also specifically for the "class" property
of PolyModel entities.

Fixes #248

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/model.py

@@ -533,20 +533,20 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
     """
     class_key = ds_entity.get("class")
     if class_key:
-        kind = class_key[-1]
+        # If this is a projection query, we'll get multiple entities with
+        # scalar values rather than single entities with array values.
+        # It's weird:
+        #   https://cloud.google.com/datastore/docs/concepts/queries#datastore-datastore-array-value-python
+        if not isinstance(class_key, list):
+            kind = class_key
+        else:
+            kind = class_key[-1]
     else:
         kind = ds_entity.kind
 
     model_class = model_class or Model._lookup_model(kind)
     entity = model_class()
 
-    # Check if we are dealing with a PolyModel, and if so get correct subclass.
-    # We need to import here to avoid circular import.
-    from google.cloud.ndb import PolyModel
-
-    if isinstance(entity, PolyModel) and "class" in ds_entity:
-        entity = entity._class_map[tuple(ds_entity["class"])]()
-
     if ds_entity.key:
         entity._key = key_module.Key._from_ds_key(ds_entity.key)
 
@@ -640,10 +640,18 @@ def new_entity(key):
 
         if value is not None:
             if prop._repeated:
-                value = [
-                    (_BaseValue(sub_value) if sub_value else None)
-                    for sub_value in value
-                ]
+                # A repeated property will have a scalar value if this is a
+                # projection query.
+                if isinstance(value, list):
+                    # Not a projection
+                    value = [
+                        (_BaseValue(sub_value) if sub_value else None)
+                        for sub_value in value
+                    ]
+                else:
+                    # Projection
+                    value = [_BaseValue(value)]
+
             else:
                 value = _BaseValue(value)
 

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

@@ -21,3 +21,8 @@ indexes:
   - name: foo
   - name: bar.one
   - name: bar.two
+
+- kind: Animal
+  properties:
+  - name: class
+  - name: foo

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

@@ -607,6 +607,46 @@ class Cat(Animal):
     assert isinstance(results[0], Cat)
 
 
+@pytest.mark.usefixtures("client_context")
+def test_polymodel_query_class_projection(ds_entity):
+    """Regression test for Issue #248
+
+    https://github.com/googleapis/python-ndb/issues/248
+    """
+
+    class Animal(ndb.PolyModel):
+        foo = ndb.IntegerProperty()
+
+    class Cat(Animal):
+        pass
+
+    animal = Animal(foo=1)
+    animal.put()
+    cat = Cat(foo=2)
+    cat.put()
+
+    query = Animal.query(projection=["class", "foo"])
+    results = eventually(query.fetch, _length_equals(3))
+
+    # Mostly reproduces odd behavior of legacy code
+    results = sorted(results, key=operator.attrgetter("foo"))
+
+    assert isinstance(results[0], Animal)
+    assert not isinstance(results[0], Cat)
+    assert results[0].foo == 1
+    assert results[0].class_ == ["Animal"]
+
+    assert isinstance(results[1], Animal)
+    assert not isinstance(results[1], Cat)
+    assert results[1].foo == 2
+    assert results[1].class_ == ["Animal"]
+
+    assert isinstance(results[2], Animal)
+    assert isinstance(results[2], Cat)  # This would be False in legacy
+    assert results[2].foo == 2
+    assert results[2].class_ == ["Cat"]
+
+
 @pytest.mark.usefixtures("client_context")
 def test_query_repeated_property(ds_entity):
     entity_id = test_utils.system.unique_resource_id()

packages/google-cloud-ndb/tests/unit/test_model.py

@@ -37,6 +37,7 @@
 from google.cloud.ndb import key as key_module
 from google.cloud.ndb import model
 from google.cloud.ndb import _options
+from google.cloud.ndb import polymodel
 from google.cloud.ndb import query as query_module
 from google.cloud.ndb import tasklets
 from google.cloud.ndb import utils as ndb_utils
@@ -5091,6 +5092,46 @@ class ThisKind(model.Model):
         assert entity.baz[2].bar == "iminjail"
         assert entity.copacetic is True
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_polymodel():
+        class Animal(polymodel.PolyModel):
+            foo = model.IntegerProperty()
+
+        class Cat(Animal):
+            bar = model.StringProperty()
+
+        key = datastore.Key("Animal", 123, project="testing")
+        datastore_entity = datastore.Entity(key=key)
+        datastore_entity.update(
+            {"foo": 42, "bar": "himom!", "class": ["Animal", "Cat"]}
+        )
+
+        entity = model._entity_from_ds_entity(datastore_entity)
+        assert isinstance(entity, Cat)
+        assert entity.foo == 42
+        assert entity.bar == "himom!"
+        assert entity.class_ == ["Animal", "Cat"]
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_polymodel_projection():
+        class Animal(polymodel.PolyModel):
+            foo = model.IntegerProperty()
+
+        class Cat(Animal):
+            bar = model.StringProperty()
+
+        key = datastore.Key("Animal", 123, project="testing")
+        datastore_entity = datastore.Entity(key=key)
+        datastore_entity.update({"foo": 42, "bar": "himom!", "class": "Cat"})
+
+        entity = model._entity_from_ds_entity(datastore_entity)
+        assert isinstance(entity, Cat)
+        assert entity.foo == 42
+        assert entity.bar == "himom!"
+        assert entity.class_ == ["Cat"]
+
 
 class Test_entity_to_protobuf:
     @staticmethod