diff --git a/mongoengine/base.py b/mongoengine/base.py index ede90835..9a6b5f12 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -306,6 +306,30 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): for spec in meta['indexes']] + base_indexes new_class._meta['indexes'] = user_indexes + unique_indexes = cls._unique_with_indexes(new_class) + new_class._meta['unique_indexes'] = unique_indexes + + for field_name, field in new_class._fields.items(): + # Check for custom primary key + if field.primary_key: + current_pk = new_class._meta['id_field'] + if current_pk and current_pk != field_name: + raise ValueError('Cannot override primary key field') + + if not current_pk: + new_class._meta['id_field'] = field_name + # Make 'Document.id' an alias to the real primary key field + new_class.id = field + + if not new_class._meta['id_field']: + new_class._meta['id_field'] = 'id' + new_class._fields['id'] = ObjectIdField(db_field='_id') + new_class.id = new_class._fields['id'] + + return new_class + + @classmethod + def _unique_with_indexes(cls, new_class, namespace=""): unique_indexes = [] for field_name, field in new_class._fields.items(): # Generate a list of indexes needed by uniqueness constraints @@ -331,28 +355,16 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): unique_fields += unique_with # Add the new index to the list - index = [(f, pymongo.ASCENDING) for f in unique_fields] + index = [("%s%s" % (namespace, f), pymongo.ASCENDING) for f in unique_fields] unique_indexes.append(index) - # Check for custom primary key - if field.primary_key: - current_pk = new_class._meta['id_field'] - if current_pk and current_pk != field_name: - raise ValueError('Cannot override primary key field') + # Grab any embedded document field unique indexes + if field.__class__.__name__ == "EmbeddedDocumentField": + field_namespace = "%s." % field_name + unique_indexes += cls._unique_with_indexes(field.document_type, + field_namespace) - if not current_pk: - new_class._meta['id_field'] = field_name - # Make 'Document.id' an alias to the real primary key field - new_class.id = field - - new_class._meta['unique_indexes'] = unique_indexes - - if not new_class._meta['id_field']: - new_class._meta['id_field'] = 'id' - new_class._fields['id'] = ObjectIdField(db_field='_id') - new_class.id = new_class._fields['id'] - - return new_class + return unique_indexes class BaseDocument(object): diff --git a/tests/document.py b/tests/document.py index 77c8269b..8f47ec3c 100644 --- a/tests/document.py +++ b/tests/document.py @@ -362,6 +362,10 @@ class DocumentTest(unittest.TestCase): post2 = BlogPost(title='test2', slug='test') self.assertRaises(OperationError, post2.save) + + def test_unique_with(self): + """Ensure that unique_with constraints are applied to fields. + """ class Date(EmbeddedDocument): year = IntField(db_field='yr') @@ -385,6 +389,63 @@ class DocumentTest(unittest.TestCase): BlogPost.drop_collection() + def test_unique_embedded_document(self): + """Ensure that uniqueness constraints are applied to fields on embedded documents. + """ + class SubDocument(EmbeddedDocument): + year = IntField(db_field='yr') + slug = StringField(unique=True) + + class BlogPost(Document): + title = StringField() + sub = EmbeddedDocumentField(SubDocument) + + BlogPost.drop_collection() + + post1 = BlogPost(title='test1', sub=SubDocument(year=2009, slug="test")) + post1.save() + + # sub.slug is different so won't raise exception + post2 = BlogPost(title='test2', sub=SubDocument(year=2010, slug='another-slug')) + post2.save() + + # Now there will be two docs with the same sub.slug + post3 = BlogPost(title='test3', sub=SubDocument(year=2010, slug='test')) + self.assertRaises(OperationError, post3.save) + + BlogPost.drop_collection() + + def test_unique_with_embedded_document_and_embedded_unique(self): + """Ensure that uniqueness constraints are applied to fields on + embedded documents. And work with unique_with as well. + """ + class SubDocument(EmbeddedDocument): + year = IntField(db_field='yr') + slug = StringField(unique=True) + + class BlogPost(Document): + title = StringField(unique_with='sub.year') + sub = EmbeddedDocumentField(SubDocument) + + BlogPost.drop_collection() + + post1 = BlogPost(title='test1', sub=SubDocument(year=2009, slug="test")) + post1.save() + + # sub.slug is different so won't raise exception + post2 = BlogPost(title='test2', sub=SubDocument(year=2010, slug='another-slug')) + post2.save() + + # Now there will be two docs with the same sub.slug + post3 = BlogPost(title='test3', sub=SubDocument(year=2010, slug='test')) + self.assertRaises(OperationError, post3.save) + + # Now there will be two docs with the same title and year + post3 = BlogPost(title='test1', sub=SubDocument(year=2009, slug='test-1')) + self.assertRaises(OperationError, post3.save) + + BlogPost.drop_collection() + def test_unique_and_indexes(self): """Ensure that 'unique' constraints aren't overridden by meta.indexes.