Allow dot notation indices of embedded documents

django_mongodb_engine/creation.py

@@ -2,12 +2,13 @@
 from djangotoolbox.db.base import NonrelDatabaseCreation
 from .utils import make_index_list
 
+
 class DatabaseCreation(NonrelDatabaseCreation):
     data_types = dict(NonrelDatabaseCreation.data_types, **{
         'AutoField': 'objectid',
         'ForeignKey': 'objectid',
         'OneToOneField': 'objectid',
-        'RelatedAutoField' : 'objectid',
+        'RelatedAutoField': 'objectid',
         'DecimalField': 'float',
     })
 
@@ -22,7 +23,8 @@ def sql_indexes_for_model(self, model, termstyle):
 
         def ensure_index(*args, **kwargs):
             if ensure_index.first_index:
-                print "Installing indices for %s.%s model" % (meta.app_label, meta.object_name)
+                print "Installing indices for %s.%s model" % (meta.app_label,
+                                                              meta.object_name)
                 ensure_index.first_index = False
             return collection.ensure_index(*args, **kwargs)
         ensure_index.first_index = True
@@ -34,6 +36,8 @@ def ensure_index(*args, **kwargs):
             self._handle_oldstyle_indexes(ensure_index, meta)
 
     def _handle_newstyle_indexes(self, ensure_index, meta, indexes):
+        from djangotoolbox.fields import AbstractIterableField, EmbeddedModelField
+
         # Django indexes
         for field in meta.local_fields:
             if not (field.unique or field.db_index):
@@ -44,21 +48,35 @@ def _handle_newstyle_indexes(self, ensure_index, meta, indexes):
 
         # Django unique_together indexes
         indexes = list(indexes)
+
         for fields in getattr(meta, 'unique_together', []):
             assert isinstance(fields, (list, tuple))
             indexes.append({'fields': make_index_list(fields), 'unique': True})
 
+        def get_column_name(field):
+            opts = meta
+            parts = field.split('.')
+            for i, part in enumerate(parts):
+                field = opts.get_field(part)
+                parts[i] = field.column
+                if isinstance(field, AbstractIterableField):
+                    field = field.item_field
+                if isinstance(field, EmbeddedModelField):
+                    opts = field.embedded_model._meta
+                else:
+                    break
+            return '.'.join(parts)
+
         for index in indexes:
             if isinstance(index, dict):
                 kwargs = index.copy()
                 fields = kwargs.pop('fields')
             else:
                 fields, kwargs = index, {}
