INTPYTHON-624 Add PolymorphicEmbeddedModelField

Author: timgraham

django_mongodb_backend/fields/__init__.py

@@ -5,6 +5,7 @@
 from .embedded_model_array import EmbeddedModelArrayField
 from .json import register_json_field
 from .objectid import ObjectIdField
+from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
 
 __all__ = [
     "register_fields",
@@ -13,6 +14,7 @@
     "EmbeddedModelField",
     "ObjectIdAutoField",
     "ObjectIdField",
+    "PolymorphicEmbeddedModelField",
 ]
 
 

django_mongodb_backend/fields/polymorphic_embedded_model.py

@@ -0,0 +1,221 @@
+import contextlib
+import difflib
+
+from django.core import checks
+from django.core.exceptions import FieldDoesNotExist, ValidationError
+from django.db import models
+from django.db.models.fields.related import lazy_related_operation
+from django.db.models.lookups import Transform
+
+
+class PolymorphicEmbeddedModelField(models.Field):
+    """Field that stores a model instance."""
+
+    value_is_model_instance = True
+
+    def __init__(self, embedded_models, *args, **kwargs):
+        """
+        `embedded_models` is a list of possible model classes to be stored.
+        Like other relational fields, each model may also be passed as a
+        string.
+        """
+        self.embedded_models = embedded_models
+        kwargs["editable"] = False
+        super().__init__(*args, **kwargs)
+
+    def db_type(self, connection):
+        return "embeddedDocuments"
+
+    def check(self, **kwargs):
+        from ..models import EmbeddedModel
+
+        errors = super().check(**kwargs)
+        for model in self.embedded_models:
+            if not issubclass(model, EmbeddedModel):
+                return [
+                    checks.Error(
+                        "Embedded models must be a subclass of "
+                        "django_mongodb_backend.models.EmbeddedModel.",
+                        obj=self,
+                        hint="{model} doesn't subclass EmbeddedModel.",
+                        id="django_mongodb_backend.embedded_model.E002",
+                    )
+                ]
+            for field in model._meta.fields:
+                if field.remote_field:
+                    errors.append(
+                        checks.Error(
+                            "Embedded models cannot have relational fields "
+                            f"({model().__class__.__name__}.{field.name} "
+                            f"is a {field.__class__.__name__}).",
+                            obj=self,
+                            id="django_mongodb_backend.embedded_model.E001",
+                        )
+                    )
+        return errors
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if path.startswith("django_mongodb_backend.fields.polymorphic_embedded_model"):
+            path = path.replace(
+                "django_mongodb_backend.fields.polymorphic_embedded_model",
+                "django_mongodb_backend.fields",
+            )
+        kwargs["embedded_models"] = self.embedded_models
+        del kwargs["editable"]
+        return name, path, args, kwargs
+
+    def get_internal_type(self):
+        return "PolymorphicEmbeddedModelField"
+
+    def _set_model(self, model):
+        """
+        Resolve embedded model classes once the field knows the model it
+        belongs to. If any of the items in __init__()'s embedded_models
+        argument are strings, resolve each to the actual model class,
+        similar to relation fields.
+        """
+        self._model = model
+        if model is not None:
+            for embedded_model in self.embedded_models:
+                if isinstance(embedded_model, str):
+
+                    def _resolve_lookup(_, *resolved_models):
+                        self.embedded_models = resolved_models
+
+                    lazy_related_operation(_resolve_lookup, model, *self.embedded_models)
+
+    model = property(lambda self: self._model, _set_model)
+
+    def from_db_value(self, value, expression, connection):
+        return self.to_python(value)
+
+    def to_python(self, value):
+        """
+        Pass embedded model fields' values through each field's to_python() and
+        reinstantiate the embedded instance.
+        """
+        if value is None:
+            return None
+        if not isinstance(value, dict):
+            return value
+        model_class = self._get_model_from_label(value.pop("_label"))
+        instance = model_class(
+            **{
+                field.attname: field.to_python(value[field.attname])
+                for field in model_class._meta.fields
+                if field.attname in value
+            }
+        )
+        instance._state.adding = False
+        return instance
+
+    def get_db_prep_save(self, embedded_instance, connection):
+        """
+        Apply pre_save() and get_db_prep_save() of embedded instance fields and
+        create the {field: value} dict to be saved.
+        """
+        if embedded_instance is None:
+            return None
+        if not isinstance(embedded_instance, self.embedded_models):
+            raise TypeError(
+                f"Expected instance of type {self.embedded_models!r}, not "
+                f"{type(embedded_instance)!r}."
+            )
+        field_values = {}
+        add = embedded_instance._state.adding
+        for field in embedded_instance._meta.fields:
+            value = field.get_db_prep_save(
+                field.pre_save(embedded_instance, add), connection=connection
+            )
+            # Exclude unset primary keys (e.g. {'id': None}).
+            if field.primary_key and value is None:
+                continue
+            field_values[field.attname] = value
+        field_values["_label"] = embedded_instance._meta.label
+        # This instance will exist in the database soon.
+        embedded_instance._state.adding = False
+        return field_values
+
+    def get_transform(self, name):
+        transform = super().get_transform(name)
+        if transform:
+            return transform
+        field = None
+        for model in self.embedded_models:
+            with contextlib.suppress(FieldDoesNotExist):
+                field = model._meta.get_field(name)
+        if field is None:
+            raise FieldDoesNotExist(
+                f"The models of field '{self.name}' have no field named '{name}'."
+            )
+        return KeyTransformFactory(name, field)
+
+    def validate(self, value, model_instance):
+        super().validate(value, model_instance)
+        if not isinstance(value, self.embedded_models):
+            raise ValidationError(
+                f"Expected instance of type {self.embedded_models!r}, not {type(value)!r}."
+            )
+        for field in value._meta.fields:
+            attname = field.attname
+            field.validate(getattr(value, attname), model_instance)
+
+    def formfield(self, **kwargs):
+        raise NotImplementedError("PolymorphicEmbeddedModelField does not support forms.")
+
+    def _get_model_from_label(self, label):
+        return {model._meta.label: model for model in self.embedded_models}[label]
+
+
+class KeyTransform(Transform):
+    def __init__(self, key_name, ref_field, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.key_name = str(key_name)
+        self.ref_field = ref_field
+
+    def get_lookup(self, name):
+        return self.ref_field.get_lookup(name)
+
+    def get_transform(self, name):
+        """
+        Validate that `name` is either a field of an embedded model or a
+        lookup on an embedded model's field.
+        """
+        if transform := self.ref_field.get_transform(name):
+            return transform
+        suggested_lookups = difflib.get_close_matches(name, self.ref_field.get_lookups())
+        if suggested_lookups:
+            suggested_lookups = " or ".join(suggested_lookups)
+            suggestion = f", perhaps you meant {suggested_lookups}?"
+        else:
+            suggestion = "."
+        raise FieldDoesNotExist(
+            f"Unsupported lookup '{name}' for "
+            f"{self.ref_field.__class__.__name__} '{self.ref_field.name}'"
+            f"{suggestion}"
+        )
+
+    def as_mql(self, compiler, connection):
+        previous = self
+        key_transforms = []
+        while isinstance(previous, KeyTransform):
+            key_transforms.insert(0, previous.key_name)
+            previous = previous.lhs
+        mql = previous.as_mql(compiler, connection)
+        for key in key_transforms:
+            mql = {"$getField": {"input": mql, "field": key}}
+        return mql
+
+    @property
+    def output_field(self):
+        return self.ref_field
+
+
+class KeyTransformFactory:
+    def __init__(self, key_name, ref_field):
+        self.key_name = key_name
+        self.ref_field = ref_field
+
+    def __call__(self, *args, **kwargs):
+        return KeyTransform(self.key_name, self.ref_field, *args, **kwargs)

django_mongodb_backend/operations.py

@@ -122,6 +122,8 @@ def get_db_converters(self, expression):
             )
         elif internal_type == "JSONField":
             converters.append(self.convert_jsonfield_value)
