From c2928d8a57460e60d4b26d5f25f965d40eb4e1a6 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Thu, 6 Jun 2013 17:16:03 -0700 Subject: [PATCH 1/4] list_indexes and compare_indexes class methods + unit tests --- mongoengine/document.py | 88 ++++++++++++++++++++++--- tests/document/class_methods.py | 111 ++++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+), 8 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 585fcf76..83f60ee1 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -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 diff --git a/tests/document/class_methods.py b/tests/document/class_methods.py index b2c72838..6bd2e3c0 100644 --- a/tests/document/class_methods.py +++ b/tests/document/class_methods.py @@ -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): From 305540f0fd1731e9dd232cb995ab457f9a0fc6f4 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Thu, 6 Jun 2013 17:21:27 -0700 Subject: [PATCH 2/4] better comment --- mongoengine/document.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 83f60ee1..95ad0dc6 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -596,7 +596,7 @@ class Document(BaseDocument): if index not in indexes: indexes.append(index) - # finish up by appending _id, if needed + # finish up by appending { '_id': 1 } and { '_cls': 1 }, if needed if go_up and go_down: if [(u'_id', 1)] not in indexes: indexes.append([(u'_id', 1)]) From a2457df45e0fcc0077dde4f7fe09891e44ad0635 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Thu, 6 Jun 2013 19:14:21 -0700 Subject: [PATCH 3/4] make sure to only search for indexes in base classes inheriting from TopLevelDocumentMetaclass --- mongoengine/document.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 95ad0dc6..3b6df4f8 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -585,9 +585,10 @@ class Document(BaseDocument): # 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) + if isinstance(base_cls, TopLevelDocumentMetaclass): + 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: From ba7101ff92588185c4fa3355351d28c57ede4f78 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Thu, 6 Jun 2013 22:22:43 -0700 Subject: [PATCH 4/4] list_indexes support for multiple inheritance --- mongoengine/document.py | 70 ++++++++++++++++++++------------- tests/document/class_methods.py | 38 +++++++++++++++++- tests/document/inheritance.py | 35 +++++++++++++++++ 3 files changed, 115 insertions(+), 28 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 3b6df4f8..6345e6da 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -571,39 +571,55 @@ class Document(BaseDocument): if cls._meta.get('abstract'): return [] - indexes = [] - index_cls = cls._meta.get('index_cls', True) + # get all the base classes, subclasses and sieblings + classes = [] + def get_classes(cls): - # 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) + if (cls not in classes and + isinstance(cls, TopLevelDocumentMetaclass)): + classes.append(cls) - # add all of the indexes from the base classes - if go_up: for base_cls in cls.__bases__: - if isinstance(base_cls, TopLevelDocumentMetaclass): - 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: + 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__(): - for index in subclass.list_indexes(go_up=False, go_down=True): - if index not in indexes: - indexes.append(index) + 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 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)]) + 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 diff --git a/tests/document/class_methods.py b/tests/document/class_methods.py index 6bd2e3c0..52e3794c 100644 --- a/tests/document/class_methods.py +++ b/tests/document/class_methods.py @@ -152,6 +152,43 @@ class ClassMethodsTest(unittest.TestCase): 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 @@ -190,7 +227,6 @@ class ClassMethodsTest(unittest.TestCase): 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)], diff --git a/tests/document/inheritance.py b/tests/document/inheritance.py index f0116311..28490c95 100644 --- a/tests/document/inheritance.py +++ b/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 """