diff --git a/mongoengine/base.py b/mongoengine/base.py index 3dd2cb02..29de82fa 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -190,6 +190,8 @@ class DocumentMetaclass(type): new_class = super_new(cls, name, bases, attrs) for field in new_class._fields.values(): field.owner_document = new_class + if hasattr(field, 'delete_rule') and field.delete_rule: + field.document_type._meta['delete_rules'][(new_class, field.name)] = field.delete_rule module = attrs.get('__module__') @@ -258,6 +260,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): 'index_drop_dups': False, 'index_opts': {}, 'queryset_class': QuerySet, + 'delete_rules': {}, } meta.update(base_meta) diff --git a/mongoengine/document.py b/mongoengine/document.py index fef737db..06867168 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -6,9 +6,16 @@ from connection import _get_db import pymongo -__all__ = ['Document', 'EmbeddedDocument', 'ValidationError', 'OperationError'] +__all__ = ['Document', 'EmbeddedDocument', 'ValidationError', 'OperationError', + 'DO_NOTHING', 'NULLIFY', 'CASCADE', 'DENY'] +# Delete rules +DO_NOTHING = 0 +NULLIFY = 1 +CASCADE = 2 +DENY = 3 + class EmbeddedDocument(BaseDocument): """A :class:`~mongoengine.Document` that isn't stored in its own collection. :class:`~mongoengine.EmbeddedDocument`\ s should be used as @@ -92,6 +99,13 @@ class Document(BaseDocument): :param safe: check if the operation succeeded before returning """ + for rule_entry in self._meta['delete_rules']: + document_cls, field_name = rule_entry + rule = self._meta['delete_rules'][rule_entry] + + if rule == CASCADE: + document_cls.objects(**{field_name: self.id}).delete(safe=safe) + id_field = self._meta['id_field'] object_id = self._fields[id_field].to_mongo(self[id_field]) try: @@ -100,6 +114,17 @@ class Document(BaseDocument): message = u'Could not delete document (%s)' % err.message raise OperationError(message) + @classmethod + def register_delete_rule(cls, document_cls, field_name, rule): + """This method registers the delete rules to apply when removing this + object. This could go into the Document class. + """ + if rule == DO_NOTHING: + return + + cls._meta['delete_rules'][(document_cls, field_name)] = rule + + def reload(self): """Reloads all attributes from the database. diff --git a/mongoengine/fields.py b/mongoengine/fields.py index e95fd65e..01ec1f7b 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -417,12 +417,13 @@ class ReferenceField(BaseField): access (lazily). """ - def __init__(self, document_type, **kwargs): + def __init__(self, document_type, delete_rule=None, **kwargs): if not isinstance(document_type, basestring): if not issubclass(document_type, (Document, basestring)): raise ValidationError('Argument to ReferenceField constructor ' 'must be a document class or a string') self.document_type_obj = document_type + self.delete_rule = delete_rule super(ReferenceField, self).__init__(**kwargs) @property diff --git a/tests/document.py b/tests/document.py index c0567632..d5807c90 100644 --- a/tests/document.py +++ b/tests/document.py @@ -624,6 +624,31 @@ class DocumentTest(unittest.TestCase): BlogPost.drop_collection() + + def test_cascade_delete(self): + """Ensure that a referenced document is also deleted upon deletion. + """ + + class BlogPost(Document): + meta = {'collection': 'blogpost_1'} + content = StringField() + author = ReferenceField(self.Person, delete_rule=CASCADE) + + self.Person.drop_collection() + BlogPost.drop_collection() + + author = self.Person(name='Test User') + author.save() + + post = BlogPost(content = 'Watched some TV') + post.author = author + post.save() + + # Delete the Person, which should lead to deletion of the BlogPost, too + author.delete() + self.assertEqual(len(BlogPost.objects), 0) + + def tearDown(self): self.Person.drop_collection()