pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

list_indexes and compare_indexes class methods + unit tests

Browse Source
This commit is contained in:
Stefan Wojcik 2013-06-06 17:16:03 -07:00
parent c3a065dd33
commit c2928d8a57
2 changed files with 191 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
mongoengine/document.py
View File

@ -20,6 +20,19 @@ __all__ = ('Document', 'EmbeddedDocument', 'DynamicDocument',
'InvalidCollectionError', 'NotUniqueError', 'MapReduceDocument') '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): class InvalidCollectionError(Exception):
pass pass
@ -529,14 +542,6 @@ class Document(BaseDocument):
# an extra index on _cls, as mongodb will use the existing # an extra index on _cls, as mongodb will use the existing
# index to service queries against _cls # index to service queries against _cls
cls_indexed = False 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 # Ensure document-defined indexes are created
if cls._meta['index_specs']: if cls._meta['index_specs']:
@ -557,6 +562,73 @@ class Document(BaseDocument):
collection.ensure_index('_cls', background=background, collection.ensure_index('_cls', background=background,
**index_opts) **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 []
indexes = []
index_cls = cls._meta.get('index_cls', True)
# 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')
indexes.append(fields)
# add all of the indexes from the base classes
if go_up:
for base_cls in cls.__bases__:
for index in base_cls.list_indexes(go_up=True, go_down=False):
if index not in indexes:
indexes.append(index)
# add all of the indexes from subclasses
if go_down:
for subclass in cls.__subclasses__():
for index in subclass.list_indexes(go_up=False, go_down=True):
if index not in indexes:
indexes.append(index)
# finish up by appending _id, if needed
if go_up and go_down:
if [(u'_id', 1)] not in indexes:
indexes.append([(u'_id', 1)])
if (index_cls 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): class DynamicDocument(Document):
"""A Dynamic Document class allowing flexible, expandable and uncontrolled """A Dynamic Document class allowing flexible, expandable and uncontrolled

111
tests/document/class_methods.py
View File

@ -85,6 +85,117 @@ class ClassMethodsTest(unittest.TestCase):
self.assertEqual(self.Person._meta['delete_rules'], self.assertEqual(self.Person._meta['delete_rules'],
{(Job, 'employee'): NULLIFY}) {(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_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())
print BlogPost.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): def test_register_delete_rule_inherited(self):
class Vaccine(Document): class Vaccine(Document):