Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: support polymodels with equal names but different roots (#961) #986

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion google/cloud/ndb/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -557,7 +557,11 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
if not isinstance(class_key, list):
kind = class_key
else:
kind = class_key[-1]
poly_model_class = Model._class_map.get(tuple(class_key))
if not model_class and poly_model_class:
model_class = poly_model_class
else:
kind = class_key[-1]
else:
kind = ds_entity.kind

Expand Down Expand Up @@ -4877,6 +4881,7 @@ class MyModel(ndb.Model):
_properties = None
_has_repeated = False
_kind_map = {} # Dict mapping {kind: Model subclass}
_class_map = {} # Map class key -> suitable subclass (populated by PolyModel)

# Defaults for instance variables.
_entity_key = None
Expand Down
5 changes: 1 addition & 4 deletions google/cloud/ndb/polymodel.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,17 +184,14 @@ class PolyModel(model.Model):

class_ = _ClassKeyProperty()

_class_map = {} # Map class key -> suitable subclass.

@classmethod
def _update_kind_map(cls):
"""Override; called by Model._fix_up_properties().

Update the kind map as well as the class map, except for PolyModel
itself (its class key is empty). Note that the kind map will
contain entries for all classes in a PolyModel hierarchy; they all
have the same kind, but different class names. PolyModel class
names, like regular Model class names, must be globally unique.
have the same kind, but different class names.
"""
cls._kind_map[cls._class_name()] = cls
Copy link
Contributor

@rwhogg rwhogg Jul 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the reason why class names have to be globally unique, unfortunately. Different model classes will overwrite each other here if they have the same class name.

You can see this by adding the following to the end of your test below:

assert isinstance(CatCommand._kind_map["Cat"].bar, model.StringProperty)
assert not isinstance(CatAnimal._kind_map["Cat"].bar, model.StringProperty)

Notice how on the second line, even though we're accessing it through CatAnimal, the model referenced by CatAnimal._kind_map["Cat"] is for a CatCommand.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, they are overwritten, but there's a code above it that uses its own dictionary using fully qualified names to look up for the real entity during model instance creation. So, yes, non-unique polymorphic classes, when accessed using _kind_map will be buggy, but why would a polymodel be looked up via global registry of top level classes? Polymodels provide their own registry and don't care what other registries think of them.
At least my legacy code with my fix doesn't experience any problems in my scenario as it did not in legacy ndb. And I use polymorphic models actively.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like I can override _class_name for PolyModel to return something like Animal.Cat and Command.Cat. In this case even _kind_map will keep its invariant.

class_key = cls._class_key()
Expand Down
41 changes: 41 additions & 0 deletions tests/unit/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -5798,6 +5798,47 @@ class Cat(Animal):
assert entity.bar == "himom!"
assert entity.class_ == ["Cat"]

@staticmethod
@pytest.mark.usefixtures("in_context")
def test_polymodel_equal_names_different_bases():
def inherit_cat(base):
class Cat(base):
data = model.StringProperty()
return Cat

class Animal(polymodel.PolyModel):
foo = model.IntegerProperty()

class Command(polymodel.PolyModel):
bar = model.StringProperty()

CatAnimal = inherit_cat(Animal)
CatCommand = inherit_cat(Command)

animal_key = datastore.Key("Animal", 123, project="testing")
datastore_entity = datastore.Entity(key=animal_key)
datastore_entity.update(
{"foo": 42, "data": "meow!", "class": ["Animal", "Cat"]}
)

entity = model._entity_from_ds_entity(datastore_entity)
assert isinstance(entity, CatAnimal)
assert entity.foo == 42
assert entity.data == "meow!"
assert entity.class_ == ["Animal", "Cat"]

command_key = datastore.Key("Command", 456, project="testing")
datastore_entity = datastore.Entity(key=command_key)
datastore_entity.update(
{"bar": "42", "data": "config.conf", "class": ["Command", "Cat"]}
)

entity = model._entity_from_ds_entity(datastore_entity)
assert isinstance(entity, CatCommand)
assert entity.bar == "42"
assert entity.data == "config.conf"
assert entity.class_ == ["Command", "Cat"]


class Test_entity_to_protobuf:
@staticmethod
Expand Down