NDB: Unable to retrieve heirs of two different PolyModel bases with the same name from a query.

[ventice11o]

Environment:
OS: Microsoft Windows 11 Pro 10.0.22631 Build 22631, Ubuntu Linux
Python 3.12.1
google-cloud-ndb==2.3.0

Steps to reproduce:
Run this code:

from google.cloud.ndb.polymodel import PolyModel
class A(PolyModel): pass
class B(PolyModel): pass

def inherit(base):
  class Same(base): pass
  return Same

client = ndb.Client(namespace='test')
with client.context():
  ndb.put_multi([inherit(A)(), inherit(B)()])
  A.query().fetch()
  B.query().fetch()

Stack trace

---------------------------------------------------------------------------
KindError                                 Traceback (most recent call last)
Cell In[5], line 13
     11 with client.context():
     12   ndb.put_multi([inherit(A)(), inherit(B)()])
---> 13   A.query().fetch()
     14   B.query().fetch()

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\query.py:1202, in _query_options.<locals>.wrapper(self, *args, **kwargs)
   1199 context = context_module.get_context()
   1200 query_options = QueryOptions(context=context, **query_arguments)
-> 1202 return wrapped(self, *dummy_args, _options=query_options)

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\utils.py:118, in keyword_only.__call__.<locals>.wrapper(*args, **kwargs)
    113         raise TypeError(
    114             "%s() got an unexpected keyword argument '%s'"
    115             % (wrapped.__name__, kwarg)
    116         )
    117 new_kwargs.update(kwargs)
--> 118 return wrapped(*args, **new_kwargs)

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\utils.py:150, in positional.<locals>.positional_decorator.<locals>.positional_wrapper(*args, **kwds)
    145         plural_s = "s"
    146     raise TypeError(
    147         "%s() takes at most %d positional argument%s (%d given)"
    148         % (wrapped.__name__, max_pos_args, plural_s, len(args))
    149     )
--> 150 return wrapped(*args, **kwds)

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\query.py:1744, in Query.fetch(self, limit, **kwargs)
   1695 @_query_options
   1696 @utils.keyword_only(
   1697     keys_only=None,
   (...)
   1713 @utils.positional(2)
   1714 def fetch(self, limit=None, **kwargs):
   1715     """Run a query, fetching results.
   1716
   1717     Args:
   (...)
   1742         List[Union[model.Model, key.Key]]: The query results.
   1743     """
-> 1744     return self.fetch_async(_options=kwargs["_options"]).result()

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\tasklets.py:210, in Future.result(self)
    201 def result(self):
    202     """Return the result of this future's task.
    203
    204     If the task is finished, this will return immediately. Otherwise, this
   (...)
    208         Any: The result
    209     """
--> 210     self.check_success()
    211     return self._result

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\tasklets.py:157, in Future.check_success(self)
    154 self.wait()
    156 if self._exception:
--> 157     raise self._exception

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\tasklets.py:323, in _TaskletFuture._advance_tasklet(***failed resolving arguments***)
    319     yielded = self.generator.throw(type(error), error, traceback)
    321 else:
    322     # send_value will be None if this is the first time
--> 323     yielded = self.generator.send(send_value)
    325 # Context may have changed in tasklet
    326 self.context = context_module.get_context()

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\_datastore_query.py:117, in fetch(query)
    115 entities = []
    116 while (yield results.has_next_async()):
--> 117     entities.append(results.next())
    119 raise tasklets.Return(entities)

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\_datastore_query.py:433, in _QueryIteratorImpl.next(self)
    430 self._cursor_after = next_result.cursor
    432 if not self._raw:
--> 433     next_result = next_result.entity()
    435 return next_result

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\_datastore_query.py:889, in _Result.entity(self)
    886 entity = self.check_cache(context)
    887 if entity is _KEY_NOT_IN_CACHE:
    888     # entity not in cache, create one, and then add it to cache
--> 889     entity = model._entity_from_protobuf(self.result_pb.entity)
    890     if context._use_cache(entity.key, self._query_options):
    891         context.cache[entity.key] = entity

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\model.py:702, in _entity_from_protobuf(protobuf)
    692 """Deserialize an entity from a protobuffer.
    693
    694 Args:
   (...)
    699     .Model: The deserialized entity.
    700 """
    701 ds_entity = helpers.entity_from_protobuf(protobuf)
--> 702 return _entity_from_ds_entity(ds_entity)

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\model.py:565, in _entity_from_ds_entity(ds_entity, model_class)
    562 else:
    563     kind = ds_entity.kind
--> 565 model_class = model_class or Model._lookup_model(kind)
    566 entity = model_class()
    568 if ds_entity.key:

File C:\work\GoogleApps\opttorg-api\temp\venv312\Lib\site-packages\google\cloud\ndb\model.py:5256, in Model._lookup_model(cls, kind, default_model)
   5254 model_class = cls._kind_map.get(kind, default_model)
   5255 if model_class is None:
-> 5256     raise KindError(
   5257         (
   5258             "No model class found for the kind '{}'. Did you forget "
   5259             "to import it?"
   5260         ).format(kind)
   5261     )
   5262 return model_class

KindError: No model class found for the kind 'Heir'. Did you forget to import it?

The reason:
The reason is that the function google.cloud.ndb.model._entity_from_ds_entity ignores the full class_key from ds_entity and reduces it to its final kind, so, two heirs having the same name overwrite each other in the _kind_map and one of them cannot be found. The solution would be to record all the polymodels into separate registry and look them up by class.

Backward compatibility
This used to work in google.appengine.ext.ndb, this is how these entities have been written to the datastore. Current code blocks their retrieval.

Workaround
This code can be worked around by monkeypatching the google.cloud.ndb.model._entity_from_ds_entity in the following way:

from google.cloud import ndb
import model # the package containing all the model classes. as an alternative, a metaclass can be used to collect them

def _get_polymodels():
    import pkgutil, os
    for info in pkgutil.iter_modules([os.path.dirname(model.__file__)]):
        for member in model.__dict__[info.name].__dict__.values():
            if isinstance(member, ndb.model.MetaModel):
                obj = member()
                if hasattr(obj, 'class_'):
                    yield tuple(obj.class_), member

_POLY_MODELS = dict(_get_polymodels())

def _patch_ds_entity_converter():
    _old_efde = ndb.model._entity_from_ds_entity

    def _entity_from_ds_entity(ds_entity, model_class=None):
        ds_class = tuple(ds_entity.get("class") or [])
        return _old_efde(ds_entity, _POLY_MODELS.get(ds_class, model_class))

    ndb.model._entity_from_ds_entity = _entity_from_ds_entity

_patch_ds_entity_converter()