Updated index creation now tied to Document class ((MongoEngine/mongoengine#102)

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 
 Changes in 0.8
 ==============
+- Updated index creation now tied to Document class ((MongoEngine/mongoengine#102)
 - Added none() to queryset (MongoEngine/mongoengine#127)
 - Updated SequenceFields to allow post processing of the calculated counter value (MongoEngine/mongoengine#141)
 - Added clean method to documents for pre validation data cleaning (MongoEngine/mongoengine#60)

docs/upgrade.rst

@@ -40,7 +40,7 @@ inherited classes like so:
         collection.drop_index(index)
 
     # 5. Recreate indexes
-    Animal.objects._ensure_indexes()
+    Animal.ensure_indexes()
 
 
 Document Definition
@@ -56,6 +56,13 @@ you will need to declare :attr:`allow_inheritance` in the meta data like so:
 
         meta = {'allow_inheritance': True}
 
+Indexes
+-------
+
+Index methods are no longer tied to querysets but rather to the document class.
+Although `QuerySet._ensure_indexes` and `QuerySet.ensure_index` still exist.
+They should be replaced with :func:`~mongoengine.Document.ensure_indexes` /
+:func:`~mongoengine.Document.ensure_index`.
 
 SequenceFields
 --------------

mongoengine/document.py

@@ -7,7 +7,7 @@ from bson.dbref import DBRef
 from mongoengine import signals, queryset
 
 from base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument,
-                  BaseDict, BaseList)
+                  BaseDict, BaseList, ALLOW_INHERITANCE)
 from queryset import OperationError, NotUniqueError
 from connection import get_db, DEFAULT_CONNECTION_NAME
 
@@ -163,6 +163,8 @@ class Document(BaseDocument):
                     )
             else:
                 cls._collection = db[collection_name]
+            if cls._meta.get('auto_create_index', True):
+                cls.ensure_indexes()
         return cls._collection
 
     def save(self, safe=True, force_insert=False, validate=True, clean=True,
@@ -418,9 +420,86 @@ class Document(BaseDocument):
         """Drops the entire collection associated with this
         :class:`~mongoengine.Document` type from the database.
         """
+        cls._collection = None
         db = cls._get_db()
         db.drop_collection(cls._get_collection_name())
-        queryset.QuerySet._reset_already_indexed(cls)
+
+    @classmethod
+    def ensure_index(cls, key_or_list, drop_dups=False, background=False,
+        **kwargs):
+        """Ensure that the given indexes are in place.
+
+        :param key_or_list: a single index key or a list of index keys (to
+            construct a multi-field index); keys may be prefixed with a **+**
+            or a **-** to determine the index ordering
+        """
+        index_spec = cls._build_index_spec(key_or_list)
+        index_spec = index_spec.copy()
+        fields = index_spec.pop('fields')
+        index_spec['drop_dups'] = drop_dups
+        index_spec['background'] = background
+        index_spec.update(kwargs)
+
+        return cls._get_collection().ensure_index(fields, **index_spec)
+
+    @classmethod
+    def ensure_indexes(cls):
+        """Checks the document meta data and ensures all the indexes exist.
+
+        .. note:: You can disable automatic index creation by setting
+                  `auto_create_index` to False in the documents meta data
+        """
+        background = cls._meta.get('index_background', False)
+        drop_dups = cls._meta.get('index_drop_dups', False)
+        index_opts = cls._meta.get('index_opts') or {}
+        index_cls = cls._meta.get('index_cls', True)
+
+        collection = cls._get_collection()
+
+        # determine if an index which we are creating includes
+        # _cls as its first field; if so, we can avoid creating
+        # 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 indexes created by uniqueness constraints
+        for index in cls._meta['unique_indexes']:
+            cls_indexed = cls_indexed or includes_cls(index)
+            collection.ensure_index(index, unique=True, background=background,
+                                           drop_dups=drop_dups, **index_opts)
+
+        # Ensure document-defined indexes are created
+        if cls._meta['index_specs']:
+            index_spec = cls._meta['index_specs']
+            for spec in index_spec:
+                spec = spec.copy()
+                fields = spec.pop('fields')
+                cls_indexed = cls_indexed or includes_cls(fields)
+                opts = index_opts.copy()
+                opts.update(spec)
+                collection.ensure_index(fields, background=background, **opts)
+
+        # If _cls is being used (for polymorphism), it needs an index,
+        # only if another index doesn't begin with _cls
+        if (index_cls and not cls_indexed and
+            cls._meta.get('allow_inheritance', ALLOW_INHERITANCE) == True):
+            collection.ensure_index('_cls', background=background,
+                                    **index_opts)
+
+        # Add geo indicies
+        for field in cls._geo_indices():
+            index_spec = [(field.db_field, pymongo.GEO2D)]
+            collection.ensure_index(index_spec, background=background,
+                                    **index_opts)
 
 
 class DynamicDocument(Document):

mongoengine/queryset/queryset.py

@@ -1,11 +1,12 @@
-import pprint
-import re
 import copy
 import itertools
 import operator
+import pprint
+import re
+import warnings
 
-import pymongo
 from bson.code import Code
+import pymongo
 from pymongo.common import validate_read_preference
 
 from mongoengine import signals
@@ -37,8 +38,6 @@ class QuerySet(object):
     """A set of results returned from a query. Wraps a MongoDB cursor,
     providing :class:`~mongoengine.Document` objects as the results.
     """
-
-    __already_indexed = set()
     __dereference = False
 
     def __init__(self, document, collection):
@@ -95,24 +94,6 @@ class QuerySet(object):
                 self._mongo_query.update(self._initial_query)
         return self._mongo_query
 
-    def ensure_index(self, key_or_list, drop_dups=False, background=False,
-        **kwargs):
-        """Ensure that the given indexes are in place.
-
-        :param key_or_list: a single index key or a list of index keys (to
-            construct a multi-field index); keys may be prefixed with a **+**
-            or a **-** to determine the index ordering
-        """
-        index_spec = self._document._build_index_spec(key_or_list)
-        index_spec = index_spec.copy()
-        fields = index_spec.pop('fields')
-        index_spec['drop_dups'] = drop_dups
-        index_spec['background'] = background
-        index_spec.update(kwargs)
-
-        self._collection.ensure_index(fields, **index_spec)
-        return self
-
     def __call__(self, q_obj=None, class_check=True, slave_okay=False,
                  read_preference=None, **query):
         """Filter the selected documents by calling the
@@ -150,87 +131,26 @@ class QuerySet(object):
         """Returns all documents."""
         return self.__call__()
 
+    def ensure_index(self, **kwargs):
+        """Deprecated use :func:`~Document.ensure_index`"""
+        msg = ("Doc.objects()._ensure_index() is deprecated. "
+              "Use Doc.ensure_index() instead.")
+        warnings.warn(msg, DeprecationWarning)
+        self._document.__class__.ensure_index(**kwargs)
+        return self
+
     def _ensure_indexes(self):
-        """Checks the document meta data and ensures all the indexes exist.
+        """Deprecated use :func:`~Document.ensure_indexes`"""
-
+        msg = ("Doc.objects()._ensure_indexes() is deprecated. "
-        .. note:: You can disable automatic index creation by setting
+              "Use Doc.ensure_indexes() instead.")
-                  `auto_create_index` to False in the documents meta data
+        warnings.warn(msg, DeprecationWarning)
-        """
+        self._document.__class__.ensure_indexes()
-        background = self._document._meta.get('index_background', False)
-        drop_dups = self._document._meta.get('index_drop_dups', False)
-        index_opts = self._document._meta.get('index_opts') or {}
-        index_cls = self._document._meta.get('index_cls', True)
-
-        # determine if an index which we are creating includes
-        # _cls as its first field; if so, we can avoid creating
-        # 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 indexes created by uniqueness constraints
-        for index in self._document._meta['unique_indexes']:
-            cls_indexed = cls_indexed or includes_cls(index)
-            self._collection.ensure_index(index, unique=True,
-                background=background, drop_dups=drop_dups, **index_opts)
-
-        # Ensure document-defined indexes are created
-        if self._document._meta['index_specs']:
-            index_spec = self._document._meta['index_specs']
-            for spec in index_spec:
-                spec = spec.copy()
-                fields = spec.pop('fields')
-                cls_indexed = cls_indexed or includes_cls(fields)
-                opts = index_opts.copy()
-                opts.update(spec)
-                self._collection.ensure_index(fields,
-                    background=background, **opts)
-
-        # If _cls is being used (for polymorphism), it needs an index,
-        # only if another index doesn't begin with _cls
-        if index_cls and '_cls' in self._query and not cls_indexed:
-            self._collection.ensure_index('_cls',
-                background=background, **index_opts)
-
-        # Add geo indicies
-        for field in self._document._geo_indices():
-            index_spec = [(field.db_field, pymongo.GEO2D)]
-            self._collection.ensure_index(index_spec,
-                background=background, **index_opts)
-
-    @classmethod
-    def _reset_already_indexed(cls, document=None):
-        """Helper to reset already indexed, can be useful for testing purposes
-        """
-        if document:
-            cls.__already_indexed.discard(document)
-        cls.__already_indexed.clear()
 
     @property
     def _collection(self):
         """Property that returns the collection object. This allows us to
         perform operations only if the collection is accessed.
         """
-        if self._document not in QuerySet.__already_indexed:
-            # Ensure collection exists
-            db = self._document._get_db()
-            if self._collection_obj.name not in db.collection_names():
-                self._document._collection = None
-                self._collection_obj = self._document._get_collection()
-
-            QuerySet.__already_indexed.add(self._document)
-
-            if self._document._meta.get('auto_create_index', True):
-                self._ensure_indexes()
-
         return self._collection_obj
 
     @property

tests/document/indexes.py

@@ -80,7 +80,7 @@ class InstanceTest(unittest.TestCase):
                                       ('addDate', -1)]}]
         self.assertEqual(expected_specs, BlogPost._meta['index_specs'])
 
-        BlogPost.objects._ensure_indexes()
+        BlogPost.ensure_indexes()
         info = BlogPost.objects._collection.index_information()
         # _id, '-date', 'tags', ('cat', 'date')
         # NB: there is no index on _cls by itself, since
@@ -100,7 +100,7 @@ class InstanceTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
-        ExtendedBlogPost.objects._ensure_indexes()
+        ExtendedBlogPost.ensure_indexes()
         info = ExtendedBlogPost.objects._collection.index_information()
         info = [value['key'] for key, value in info.iteritems()]
         for expected in expected_specs:
@@ -141,7 +141,7 @@ class InstanceTest(unittest.TestCase):
                          [{'fields': [('keywords', 1)]}])
 
         # Force index creation
-        MyDoc.objects._ensure_indexes()
+        MyDoc.ensure_indexes()
 
         self.assertEqual(MyDoc._meta['index_specs'],
                         [{'fields': [('keywords', 1)]}])
@@ -189,7 +189,7 @@ class InstanceTest(unittest.TestCase):
         self.assertEqual([{'fields': [('location.point', '2d')]}],
                         Place._meta['index_specs'])
 
-        Place.objects()._ensure_indexes()
+        Place.ensure_indexes()
         info = Place._get_collection().index_information()
         info = [value['key'] for key, value in info.iteritems()]
         self.assertTrue([('location.point', '2d')] in info)
@@ -335,7 +335,7 @@ class InstanceTest(unittest.TestCase):
             recursive_obj = EmbeddedDocumentField(RecursiveObject)
             meta = {'allow_inheritance': True}
 
-        RecursiveDocument.objects._ensure_indexes()
+        RecursiveDocument.ensure_indexes()
         info = RecursiveDocument._get_collection().index_information()
         self.assertEqual(info.keys(), ['_id_', '_cls_1'])
 

tests/migration/test_convert_to_new_inheritance_model.py

@@ -48,4 +48,4 @@ class ConvertToNewInheritanceModel(unittest.TestCase):
             collection.drop_index(index)
 
         # 5. Recreate indexes
-        Animal.objects._ensure_indexes()
+        Animal.ensure_indexes()

tests/migration/turn_off_inheritance.py

@@ -59,4 +59,4 @@ class TurnOffInheritanceTest(unittest.TestCase):
             collection.drop_index(index)
 
         # 5. Recreate indexes
-        Animal.objects._ensure_indexes()
+        Animal.ensure_indexes()

tests/test_queryset.py

@@ -3078,7 +3078,6 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual([1, 2, 3], numbers)
         Number.drop_collection()
 
-
     def test_ensure_index(self):
         """Ensure that manual creation of indexes works.
         """
@@ -3086,7 +3085,7 @@ class QuerySetTest(unittest.TestCase):
             message = StringField()
             meta = {'allow_inheritance': True}
 
-        Comment.objects.ensure_index('message')
+        Comment.ensure_index('message')
 
         info = Comment.objects._collection.index_information()
         info = [(value['key'],