Added single and multifield uniqueness constraints

mongoengine/base.py

@@ -1,4 +1,4 @@
-from queryset import QuerySetManager
+from queryset import QuerySet, QuerySetManager
 
 import pymongo
 
@@ -12,10 +12,13 @@ class BaseField(object):
     may be added to subclasses of `Document` to define a document's schema.
     """
     
-    def __init__(self, name=None, required=False, default=None):
+    def __init__(self, name=None, required=False, default=None, unique=False,
+                 unique_with=None):
         self.name = name
         self.required = required
         self.default = default
+        self.unique = bool(unique or unique_with)
+        self.unique_with = unique_with
 
     def __get__(self, instance, owner):
         """Descriptor for retrieving a value from a field in a document. Do 
@@ -176,6 +179,35 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
         new_class = super_new(cls, name, bases, attrs)
         new_class.objects = QuerySetManager()
 
+        # Generate a list of indexes needed by uniqueness constraints
+        unique_indexes = []
+        for field_name, field in new_class._fields.items():
+            if field.unique:
+                field.required = True
+                unique_fields = [field_name]
+
+                # Add any unique_with fields to the back of the index spec
+                if field.unique_with:
+                    if isinstance(field.unique_with, basestring):
+                        field.unique_with = [field.unique_with]
+
+                    # Convert unique_with field names to real field names
+                    unique_with = []
+                    for other_name in field.unique_with:
+                        parts = other_name.split('.')
+                        # Lookup real name
+                        parts = QuerySet._lookup_field(new_class, parts)
+                        name_parts = [part.name for part in parts]
+                        unique_with.append('.'.join(name_parts))
+                        # Unique field should be required
+                        parts[-1].required = True
+                    unique_fields += unique_with
+
+                # Add the new index to the list
+                index = [(field, pymongo.ASCENDING) for field in unique_fields]
+                unique_indexes.append(index)
+        new_class._meta['unique_indexes'] = unique_indexes
+
         return new_class
 
 

mongoengine/document.py

@@ -2,8 +2,10 @@ from base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument,
                   ValidationError)
 from connection import _get_db
 
+import pymongo
 
-__all__ = ['Document', 'EmbeddedDocument']
+
+__all__ = ['Document', 'EmbeddedDocument', 'ValidationError']
 
 
 class EmbeddedDocument(BaseDocument):
@@ -53,13 +55,18 @@ class Document(BaseDocument):
 
     __metaclass__ = TopLevelDocumentMetaclass
 
-    def save(self):
+    def save(self, safe=True):
         """Save the :class:`~mongoengine.Document` to the database. If the
         document already exists, it will be updated, otherwise it will be
         created.
         """
         self.validate()
-        object_id = self.__class__.objects._collection.save(self.to_mongo())
+        doc = self.to_mongo()
+        try:
+            object_id = self.__class__.objects._collection.save(doc, safe=safe)
+        except pymongo.errors.OperationFailure, err:
+            raise ValidationError('Tried to safe duplicate unique keys (%s)'
+                                  % str(err))
         self.id = self._fields['id'].to_python(object_id)
 
     def delete(self):

mongoengine/queryset.py

@@ -17,7 +17,8 @@ class QuerySet(object):
     
     def __init__(self, document, collection):
         self._document = document
-        self._collection = collection
+        self._collection_obj = collection
+        self._accessed_collection = False
         self._query = {}
 
         # If inheritance is allowed, only return instances and instances of
@@ -59,17 +60,30 @@ class QuerySet(object):
         return self
 
     @property
-    def _cursor(self):
-        if not self._cursor_obj:
+    def _collection(self):
+        """Property that returns the collection object. This allows us to
+        perform operations only if the collection is accessed.
+        """
+        if not self._accessed_collection:
+            self._accessed_collection = True
+
             # Ensure document-defined indexes are created
             if self._document._meta['indexes']:
                 for key_or_list in self._document._meta['indexes']:
                     self.ensure_index(key_or_list)
 
+            # Ensure indexes created by uniqueness constraints
+            for index in self._document._meta['unique_indexes']:
+                self._collection.ensure_index(index, unique=True)
+
             # If _types is being used (for polymorphism), it needs an index
             if '_types' in self._query:
                 self._collection.ensure_index('_types')
+        return self._collection_obj
 
+    @property
+    def _cursor(self):
+        if not self._cursor_obj:
             self._cursor_obj = self._collection.find(self._query)
             
             # apply default ordering