+        elif internal_type == "PolymorphicEmbeddedModelField":
+            converters.append(self.convert_polymorphicembeddedmodelfield_value)
         elif internal_type == "TimeField":
             # Trunc(... output_field="TimeField") values must remain datetime
             # until Trunc.convert_value() so they can be converted from UTC
@@ -182,6 +184,19 @@ def convert_jsonfield_value(self, value, expression, connection):
         """
         return json.dumps(value)
 
+    def convert_polymorphicembeddedmodelfield_value(self, value, expression, connection):
+        if value is not None:
+            model_class = expression.output_field._get_model_from_label(value["_label"])
+            # Apply database converters to each field of the embedded model.
+            for field in model_class._meta.fields:
+                field_expr = Expression(output_field=field)
+                converters = connection.ops.get_db_converters(
+                    field_expr
+                ) + field_expr.get_db_converters(connection)
+                for converter in converters:
+                    value[field.attname] = converter(value[field.attname], field_expr, connection)
+        return value
+
     def convert_timefield_value(self, value, expression, connection):
         if value is not None:
             value = value.time()

docs/source/ref/models/fields.rst

@@ -313,3 +313,32 @@ These indexes use 0-based indexing.
 .. class:: ObjectIdField
 
     Stores an :class:`~bson.objectid.ObjectId`.
+
+``PolymorphicEmbeddedModelField``
+---------------------------------
+
+.. class:: PolymorphicEmbeddedModelField(embedded_models, **kwargs)
+
+    .. versionadded:: 5.2.0b2
+
+    Stores a model of type ``embedded_models``.
+
+    .. attribute:: embedded_models
+
+        This is a required argument that specifies a list of model classes
+        that may be embedded.
+
+        Each model class reference works just like
+        :attr:`.EmbeddedModelField.embedded_model`.
+
+    See :ref:`the embedded model topic guide
+    <polymorphic-embedded-model-field-example>` for more details and examples.
+
+.. admonition:: Migrations support is limited
+
+    :djadmin:`makemigrations` does not yet detect changes to embedded models,
+    nor does it create indexes or constraints for embedded models.
+
+.. admonition:: Forms are not supported
+
+    ``PolymorphicEmbeddedModelField``\s don't appear in model forms.

docs/source/topics/embedded-models.rst

@@ -181,3 +181,84 @@ For example, if the ``Tag`` model had an ``EmbeddedModelArrayField`` called
     >>> Post.objects.filter(tags__colors__name="blue")
     ...
     ValueError: Cannot perform multiple levels of array traversal in a query.
+
+.. _polymorphic-embedded-model-field-example:
+
+``PolymorphicEmbeddedModelField``
+---------------------------------
+
+The basics
+~~~~~~~~~~
+
+Let's consider this example::
+
+    from django.db import models
+
+    from django_mongodb_backend.fields import PolymorphicEmbeddedModelField
+    from django_mongodb_backend.models import EmbeddedModel
+
+
+    class Person(models.Model):
+        name = models.CharField(max_length=255)
+        pet = PolymorphicEmbeddedModelField(["Cat", "Dog"])
+
+        def __str__(self):
+            return self.name
+
+
+    class Cat(EmbeddedModel):
+        name = models.CharField(max_length=255)
+        purrs = models.BooleanField(default=True)
+
+        def __str__(self):
+            return self.name
+
+
+    class Dog(EmbeddedModel):
+        name = models.CharField(max_length=255)
+        barks = models.BooleanField(default=True)
+
+        def __str__(self):
+            return self.name
+
+
+The API is similar to that of Django's relational fields::
+
+    >>> bob = Person.objects.create(name="Bob", pet=Dog(name="Woofer"))
+    >>> bob.pet
+    <Dog: Woofer>
+    >>> bob.pet.name
+    'Woofer'
+    >>> bob = Person.objects.create(name="Fred", pet=Cat(name="Pheobe"))
+
+Represented in BSON, the person structures looks like this:
+
+.. code-block:: js
+
+    {
+      _id: ObjectId('685da4895e42adade0c8db29'),
+      name: 'Bob',
+     pet: { name: 'Woofer', barks: true, _label: 'myapp.Dog' }
+    },
+    {
+      _id: ObjectId('685da4925e42adade0c8db2a'),
+      name: 'Fred',
+      pet: { name: 'Pheobe', purrs: true, _label: 'myapp.Cat' }
+    }
+
+The ``_label`` field contains the model's
+:attr:`~django.db.models.Options.label`.
+
+Querying ``PolymorphicEmbeddedModelField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can query into a polymorphic embedded model field using the same double
+underscore syntax as relational fields. For example, to retrieve all people
+who have a pet named "Lassy"::
+
+    >>> Person.objects.filter(pet__name="Lassy")
+
+You can also filter on fields that aren't shared among the embedded models. For
+example, if you filter on ``barks``, you'll only get back people with dogs::
+
+    >>> Person.objects.filter(pet__barks=True)