Added single and multifield uniqueness constraints

docs/userguide.rst

@@ -138,6 +138,21 @@ field::
 The :class:`User` object is automatically turned into a reference behind the
 scenes, and dereferenced when the :class:`Page` object is retrieved.
 
+Uniqueness constraints
+^^^^^^^^^^^^^^^^^^^^^^
+MongoEngine allows you to specify that a field should be unique across a
+collection by providing ``unique=True`` to a :class:`~mongoengine.Field`\ 's
+constructor. If you try to save a document that has the same value for a unique
+field as a document that is already in the database, a 
+:class:`~mongoengine.ValidationError` will be raised. You may also specify
+multi-field uniqueness constraints by using :attr:`unique_with`, which may be
+either a single field name, or a list or tuple of field names::
+
+    class User(Document):
+        username = StringField(unique=True)
+        first_name = StringField()
+        last_name = StringField(unique_with='last_name')
+
 Document collections
 --------------------
 Document classes that inherit **directly** from :class:`~mongoengine.Document`
@@ -172,10 +187,10 @@ Indexes
 -------
 You can specify indexes on collections to make querying faster. This is done
 by creating a list of index specifications called :attr:`indexes` in the
-:attr:`~Document.meta` dictionary, where an index specification may either be
+:attr:`~mongoengine.Document.meta` dictionary, where an index specification may
-a single field name, or a tuple containing multiple field names. A direction
+either be a single field name, or a tuple containing multiple field names. A
-may be specified on fields by prefixing the field name with a **+** or a **-**
+direction may be specified on fields by prefixing the field name with a **+**
-sign. Note that direction only matters on multi-field indexes. ::
+or a **-** sign. Note that direction only matters on multi-field indexes. ::
 
     class Page(Document):
         title = StringField()
@@ -186,11 +201,11 @@ sign. Note that direction only matters on multi-field indexes. ::
         
 Ordering
 --------
-
+A default ordering can be specified for your
-A default ordering can be specified for your :class:`~mongoengine.queryset.QuerySet`
+:class:`~mongoengine.queryset.QuerySet` using the :attr:`ordering` attribute of
-using the :attr:`ordering` attributeof :attr:`~Document.meta`. 
+:attr:`~mongoengine.Document.meta`.  Ordering will be applied when the
-Ordering will be applied when the ``QuerySet`` is created, and can be 
+:class:`~mongoengine.queryset.QuerySet` is created, and can be overridden by
-overridden by subsequent calls to :meth:`~mongoengine.QuerySet.order_by`. ::
+subsequent calls to :meth:`~mongoengine.queryset.QuerySet.order_by`. ::
 
     from datetime import datetime
 
@@ -202,9 +217,14 @@ overridden by subsequent calls to :meth:`~mongoengine.QuerySet.order_by`. ::
             'ordering': ['-published_date']
         }
 
-    blog_post_1 = BlogPost(title="Blog Post #1", published_date=datetime(2010, 1, 5, 0, 0 ,0))
+    blog_post_1 = BlogPost(title="Blog Post #1")
-    blog_post_2 = BlogPost(title="Blog Post #2", published_date=datetime(2010, 1, 6, 0, 0 ,0))
+    blog_post_1.published_date = datetime(2010, 1, 5, 0, 0 ,0))
-    blog_post_3 = BlogPost(title="Blog Post #3", published_date=datetime(2010, 1, 7, 0, 0 ,0))
+
+    blog_post_2 = BlogPost(title="Blog Post #2") 
+    blog_post_2.published_date = datetime(2010, 1, 6, 0, 0 ,0))
+
+    blog_post_3 = BlogPost(title="Blog Post #3")
+    blog_post_3.published_date = datetime(2010, 1, 7, 0, 0 ,0))
 
     blog_post_1.save()
     blog_post_2.save()

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):
+    def _collection(self):
-        if not self._cursor_obj:
+        """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

tests/document.py

@@ -237,7 +237,7 @@ class DocumentTest(unittest.TestCase):
         BlogPost.drop_collection()
 
         info = BlogPost.objects._collection.index_information()
-        self.assertEqual(len(info), 0)
+        self.assertEqual(len(info), 4) # _id, types, '-date', ('cat', 'date')
 
         # Indexes are lazy so use list() to perform query
         list(BlogPost.objects)
@@ -248,6 +248,45 @@ class DocumentTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+    def test_unique(self):
+        """Ensure that uniqueness constraints are applied to fields.
+        """
+        class BlogPost(Document):
+            title = StringField()
+            slug = StringField(unique=True)
+
+        BlogPost.drop_collection()
+
+        post1 = BlogPost(title='test1', slug='test')
+        post1.save()
+
+        # Two posts with the same slug is not allowed
+        post2 = BlogPost(title='test2', slug='test')
+        self.assertRaises(ValidationError, post2.save)
+
+        class Date(EmbeddedDocument):
+            year = IntField(name='yr')
+
+        class BlogPost(Document):
+            title = StringField()
+            date = EmbeddedDocumentField(Date)
+            slug = StringField(unique_with='date.year')
+
+        BlogPost.drop_collection()
+
+        post1 = BlogPost(title='test1', date=Date(year=2009), slug='test')
+        post1.save()
+
+        # day is different so won't raise exception
+        post2 = BlogPost(title='test2', date=Date(year=2010), slug='test')
+        post2.save()
+
+        # Now there will be two docs with the same slug and the same day: fail
+        post3 = BlogPost(title='test3', date=Date(year=2010), slug='test')
+        self.assertRaises(ValidationError, post3.save)
+
+        BlogPost.drop_collection()
+
     def test_creation(self):
         """Ensure that document may be created using keyword arguments.
         """

tests/fields.py

@@ -231,6 +231,9 @@ class FieldTest(unittest.TestCase):
             content = StringField()
             author = ReferenceField(User)
 
+        User.drop_collection()
+        BlogPost.drop_collection()
+
         self.assertRaises(ValidationError, ReferenceField, EmbeddedDocument)
 
         user = User(name='Test User')

tests/queryset.py

@@ -144,9 +144,12 @@ class QuerySetTest(unittest.TestCase):
                 'ordering': ['-published_date']
             }
 
-        blog_post_1 = BlogPost(title="Blog Post #1", published_date=datetime(2010, 1, 5, 0, 0 ,0))
+        blog_post_1 = BlogPost(title="Blog Post #1", 
-        blog_post_2 = BlogPost(title="Blog Post #2", published_date=datetime(2010, 1, 6, 0, 0 ,0))
+                               published_date=datetime(2010, 1, 5, 0, 0 ,0))
-        blog_post_3 = BlogPost(title="Blog Post #3", published_date=datetime(2010, 1, 7, 0, 0 ,0))
+        blog_post_2 = BlogPost(title="Blog Post #2", 
+                               published_date=datetime(2010, 1, 6, 0, 0 ,0))
+        blog_post_3 = BlogPost(title="Blog Post #3", 
+                               published_date=datetime(2010, 1, 7, 0, 0 ,0))
 
         blog_post_1.save()
         blog_post_2.save()
@@ -161,6 +164,8 @@ class QuerySetTest(unittest.TestCase):
         first_post = BlogPost.objects.order_by("+published_date").first()
         self.assertEqual(first_post.title, "Blog Post #1")
 
+        BlogPost.drop_collection()
+
     def test_find_embedded(self):
         """Ensure that an embedded document is properly returned from a query.
         """