Merge remote branch 'nvie/dev' into dev

mongoengine/base.py

@@ -1,5 +1,6 @@
 from queryset import QuerySet, QuerySetManager
 from queryset import DoesNotExist, MultipleObjectsReturned
+from queryset import DO_NOTHING
 
 import sys
 import pymongo
@@ -203,6 +204,10 @@ class DocumentMetaclass(type):
         new_class = super_new(cls, name, bases, attrs)
         for field in new_class._fields.values():
             field.owner_document = new_class
+            delete_rule = getattr(field, 'reverse_delete_rule', DO_NOTHING)
+            if delete_rule != DO_NOTHING:
+                field.document_type.register_delete_rule(new_class, field.name,
+                        delete_rule)
 
         module = attrs.get('__module__')
 
@@ -271,6 +276,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
             'index_drop_dups': False,
             'index_opts': {},
             'queryset_class': QuerySet,
+            'delete_rules': {},
         }
         meta.update(base_meta)
 
@@ -506,6 +512,7 @@ if sys.version_info < (2, 5):
     # Prior to Python 2.5, Exception was an old-style class
     import types
     def subclass_exception(name, parents, unused):
+        import types
         return types.ClassType(name, parents, {})
 else:
     def subclass_exception(name, parents, module):

mongoengine/document.py

@@ -100,6 +100,14 @@ 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.
+        """
+        cls._meta['delete_rules'][(document_cls, field_name)] = rule
+
+
     def reload(self):
         """Reloads all attributes from the database.
 

mongoengine/fields.py

@@ -1,4 +1,5 @@
 from base import BaseField, ObjectIdField, ValidationError, get_document
+from queryset import DO_NOTHING
 from document import Document, EmbeddedDocument
 from connection import _get_db
 from operator import itemgetter
@@ -455,12 +456,13 @@ class ReferenceField(BaseField):
     access (lazily).
     """
 
-    def __init__(self, document_type, **kwargs):
+    def __init__(self, document_type, reverse_delete_rule=DO_NOTHING, **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.reverse_delete_rule = reverse_delete_rule
         super(ReferenceField, self).__init__(**kwargs)
 
     @property

mongoengine/queryset.py

@@ -10,11 +10,18 @@ import copy
 import itertools
 
 __all__ = ['queryset_manager', 'Q', 'InvalidQueryError',
-           'InvalidCollectionError']
+           'InvalidCollectionError', 'DO_NOTHING', 'NULLIFY', 'CASCADE', 'DENY']
+
 
 # The maximum number of items to display in a QuerySet.__repr__
 REPR_OUTPUT_SIZE = 20
 
+# Delete rules
+DO_NOTHING = 0
+NULLIFY = 1
+CASCADE = 2
+DENY = 3
+
 
 class DoesNotExist(Exception):
     pass
@@ -947,6 +954,28 @@ class QuerySet(object):
 
         :param safe: check if the operation succeeded before returning
         """
+        doc = self._document
+
+        # Check for DENY rules before actually deleting/nullifying any other
+        # references
+        for rule_entry in doc._meta['delete_rules']:
+            document_cls, field_name = rule_entry
+            rule = doc._meta['delete_rules'][rule_entry]
+            if rule == DENY and document_cls.objects(**{field_name + '__in': self}).count() > 0:
+                msg = u'Could not delete document (at least %s.%s refers to it)' % \
+                        (document_cls.__name__, field_name)
+                raise OperationError(msg)
+
+        for rule_entry in doc._meta['delete_rules']:
+            document_cls, field_name = rule_entry
+            rule = doc._meta['delete_rules'][rule_entry]
+            if rule == CASCADE:
+                document_cls.objects(**{field_name + '__in': self}).delete(safe=safe)
+            elif rule == NULLIFY:
+                document_cls.objects(**{field_name + '__in': self}).update(
+                        safe_update=safe,
+                        **{'unset__%s' % field_name: 1})
+
         self._collection.remove(self._query, safe=safe)
 
     @classmethod