Merge remote-tracking branch 'origin/pr/364' into 364

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

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):

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
         """