Fixed Document with Proxied Fields Deletion Logic for EmbeddedDocuments

mongoengine/base.py

@@ -618,10 +618,6 @@ def _get_mixin_fields(base):
                     raise InvalidDocumentError("Reverse delete rules are not supported for EmbeddedDocuments (field: %s)" % field.name)
                 f.document_type.register_delete_rule(new_class, field.name, delete_rule)
 
-            proxy_class = getattr(field, 'proxy_class', None)
-            if proxy_class is not None:
-                new_class.register_proxy_field(field.name, proxy_class)
-
             if field.name and hasattr(Document, field.name) and EmbeddedDocument not in new_class.mro():
                 raise InvalidDocumentError("%s is a document method and not a valid field name" % field.name)
 
@@ -758,6 +754,8 @@ def __new__(cls, name, bases, attrs):
         unique_indexes = cls._unique_with_indexes(new_class)
         new_class._meta['unique_indexes'] = unique_indexes
 
+        from mongoengine import EmbeddedDocumentField
+
         for field_name, field in new_class._fields.items():
             # Check for custom primary key
             if field.primary_key:
@@ -770,6 +768,19 @@ def __new__(cls, name, bases, attrs):
                     # Make 'Document.id' an alias to the real primary key field
                     new_class.id = field
 
+            f = field
+            if isinstance(f, EmbeddedDocumentField):
+                for embedded_field_name, embedded_field in f.document_type._fields.items():
+                    proxy_class = getattr(embedded_field, 'proxy_class', None)
+                    if proxy_class is not None:
+                        composite_field_name = field_name + "." + embedded_field_name
+                        new_class._meta['proxy_fields'][composite_field_name] = proxy_class
+
+
+            proxy_class = getattr(field, 'proxy_class', None)
+            if proxy_class is not None:
+                new_class._meta['proxy_fields'][field_name] = proxy_class
+
         if not new_class._meta['id_field']:
             new_class._meta['id_field'] = 'id'
             new_class._fields['id'] = ObjectIdField(db_field='_id')

mongoengine/document.py

@@ -281,8 +281,13 @@ def delete(self, safe=False):
             for field_name in self._meta['proxy_fields']:
                 proxy_class = self._meta['proxy_fields'][field_name]
                 if hasattr(proxy_class, 'delete'):
-                    proxy = getattr(self, field_name)
-                    proxy.delete()
+                    sub_fields = field_name.split(".")
+                    obj = self
+                    while sub_fields:
+                        sub_f = sub_fields.pop(0)
+                        attr = getattr(obj, sub_f)
+                        obj = attr
+                    obj.delete()
             self.__class__.objects(pk=self.pk).delete(safe=safe)
         except pymongo.errors.OperationFailure, err:
             message = u'Could not delete document (%s)' % err.message
@@ -347,13 +352,6 @@ def register_delete_rule(cls, document_cls, field_name, rule):
         """
         cls._meta['delete_rules'][(document_cls, field_name)] = rule
 
-    @classmethod
-    def register_proxy_field(cls, field_name, proxy_class):
-        """This method registers fields with proxy classes to delete them when
-        removing this object.
-        """
-        cls._meta['proxy_fields'][field_name] = proxy_class
-
     @classmethod
     def drop_collection(cls):
         """Drops the entire collection associated with this

tests/fields.py

@@ -1620,6 +1620,22 @@ class DemoFile(Document):
             file = FileField()
         DemoFile.objects.create()
 
+    def _verify_file_delete(self, doc, file_attr_name):
+        obj = doc
+        sub_file_attrs = file_attr_name.split('.')
+        while sub_file_attrs:
+            sub_f_attr = sub_file_attrs.pop(0)
+            attr = getattr(obj, sub_f_attr)
+            obj = attr
+
+        grid_id = attr.grid_id
+        fs = attr.fs
+
+        doc.delete()
+
+        if grid_id:
+            self.assertFalse(fs.exists(grid_id))
+
     def test_file_delete_cleanup(self):
         """Ensure that the gridfs file is deleted when a document
         with a GridFSProxied Field is deleted"""
@@ -1635,23 +1651,77 @@ class TestImage(Document):
         testfile.file.put('Hello, World!')
         testfile.save()
 
-        testfile_grid_id = testfile.file.grid_id
-        testfile_fs = testfile.file.fs
-
-        testfile.delete()
-        self.assertFalse(testfile_fs.exists(testfile_grid_id))
+        self._verify_file_delete(testfile, 'file')
 
         TestImage.drop_collection()
 
         testimage = TestImage()
         testimage.image.put(open(TEST_IMAGE_PATH, 'r'))
         testimage.save()
 
-        testimage_grid_id = testimage.image.grid_id
-        testimage_fs = testimage.image.fs
+        self._verify_file_delete(testimage, 'image')
+
+    def test_document_delete_no_file_cleanup(self):
+        """Ensure that document deletion when a GridFSProxied
+        Field exists but is not present works as expected"""
+        class TestFile(Document):
+            f = FileField()
+
+        class TestImage(Document):
+            image = ImageField()
+
+        TestFile.drop_collection()
+
+        testfile = TestFile()
+        testfile.save()
+
+        self._verify_file_delete(testfile, 'f')
+
+        TestImage.drop_collection()
+
+        testimage = TestImage()
+        testimage.save()
+
+        self._verify_file_delete(testimage, 'image')
+
+    def test_embedded_file_delete_cleanup(self):
+        """Ensure that the gridfs file is deleted when a document
+        with an EmbeddedDocument and GridFSProxied Field is deleted"""
+        class EmbeddedFile(EmbeddedDocument):
+            f = FileField()
+
+        class TestFile(Document):
+            embed_file = EmbeddedDocumentField(EmbeddedFile)
+
+        class EmbeddedImage(EmbeddedDocument):
+            image = ImageField()
+
+        class TestImage(Document):
+            embed_image = EmbeddedDocumentField(EmbeddedImage)
+
+        TestFile.drop_collection()
+
+        testfile = TestFile()
+        testfile.embed_file = EmbeddedFile()
+        testfile.embed_file.f.put('Hello, World!')
+        testfile.save()
+
+        self._verify_file_delete(testfile, 'embed_file.f')
+
+        TestImage.drop_collection()
+
+        testimage = TestImage()
+        testimage.embed_image = EmbeddedImage()
+        testimage.embed_image.image.put(open(TEST_IMAGE_PATH, 'r'))
+        testimage.save()
+
+        self._verify_file_delete(testimage, 'embed_image.image')
+
+        testfile = TestFile()
+        testfile.embed_file = EmbeddedFile()
+        testfile.save()
 
-        testimage.delete()
-        self.assertFalse(testimage_fs.exists(testimage_grid_id))
+        self._verify_file_delete(testfile, 'embed_file.f')
 
     def test_file_field_no_default(self):