-            fields = [(meta.get_field(name).column, direction)
+            fields = [(get_column_name(name), direction)
                       for name, direction in make_index_list(fields)]
             ensure_index(fields, **kwargs)
 
-
     def _handle_oldstyle_indexes(self, ensure_index, meta):
         from warnings import warn
         warn("'descending_indexes', 'sparse_indexes' and 'index_together' are"
@@ -69,7 +87,8 @@ def _handle_oldstyle_indexes(self, ensure_index, meta):
 
         # Lets normalize the sparse_index values changing [], set() to ()
         for idx in set(getattr(meta, 'sparse_indexes', ())):
-            sparse_indexes.append(isinstance(idx, (tuple, set, list)) and tuple(idx) or idx )
+            sparse_indexes.append(isinstance(idx, (tuple, set, list))
+                                  and tuple(idx) or idx)
 
         # Ordinary indexes
         for field in meta.local_fields:
@@ -82,7 +101,6 @@ def _handle_oldstyle_indexes(self, ensure_index, meta):
             ensure_index(column, unique=field.unique,
                          sparse=field.name in sparse_indexes)
 
-
         def create_compound_indexes(indexes, **kwargs):
             # indexes: (field1, field2, ...)
             if not indexes:
@@ -112,7 +130,8 @@ def create_compound_indexes(indexes, **kwargs):
         return []
 
     def sql_create_model(self, model, *unused):
-        """ Creates the collection for model. Mostly used for capped collections. """
+        """ Creates the collection for model.
+            Mostly used for capped collections. """
         kwargs = {}
         for option in ('capped', 'collection_max', 'collection_size'):
             x = getattr(model._meta, option, None)
@@ -147,14 +166,15 @@ def create_test_db(self, verbosity=1, autoclobber=False):
         self._drop_database(test_database_name)
 
         from django.core.management import call_command
-        call_command('syncdb', verbosity=max(verbosity-1, 0),
+        call_command('syncdb', verbosity=max(verbosity - 1, 0),
                      interactive=False, database=self.connection.alias)
 
         return test_database_name
 
     def destroy_test_db(self, old_database_name, verbosity=1):
         if verbosity >= 1:
-            print "Destroying test database for alias '%s'..." % self.connection.alias
+            print ("Destroying test database for alias '%s'..."
+                   % self.connection.alias)
         test_database_name = self.connection.settings_dict['NAME']
         self._drop_database(test_database_name)
         self.connection.settings_dict['NAME'] = old_database_name

tests/mongodb/models.py

@@ -1,5 +1,5 @@
 from django.db import models
-from djangotoolbox.fields import RawField, ListField, EmbeddedModelField
+from djangotoolbox.fields import RawField, ListField, EmbeddedModelField, DictField
 from django_mongodb_engine.fields import GridFSField, GridFSString
 from query.models import Post
 
@@ -40,26 +40,40 @@ class IndexTestModel2(models.Model):
     class MongoMeta:
         index_together = ['a', ('b', -1)]
 
+class CustomColumnEmbeddedModel(models.Model):
+    a = models.IntegerField(db_column='a2')
+
 class NewStyleIndexesTestModel(models.Model):
-    a = models.IntegerField()
-    b = models.IntegerField(db_column='b2')
-    c = models.IntegerField(db_index=True)
-    d = models.IntegerField()
-    e = models.IntegerField()
-    f = models.IntegerField(unique=True)
+    f1 = models.IntegerField()
+    f2 = models.IntegerField()
+    f3 = models.IntegerField()
+
+    db_index = models.IntegerField(db_index=True)
+    unique = models.IntegerField(unique=True)
+    custom_column = models.IntegerField(db_column='custom')
     geo = models.IntegerField()
-    geo2 = models.IntegerField(db_column='geo')
+    geo_custom_column = models.IntegerField(db_column='geo')
+
+    dict1 = DictField()
+    dict_custom_column = DictField(db_column='dict_custom')
+    embedded = EmbeddedModelField(CustomColumnEmbeddedModel)
+    embedded_list = ListField(EmbeddedModelField(CustomColumnEmbeddedModel))
 
     class Meta:
-        unique_together = [('a', 'b'), ('a', 'd')]
+        unique_together = [('f2', 'custom_column'), ('f2', 'f3')]
 
     class MongoMeta:
         indexes = [
-            [('e', -1)],
-            {'fields': 'a', 'sparse': True},
-            {'fields': [('b', -1), 'd']},
+            [('f1', -1)],
+            {'fields': 'f2', 'sparse': True},
+            {'fields': [('custom_column', -1), 'f3']},
             [('geo', '2d')],
-            {'fields': [('geo2', '2d'), 'a'], 'min': 42, 'max': 21}
+            {'fields': [('geo_custom_column', '2d'), 'f2'], 'min': 42, 'max': 21},
+            {'fields': [('dict1.foo', 1)]},
+            {'fields': [('dict_custom_column.foo', 1)]},
+            {'fields' :[('embedded.a', 1)]},
+            {'fields' :[('embedded_list.a', 1)]},
+
         ]
 
 class GridFSFieldTestModel(models.Model):

tests/mongodb/tests.py

@@ -337,15 +337,19 @@ def assertHaveIndex(self, key, **properties):
         self.assertEqual(info[index_name], dict(default_properties, **properties))
 
     def test_indexes(self):
-        self.assertHaveIndex([('c', 1)])
-        self.assertHaveIndex([('f', 1)], unique=True)
-        self.assertHaveIndex([('a', 1), ('b2', 1)], unique=True)
-        self.assertHaveIndex([('a', 1), ('d', 1)], unique=True)
-        self.assertHaveIndex([('e', -1)])
-        self.assertHaveIndex([('a', 1)], sparse=True)
-        self.assertHaveIndex([('b2', -1), ('d', 1)])
+        self.assertHaveIndex([('db_index', 1)])
+        self.assertHaveIndex([('unique', 1)], unique=True)
+        self.assertHaveIndex([('f2', 1), ('custom', 1)], unique=True)
+        self.assertHaveIndex([('f2', 1), ('f3', 1)], unique=True)
+        self.assertHaveIndex([('f1', -1)])
+        self.assertHaveIndex([('f2', 1)], sparse=True)
+        self.assertHaveIndex([('custom', -1), ('f3', 1)])
         self.assertHaveIndex([('geo', '2d')])
-        self.assertHaveIndex([('geo', '2d'), ('a', 1)], min=42, max=21)
+        self.assertHaveIndex([('geo', '2d'), ('f2', 1)], min=42, max=21)
+        self.assertHaveIndex([('dict1.foo', 1)])
+        self.assertHaveIndex([('dict_custom.foo', 1)])
+        self.assertHaveIndex([('embedded.a2', 1)])
+        self.assertHaveIndex([('embedded_list.a2', 1)])
 
 class GridFSFieldTests(TestCase):
     def tearDown(self):