From 4d695a3544cf446be9210e9d85977041fb1893e0 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Fri, 8 Jan 2010 12:04:11 +0000 Subject: [PATCH] Added single and multifield uniqueness constraints --- docs/userguide.rst | 44 ++++++++++++++++++++++++++++++----------- mongoengine/base.py | 36 +++++++++++++++++++++++++++++++-- mongoengine/document.py | 13 +++++++++--- mongoengine/queryset.py | 20 ++++++++++++++++--- tests/document.py | 41 +++++++++++++++++++++++++++++++++++++- tests/fields.py | 3 +++ tests/queryset.py | 11 ++++++++--- 7 files changed, 144 insertions(+), 24 deletions(-) diff --git a/docs/userguide.rst b/docs/userguide.rst index 405f7418..331980e4 100644 --- a/docs/userguide.rst +++ b/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 -a single field name, or a tuple containing multiple field names. A direction -may be specified on fields by prefixing the field name with a **+** or a **-** -sign. Note that direction only matters on multi-field indexes. :: +:attr:`~mongoengine.Document.meta` dictionary, where an index specification may +either be a single field name, or a tuple containing multiple field names. A +direction may be specified on fields by prefixing the field name with a **+** +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 :class:`~mongoengine.queryset.QuerySet` -using the :attr:`ordering` attributeof :attr:`~Document.meta`. -Ordering will be applied when the ``QuerySet`` is created, and can be -overridden by subsequent calls to :meth:`~mongoengine.QuerySet.order_by`. :: +A default ordering can be specified for your +:class:`~mongoengine.queryset.QuerySet` using the :attr:`ordering` attribute of +:attr:`~mongoengine.Document.meta`. Ordering will be applied when the +:class:`~mongoengine.queryset.QuerySet` is created, and can be overridden 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_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 = BlogPost(title="Blog Post #1") + blog_post_1.published_date = datetime(2010, 1, 5, 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() diff --git a/mongoengine/base.py b/mongoengine/base.py index f930f06c..76a2c0e9 100644 --- a/mongoengine/base.py +++ b/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 diff --git a/mongoengine/document.py b/mongoengine/document.py index 1afa1e41..daae230c 100644 --- a/mongoengine/document.py +++ b/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): diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 66ab543c..cb8d826b 100644 --- a/mongoengine/queryset.py +++ b/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 diff --git a/tests/document.py b/tests/document.py index 1ba590e3..a2f74649 100644 --- a/tests/document.py +++ b/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. """ diff --git a/tests/fields.py b/tests/fields.py index f61d45a8..affb0c9b 100644 --- a/tests/fields.py +++ b/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') diff --git a/tests/queryset.py b/tests/queryset.py index 6bee363a..12e95934 100644 --- a/tests/queryset.py +++ b/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_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 = BlogPost(title="Blog Post #1", + published_date=datetime(2010, 1, 5, 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. """