fix: support polymodel in local structured property (#497)

cguardia · web-flow · commit 8516478f19a9 · 2020-07-31T14:22:01.000-05:00
refs #481
diff --git a/packages/google-cloud-ndb/google/cloud/ndb/model.py b/packages/google-cloud-ndb/google/cloud/ndb/model.py
@@ -636,6 +636,12 @@ def new_entity(key):
 
                 continue
 
+        if prop is None and kind != model_class.__name__:
+            # kind and model_class name do not match, so this is probably a
+            # polymodel. We need to check if the prop belongs to the subclass.
+            model_subclass = Model._lookup_model(kind)
+            prop = getattr(model_subclass, name, None)
+
         def base_value_or_none(value):
             return None if value is None else _BaseValue(value)
 
@@ -4357,7 +4363,14 @@ def _from_base_type(self, value):
             value = entity_value
         if not self._keep_keys and value.key:
             value.key = None
-        return _entity_from_ds_entity(value, model_class=self._model_class)
+        model_class = self._model_class
+        kind = self._model_class.__name__
+        if "class" in value and value["class"]:
+            kind = value["class"][-1] or model_class
+        if kind != self._model_class.__name__:
+            # if this is a polymodel, find correct subclass.
+            model_class = Model._lookup_model(kind)
+        return _entity_from_ds_entity(value, model_class=model_class)
 
     def _prepare_for_put(self, entity):
         values = self._get_user_value(entity)
diff --git a/packages/google-cloud-ndb/tests/system/test_crud.py b/packages/google-cloud-ndb/tests/system/test_crud.py
@@ -1473,3 +1473,27 @@ class SomeKind(ndb.Model):
     ourkind.bar = "confusing"
 
     assert somekind.bar is None
+
+
+@pytest.mark.usefixtures("client_context")
+def test_local_structured_property_with_polymodel(dispose_of):
+    """Regression test for #481
+
+    https://github.com/googleapis/python-ndb/issues/481
+    """
+
+    class Base(ndb.PolyModel):
+        pass
+
+    class SubKind(Base):
+        foo = ndb.StringProperty()
+
+    class Container(ndb.Model):
+        child = ndb.LocalStructuredProperty(Base)
+
+    entity = Container(child=SubKind(foo="bar"))
+    key = entity.put()
+    dispose_of(key._key)
+
+    entity = entity.key.get()
+    assert entity.child.foo == "bar"
diff --git a/packages/google-cloud-ndb/tests/unit/test_model.py b/packages/google-cloud-ndb/tests/unit/test_model.py
@@ -3863,6 +3863,34 @@ class ContainerA(model.Model):
             assert data.pop("_exclude_from_indexes") == ["child_a"]
             assert data["child_a"]["child_b"] is None
 
+    @staticmethod
+    def test_local_structured_property_with_polymodel(in_context):
+        class Base(polymodel.PolyModel):
+            pass
+
+        class SubKind(Base):
+            foo = model.StringProperty()
+
+        class Container(model.Model):
+            child = model.LocalStructuredProperty(Base)
+
+        entity = Container(child=SubKind(foo="bar"))
+        value = b"".join(
+            [
+                b"\x1a \n\x05class\x12\x17J\x15\n\x07\x8a\x01\x04Base\n\n",
+                b"\x8a\x01\x07SubKind\x1a\r\n\x03foo\x12\x06\x8a\x01\x03bar",
+            ]
+        )
+
+        child = entity._properties["child"]._from_base_type(value)
+        assert child.foo == "bar"
+
+        pb = entity_pb2.Entity()
+        pb.MergeFromString(value)
+        value = helpers.entity_from_protobuf(pb)
+        child = model._entity_from_ds_entity(value, model_class=Base)
+        assert child._values["foo"].b_val == "bar"
+
 
 class TestGenericProperty:
     @staticmethod
