Improve Reverse Delete Rules

AUTHORS

@@ -264,3 +264,4 @@ that much better:
  * oleksandr-l5 (https://github.com/oleksandr-l5)
  * Ido Shraga (https://github.com/idoshr)
  * Terence Honles (https://github.com/terencehonles)
+ * Chris Combs (https://github.com/combscCode)

mongoengine/base/__init__.py

@@ -14,7 +14,9 @@
     # common
     "UPDATE_OPERATORS",
     "_document_registry",
+    "_undefined_document_delete_rules",
     "get_document",
+    "update_document_if_registered",
     # datastructures
     "BaseDict",
     "BaseList",

mongoengine/base/common.py

@@ -1,6 +1,14 @@
+from collections import defaultdict
+
 from mongoengine.errors import NotRegistered
 
-__all__ = ("UPDATE_OPERATORS", "get_document", "_document_registry")
+__all__ = (
+    "UPDATE_OPERATORS",
+    "get_document",
+    "update_document_if_registered",
+    "_document_registry",
+    "_undefined_document_delete_rules",
+)
 
 
 UPDATE_OPERATORS = {
@@ -23,6 +31,7 @@
 
 
 _document_registry = {}
+_undefined_document_delete_rules = defaultdict(list)
 
 
 def get_document(name):
@@ -60,3 +69,13 @@ def get_doc_alias(doc_cls):
         for doc_cls in _document_registry.values()
         if get_doc_alias(doc_cls) == connection_alias
     ]
+
+
+def update_document_if_registered(document):
+    """Converts a string to a Document if registered in the registry."""
+    if isinstance(document, str):
+        try:
+            return get_document(document)
+        except NotRegistered:
+            pass
+    return document

mongoengine/base/metaclasses.py

@@ -1,14 +1,17 @@
 import itertools
 import warnings
 
-from mongoengine.base.common import _document_registry
+from mongoengine.base.common import (
+    _document_registry,
+    _undefined_document_delete_rules,
+)
 from mongoengine.base.fields import (
     BaseField,
     ComplexBaseField,
     ObjectIdField,
 )
 from mongoengine.common import _import_class
-from mongoengine.errors import InvalidDocumentError
+from mongoengine.errors import InvalidDocumentError, NotRegistered
 from mongoengine.queryset import (
     DO_NOTHING,
     DoesNotExist,
@@ -206,7 +209,14 @@ def __new__(mcs, name, bases, attrs):
                         "EmbeddedDocuments (field: %s)" % field.name
                     )
                     raise InvalidDocumentError(msg)
-                f.document_type.register_delete_rule(new_class, field.name, delete_rule)
+                try:
+                    f.document_type.register_delete_rule(
+                        new_class, field.name, delete_rule
+                    )
+                except NotRegistered:
+                    _undefined_document_delete_rules[f.document_type_obj].append(
+                        (new_class, field.name, delete_rule)
+                    )
 
             if (
                 field.name
@@ -362,6 +372,13 @@ def __new__(mcs, name, bases, attrs):
 
         # Call super and get the new class
         new_class = super_new(mcs, name, bases, attrs)
+        # Find any lazy delete rules and apply to current doc.
+        if new_class._class_name in _undefined_document_delete_rules:
+            rules_tuple_list = _undefined_document_delete_rules.pop(
+                new_class._class_name
+            )
+            for document_cls, field_name, rule in rules_tuple_list:
+                new_class.register_delete_rule(document_cls, field_name, rule)
 
         meta = new_class._meta
 

mongoengine/fields.py

@@ -30,7 +30,9 @@
     GeoJsonBaseField,
     LazyReference,
     ObjectIdField,
+    _undefined_document_delete_rules,
     get_document,
+    update_document_if_registered,
 )
 from mongoengine.base.utils import LazyRegexCompiler
 from mongoengine.common import _import_class
@@ -1435,6 +1437,39 @@ def sync_all(self):
             self.owner_document.objects(**filter_kwargs).update(**update_kwargs)
 
 
+class GenericReferenceDeleteHandler:
+    """Used to make delete rules work for GenericReferenceFields.
+
+    Since delete rules are registered on single documents, we'll always need
+    something like this to make a generic reference (AKA, a reference to
+    multiple documents) with delete rules work.
+    """
+
+    def __init__(self, documents):
+        self.documents = documents
+
+    def __getattr__(self, name):
+        raise NotImplementedError(
+            f"{self.__name__} is intended only to be used "
+            "to enable generic reference delete rules. You "
+            "are trying to access undefined attributes."
+        )
+
+    def register_delete_rule(self, document_cls, field_name, rule):
+        for doc in self.documents:
+            doc = update_document_if_registered(doc)
+            if isinstance(doc, str):
+                _undefined_document_delete_rules[doc].append(
+                    (
+                        document_cls,
+                        field_name,
+                        rule,
+                    )
+                )
+            else:
+                doc.register_delete_rule(document_cls, field_name, rule)
+
+
 class GenericReferenceField(BaseField):
     """A reference to *any* :class:`~mongoengine.document.Document` subclass
     that will be automatically dereferenced on access (lazily).
@@ -1453,8 +1488,11 @@ class GenericReferenceField(BaseField):
         * You can use the choices param to limit the acceptable Document types
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, reverse_delete_rule=DO_NOTHING, **kwargs):
         choices = kwargs.pop("choices", None)
+        self.reverse_delete_rule = reverse_delete_rule
+        if self.reverse_delete_rule is not DO_NOTHING and not choices:
+            raise ValidationError("choices must be set to use reverse_delete_rules")
         super().__init__(*args, **kwargs)
         self.choices = []
         # Keep the choices as a list of allowed Document class names
@@ -1472,6 +1510,15 @@ def __init__(self, *args, **kwargs):
                         "Document subclasses and/or str"
                     )
 
+    @property
+    def document_type(self):
+        # This property is exposed purely for enabling reverse_delete_rule
+        # on this class. Do not attempt to use it in any other way, if you
+        # do a NotImplementedError will be raised.
+        if not self.choices:
+            return None
+        return GenericReferenceDeleteHandler(self.choices)
+
     def _validate_choices(self, value):
         if isinstance(value, dict):
             # If the field has not been dereferenced, it is still a dict

tests/document/test_instance.py

@@ -13,7 +13,11 @@
 
 from mongoengine import *
 from mongoengine import signals
-from mongoengine.base import _document_registry, get_document
+from mongoengine.base import (
+    _document_registry,
+    _undefined_document_delete_rules,
+    get_document,
+)
 from mongoengine.connection import get_db
 from mongoengine.context_managers import query_counter, switch_db
 from mongoengine.errors import (
@@ -2516,6 +2520,52 @@ class BlogPost(Document):
         author.delete()
         assert self.Person.objects.count() == 1
 
+    def test_lazy_delete_rules(self):
+        """Ensure that a document does not need to be defined to reference it
+        in a ReferenceField."""
+
+        assert not _undefined_document_delete_rules.get("BlogPostLazy")
+
+        # named "lazy" to ensure these Documents don't exist in the
+        # document registry
+        class CommentLazy(Document):
+            text = StringField()
+            post = ReferenceField("BlogPostLazy", reverse_delete_rule=CASCADE)
+
+        assert len(_undefined_document_delete_rules.get("BlogPostLazy")) == 1
+
+        class CommentDosLazy(Document):
+            textdos = StringField()
+            postdos = ReferenceField("BlogPostLazy", reverse_delete_rule=CASCADE)
+
+        assert len(_undefined_document_delete_rules.get("BlogPostLazy")) == 2
+
+        class BlogPostLazy(Document):
+            content = StringField()
+            author = ReferenceField(self.Person, reverse_delete_rule=CASCADE)
+
+        assert not _undefined_document_delete_rules.get("BlogPostLazy")
+
+        self.Person.drop_collection()
+        BlogPostLazy.drop_collection()
+        CommentLazy.drop_collection()
+
+        author = self.Person(name="Test User")
+        author.save()
+
+        post = BlogPostLazy(content="Watched some TV")
+        post.author = author
+        post.save()
+
+        comment = CommentLazy(text="Kudos.")
+        comment.post = post
+        comment.save()
+
+        # Delete the Person, which should lead to deletion of the BlogPost,
+        # and, recursively to the Comment, too
+        author.delete()
+        assert CommentLazy.objects.count() == 0
+
     def subclasses_and_unique_keys_works(self):
         class A(Document):
             pass

tests/fields/test_fields.py

@@ -5,6 +5,9 @@
 from bson import SON, DBRef, ObjectId
 
 from mongoengine import (
+    CASCADE,
+    DENY,
+    PULL,
     BooleanField,
     ComplexDateTimeField,
     DateField,
@@ -1623,6 +1626,112 @@ class Bookmark(Document):
         assert bm.bookmark_object == link_1
         assert isinstance(bm.bookmark_object, Link)
 
+    def test_generic_reference_deny(self):
+        """Ensure that a GenericReferenceField properly enforces DENY delete rules"""
+
+        class Link(Document):
+            title = StringField()
+            meta = {"allow_inheritance": False}
+
+        class Bookmark(Document):
+            bookmark_object = GenericReferenceField(
+                reverse_delete_rule=DENY, choices=[Link]
+            )
+
+        Link.drop_collection()
+        Bookmark.drop_collection()
+
+        link_1 = Link(title="Pitchfork")
+        link_1.save()
+
+        bm = Bookmark(bookmark_object=link_1)
+        bm.save()
+
+        with pytest.raises(OperationError):
+            link_1.delete()
+        # once bm is gone, we should be able to delete link_1 just fine
+        bm.delete()
+        link_1.delete()
+
+    def test_generic_reference_cascade(self):
+        """Ensure that a GenericReferenceField properly enforces CASCADE delete rules"""
+
+        class Link(Document):
+            title = StringField()
+            meta = {"allow_inheritance": False}
+
+        class Bookmark(Document):
+            bookmark_object = GenericReferenceField(
+                reverse_delete_rule=CASCADE, choices=[Link]
+            )
+
+        Link.drop_collection()
+        Bookmark.drop_collection()
+
+        link_1 = Link(title="Pitchfork")
+        link_1.save()
+
+        bm = Bookmark(bookmark_object=link_1)
+        bm.save()
+
+        assert Bookmark.objects.count() == 1
+        link_1.delete()
+        assert Bookmark.objects.count() == 0
+
+    def test_generic_reference_pull(self):
+        """Ensure that a GenericReferenceField properly enforces PULL delete rules"""
+
+        class Link(Document):
+            title = StringField()
+            meta = {"allow_inheritance": False}
+
+        class Blog(Document):
+            links = ListField(
+                GenericReferenceField(reverse_delete_rule=PULL, choices=[Link])
+            )
+
+        Link.drop_collection()
+        Blog.drop_collection()
+
+        link_1 = Link(title="Pitchfork")
+        link_1.save()
+
+        link_2 = Link(title="Pitchfork 2: Electric Boogaloo")
+        link_2.save()
+        blog = Blog(links=[link_1, link_2])
+        blog.save()
+
+        assert len(blog.links) == 2
+        link_1.delete()
+        blog.reload()
+        assert len(blog.links) == 1
+        link_2.delete()
+        blog.reload()
+        assert len(blog.links) == 0
+
+    def test_generic_reference_deny_errors(self):
+        class Link(Document):
+            title = StringField()
+            meta = {"allow_inheritance": False}
+
+        class Bookmark(Document):
+            bookmark_object = GenericReferenceField(
+                reverse_delete_rule=DENY, choices=[Link]
+            )
+
+        class NotRegisteredInChoices(Document):
+            title = StringField()
+
+        with pytest.raises(ValidationError):
+            nr = NotRegisteredInChoices(title="hello")
+            bm = Bookmark(bookmark_object=nr)
+            bm.validate()
+
+        with pytest.raises(ValidationError):
+
+            class Oops(Document):
+                bad_field = GenericReferenceField(reverse_delete_rule=DENY)
+
     def test_generic_reference_list(self):
         """Ensure that a ListField properly dereferences generic references."""