diff --git a/mongoengine/document.py b/mongoengine/document.py index 8b152d5e..a61ed079 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -21,6 +21,19 @@ __all__ = ('Document', 'EmbeddedDocument', 'DynamicDocument', 'InvalidCollectionError', 'NotUniqueError', 'MapReduceDocument') +def includes_cls(fields): + """ Helper function used for ensuring and comparing indexes + """ + + first_field = None + if len(fields): + if isinstance(fields[0], basestring): + first_field = fields[0] + elif isinstance(fields[0], (list, tuple)) and len(fields[0]): + first_field = fields[0][0] + return first_field == '_cls' + + class InvalidCollectionError(Exception): pass @@ -536,14 +549,6 @@ class Document(BaseDocument): # an extra index on _cls, as mongodb will use the existing # index to service queries against _cls cls_indexed = False - def includes_cls(fields): - first_field = None - if len(fields): - if isinstance(fields[0], basestring): - first_field = fields[0] - elif isinstance(fields[0], (list, tuple)) and len(fields[0]): - first_field = fields[0][0] - return first_field == '_cls' # Ensure document-defined indexes are created if cls._meta['index_specs']: @@ -564,6 +569,90 @@ class Document(BaseDocument): collection.ensure_index('_cls', background=background, **index_opts) + @classmethod + def list_indexes(cls, go_up=True, go_down=True): + """ Lists all of the indexes that should be created for given + collection. It includes all the indexes from super- and sub-classes. + """ + + if cls._meta.get('abstract'): + return [] + + # get all the base classes, subclasses and sieblings + classes = [] + def get_classes(cls): + + if (cls not in classes and + isinstance(cls, TopLevelDocumentMetaclass)): + classes.append(cls) + + for base_cls in cls.__bases__: + if (isinstance(base_cls, TopLevelDocumentMetaclass) and + base_cls != Document and + not base_cls._meta.get('abstract') and + base_cls._get_collection().full_name == cls._get_collection().full_name and + base_cls not in classes): + classes.append(base_cls) + get_classes(base_cls) + for subclass in cls.__subclasses__(): + if (isinstance(base_cls, TopLevelDocumentMetaclass) and + subclass._get_collection().full_name == cls._get_collection().full_name and + subclass not in classes): + classes.append(subclass) + get_classes(subclass) + + get_classes(cls) + + # get the indexes spec for all of the gathered classes + def get_indexes_spec(cls): + indexes = [] + + if cls._meta['index_specs']: + index_spec = cls._meta['index_specs'] + for spec in index_spec: + spec = spec.copy() + fields = spec.pop('fields') + indexes.append(fields) + return indexes + + indexes = [] + for cls in classes: + for index in get_indexes_spec(cls): + if index not in indexes: + indexes.append(index) + + # finish up by appending { '_id': 1 } and { '_cls': 1 }, if needed + if [(u'_id', 1)] not in indexes: + indexes.append([(u'_id', 1)]) + if (cls._meta.get('index_cls', True) and + cls._meta.get('allow_inheritance', ALLOW_INHERITANCE) is True): + indexes.append([(u'_cls', 1)]) + + return indexes + + @classmethod + def compare_indexes(cls): + """ Compares the indexes defined in MongoEngine with the ones existing + in the database. Returns any missing/extra indexes. + """ + + required = cls.list_indexes() + existing = [info['key'] for info in cls._get_collection().index_information().values()] + missing = [index for index in required if index not in existing] + extra = [index for index in existing if index not in required] + + # if { _cls: 1 } is missing, make sure it's *really* necessary + if [(u'_cls', 1)] in missing: + cls_obsolete = False + for index in existing: + if includes_cls(index) and index not in extra: + cls_obsolete = True + break + if cls_obsolete: + missing.remove([(u'_cls', 1)]) + + return {'missing': missing, 'extra': extra} + class DynamicDocument(Document): """A Dynamic Document class allowing flexible, expandable and uncontrolled diff --git a/tests/document/class_methods.py b/tests/document/class_methods.py index b2c72838..52e3794c 100644 --- a/tests/document/class_methods.py +++ b/tests/document/class_methods.py @@ -85,6 +85,153 @@ class ClassMethodsTest(unittest.TestCase): self.assertEqual(self.Person._meta['delete_rules'], {(Job, 'employee'): NULLIFY}) + def test_compare_indexes(self): + """ Ensure that the indexes are properly created and that + compare_indexes identifies the missing/extra indexes + """ + + class BlogPost(Document): + author = StringField() + title = StringField() + description = StringField() + tags = StringField() + + meta = { + 'indexes': [('author', 'title')] + } + + BlogPost.drop_collection() + + BlogPost.ensure_indexes() + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [] }) + + BlogPost.ensure_index(['author', 'description']) + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [[('author', 1), ('description', 1)]] }) + + BlogPost._get_collection().drop_index('author_1_description_1') + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [] }) + + BlogPost._get_collection().drop_index('author_1_title_1') + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [[('author', 1), ('title', 1)]], 'extra': [] }) + + def test_compare_indexes_inheritance(self): + """ Ensure that the indexes are properly created and that + compare_indexes identifies the missing/extra indexes for subclassed + documents (_cls included) + """ + + class BlogPost(Document): + author = StringField() + title = StringField() + description = StringField() + + meta = { + 'allow_inheritance': True + } + + class BlogPostWithTags(BlogPost): + tags = StringField() + tag_list = ListField(StringField()) + + meta = { + 'indexes': [('author', 'tags')] + } + + BlogPost.drop_collection() + + BlogPost.ensure_indexes() + BlogPostWithTags.ensure_indexes() + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [] }) + + BlogPostWithTags.ensure_index(['author', 'tag_list']) + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [[('_cls', 1), ('author', 1), ('tag_list', 1)]] }) + + BlogPostWithTags._get_collection().drop_index('_cls_1_author_1_tag_list_1') + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [] }) + + BlogPostWithTags._get_collection().drop_index('_cls_1_author_1_tags_1') + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [[('_cls', 1), ('author', 1), ('tags', 1)]], 'extra': [] }) + + def test_compare_indexes_multiple_subclasses(self): + """ Ensure that compare_indexes behaves correctly if called from a + class, which base class has multiple subclasses + """ + + class BlogPost(Document): + author = StringField() + title = StringField() + description = StringField() + + meta = { + 'allow_inheritance': True + } + + class BlogPostWithTags(BlogPost): + tags = StringField() + tag_list = ListField(StringField()) + + meta = { + 'indexes': [('author', 'tags')] + } + + class BlogPostWithCustomField(BlogPost): + custom = DictField() + + meta = { + 'indexes': [('author', 'custom')] + } + + BlogPost.ensure_indexes() + BlogPostWithTags.ensure_indexes() + BlogPostWithCustomField.ensure_indexes() + + self.assertEqual(BlogPost.compare_indexes(), { 'missing': [], 'extra': [] }) + self.assertEqual(BlogPostWithTags.compare_indexes(), { 'missing': [], 'extra': [] }) + self.assertEqual(BlogPostWithCustomField.compare_indexes(), { 'missing': [], 'extra': [] }) + + def test_list_indexes_inheritance(self): + """ ensure that all of the indexes are listed regardless of the super- + or sub-class that we call it from + """ + + class BlogPost(Document): + author = StringField() + title = StringField() + description = StringField() + + meta = { + 'allow_inheritance': True + } + + class BlogPostWithTags(BlogPost): + tags = StringField() + + meta = { + 'indexes': [('author', 'tags')] + } + + class BlogPostWithTagsAndExtraText(BlogPostWithTags): + extra_text = StringField() + + meta = { + 'indexes': [('author', 'tags', 'extra_text')] + } + + BlogPost.drop_collection() + + BlogPost.ensure_indexes() + BlogPostWithTags.ensure_indexes() + BlogPostWithTagsAndExtraText.ensure_indexes() + + self.assertEqual(BlogPost.list_indexes(), + BlogPostWithTags.list_indexes()) + self.assertEqual(BlogPost.list_indexes(), + BlogPostWithTagsAndExtraText.list_indexes()) + self.assertEqual(BlogPost.list_indexes(), + [[('_cls', 1), ('author', 1), ('tags', 1)], + [('_cls', 1), ('author', 1), ('tags', 1), ('extra_text', 1)], + [(u'_id', 1)], [('_cls', 1)]]) + def test_register_delete_rule_inherited(self): class Vaccine(Document): diff --git a/tests/document/inheritance.py b/tests/document/inheritance.py index f0116311..28490c95 100644 --- a/tests/document/inheritance.py +++ b/tests/document/inheritance.py @@ -189,6 +189,41 @@ class InheritanceTest(unittest.TestCase): self.assertEqual(Employee._get_collection_name(), Person._get_collection_name()) + def test_indexes_and_multiple_inheritance(self): + """ Ensure that all of the indexes are created for a document with + multiple inheritance. + """ + + class A(Document): + a = StringField() + + meta = { + 'allow_inheritance': True, + 'indexes': ['a'] + } + + class B(Document): + b = StringField() + + meta = { + 'allow_inheritance': True, + 'indexes': ['b'] + } + + class C(A, B): + pass + + A.drop_collection() + B.drop_collection() + C.drop_collection() + + C.ensure_indexes() + + self.assertEqual( + [idx['key'] for idx in C._get_collection().index_information().values()], + [[(u'_cls', 1), (u'b', 1)], [(u'_id', 1)], [(u'_cls', 1), (u'a', 1)]] + ) + def test_polymorphic_queries(self): """Ensure that the correct subclasses are returned from a query """