Merge pull request #88 from akolechkin/master

Fix for UnboundLocalError in composite index with pk field

mongoengine/queryset.py

@@ -501,7 +501,9 @@ class QuerySet(object):
         direction = None
 
         allow_inheritance = doc_cls._meta.get('allow_inheritance') != False
-        use_types = allow_inheritance
+
+        # If sparse - dont include types
+        use_types = allow_inheritance and not spec.get('sparse', False)
 
         for key in spec['fields']:
             # Get ASCENDING direction from +, DESCENDING from -, and GEO2D from *
@@ -518,6 +520,7 @@ class QuerySet(object):
             parts = key.split('.')
             if parts in (['pk'], ['id'], ['_id']):
                 key = '_id'
+                fields = []
             else:
                 fields = QuerySet._lookup_field(doc_cls, parts)
                 parts = [field if field == '_id' else field.db_field
@@ -525,10 +528,6 @@ class QuerySet(object):
                 key = '.'.join(parts)
             index_list.append((key, direction))
 
-            # If sparse - dont include types
-            if spec.get('sparse', False):
-                use_types = False
-
             # Check if a list field is being used, don't use _types if it is
             if use_types and not all(f._index_with_types for f in fields):
                 use_types = False
@@ -536,7 +535,7 @@ class QuerySet(object):
         # If _types is being used, prepend it to every specified index
         index_types = doc_cls._meta.get('index_types', True)
 
-        if (spec.get('types', index_types) and allow_inheritance and use_types
+        if (spec.get('types', index_types) and use_types
             and direction is not pymongo.GEO2D):
             index_list.insert(0, ('_types', 1))
 

tests/test_queryset.py

@@ -2524,6 +2524,24 @@ class QuerySetTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+    def test_types_index_with_pk(self):
+
+        class Comment(EmbeddedDocument):
+            comment_id = IntField(required=True)
+
+        try:
+            class BlogPost(Document):
+                comments = EmbeddedDocumentField(Comment)
+                meta = {'indexes': [{'fields': ['pk', 'comments.comment_id'],
+                    'unique': True}]}
+        except UnboundLocalError:
+            self.fail('Unbound local error at types index + pk definition')
+
+        info = BlogPost.objects._collection.index_information()
+        info = [value['key'] for key, value in info.iteritems()]
+        index_item = [(u'_types', 1), (u'_id', 1), (u'comments.comment_id', 1)]
+        self.assertTrue(index_item in info)
+
     def test_dict_with_custom_baseclass(self):
         """Ensure DictField working with custom base clases.
         """