list_indexes and compare_indexes class methods + unit tests
This commit is contained in:
Stefan Wojcik
parent
c3a065dd33
commit
c2928d8a57
@ -20,6 +20,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
|
||||
|
||||
@ -529,14 +542,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']:
|
||||
@ -557,6 +562,73 @@ 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 []
|
||||
|
||||
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):
|
||||
"""A Dynamic Document class allowing flexible, expandable and uncontrolled
|
||||
|
||||
@ -85,6 +85,117 @@ 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_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):
|
||||
|
||||
class Vaccine(Document):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user