diff --git a/mongoengine/document.py b/mongoengine/document.py index 3d919e7d..e4507c1c 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -1,4 +1,4 @@ - +import warnings import pymongo import re @@ -17,6 +17,7 @@ from mongoengine.base import ( get_document ) from mongoengine.errors import InvalidQueryError, InvalidDocumentError +from mongoengine.python_support import IS_PYMONGO_3 from mongoengine.queryset import (OperationError, NotUniqueError, QuerySet, transform) from mongoengine.connection import get_db, DEFAULT_CONNECTION_NAME @@ -629,28 +630,51 @@ class Document(BaseDocument): db = cls._get_db() db.drop_collection(cls._get_collection_name()) + @classmethod + def create_index(cls, keys, background=False, **kwargs): + """Creates the given indexes if required. + + :param keys: 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 + :param background: Allows index creation in the background + """ + index_spec = cls._build_index_spec(keys) + index_spec = index_spec.copy() + fields = index_spec.pop('fields') + drop_dups = kwargs.get('drop_dups', False) + if IS_PYMONGO_3 and drop_dups: + msg = "drop_dups is deprecated and is removed when using PyMongo 3+." + warnings.warn(msg, DeprecationWarning) + elif not IS_PYMONGO_3: + index_spec['drop_dups'] = drop_dups + index_spec['background'] = background + index_spec.update(kwargs) + + if IS_PYMONGO_3: + return cls._get_collection().create_index(fields, **index_spec) + else: + return cls._get_collection().ensure_index(fields, **index_spec) + @classmethod def ensure_index(cls, key_or_list, drop_dups=False, background=False, **kwargs): - """Ensure that the given indexes are in place. + """Ensure that the given indexes are in place. Deprecated in favour + of create_index. :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 :param background: Allows index creation in the background - :param drop_dups: Was removed with MongoDB 3. The value will be - removed if PyMongo3+ is used + :param drop_dups: Was removed/ignored with MongoDB >2.7.5. The value + will be removed if PyMongo3+ is used """ - 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 - # TODO: raise warning if dropdups given and remove with PyMongo3+ - index_spec['background'] = background - index_spec.update(kwargs) - - # TODO: ensure_index is deprecated - return cls._get_collection().ensure_index(fields, **index_spec) + if IS_PYMONGO_3 and drop_dups: + msg = "drop_dups is deprecated and is removed when using PyMongo 3+." + warnings.warn(msg, DeprecationWarning) + elif not IS_PYMONGO_3: + kwargs.update({'drop_dups': drop_dups}) + return cls.create_index(key_or_list, background=background, **kwargs) @classmethod def ensure_indexes(cls): @@ -665,6 +689,9 @@ class Document(BaseDocument): 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) + if IS_PYMONGO_3 and drop_dups: + msg = "drop_dups is deprecated and is removed when using PyMongo 3+." + warnings.warn(msg, DeprecationWarning) collection = cls._get_collection() # 746: when connection is via mongos, the read preference is not necessarily an indication that @@ -693,9 +720,11 @@ class Document(BaseDocument): if 'cls' in opts: del opts['cls'] - # TODO: ensure_index is deprecated in PyMongo 3+ and drop_dups removed - collection.ensure_index(fields, background=background, - drop_dups=drop_dups, **opts) + if IS_PYMONGO_3: + collection.create_index(fields, background=background, **opts) + else: + collection.ensure_index(fields, background=background, + drop_dups=drop_dups, **opts) # If _cls is being used (for polymorphism), it needs an index, # only if another index doesn't begin with _cls @@ -707,9 +736,12 @@ class Document(BaseDocument): if 'cls' in index_opts: del index_opts['cls'] - # TODO: ensure_index is deprecated in PyMongo 3+ - collection.ensure_index('_cls', background=background, - **index_opts) + if IS_PYMONGO_3: + collection.create_index('_cls', background=background, + **index_opts) + else: + collection.ensure_index('_cls', background=background, + **index_opts) @classmethod def list_indexes(cls, go_up=True, go_down=True): diff --git a/tests/document/indexes.py b/tests/document/indexes.py index cecb8bf2..f9f68e57 100644 --- a/tests/document/indexes.py +++ b/tests/document/indexes.py @@ -298,19 +298,18 @@ class IndexesTest(unittest.TestCase): def test_explicit_geohaystack_index(self): """Ensure that geohaystack indexes work when created via meta[indexes] """ - raise SkipTest('GeoHaystack index creation seems broken on PyMongo' - 'side, as it requires a bucketSize parameter.') + raise SkipTest('GeoHaystack index creation is not supported for now' + 'from meta, as it requires a bucketSize parameter.') class Place(Document): location = DictField() + name = StringField() meta = { - 'allow_inheritance': True, 'indexes': [ - ')location.point', + (')location.point', 'name') ] } - - self.assertEqual([{'fields': [('location.point', 'geoHaystack')]}], + self.assertEqual([{'fields': [('location.point', 'geoHaystack'), ('name', 1)]}], Place._meta['index_specs']) Place.ensure_indexes() @@ -318,6 +317,18 @@ class IndexesTest(unittest.TestCase): info = [value['key'] for key, value in info.iteritems()] self.assertTrue([('location.point', 'geoHaystack')] in info) + def test_create_geohaystack_index(self): + """Ensure that geohaystack indexes can be created + """ + class Place(Document): + location = DictField() + name = StringField() + + Place.ensure_index({'fields': (')location.point', 'name')}, bucketSize=10) + info = Place._get_collection().index_information() + info = [value['key'] for key, value in info.iteritems()] + self.assertTrue([('location.point', 'geoHaystack'), ('name', 1)] in info) + def test_dictionary_indexes(self): """Ensure that indexes are used when meta[indexes] contains dictionaries instead of lists.