Allow shard key to be in an embedded document (#551)

Author: gmacon

mongoengine/document.py

@@ -341,8 +341,12 @@ def save(self, force_insert=False, validate=True, clean=True,
                 select_dict['_id'] = object_id
                 shard_key = self.__class__._meta.get('shard_key', tuple())
                 for k in shard_key:
-                    actual_key = self._db_field_map.get(k, k)
-                    select_dict[actual_key] = doc[actual_key]
+                    path = self._lookup_field(k.split('.'))
+                    actual_key = [p.db_field for p in path]
+                    val = doc
+                    for ak in actual_key:
+                        val = val[ak]
+                    select_dict['.'.join(actual_key)] = val
 
                 def is_new_object(last_error):
                     if last_error is not None:
@@ -444,7 +448,12 @@ def _object_key(self):
         select_dict = {'pk': self.pk}
         shard_key = self.__class__._meta.get('shard_key', tuple())
         for k in shard_key:
-            select_dict[k] = getattr(self, k)
+            path = self._lookup_field(k.split('.'))
+            actual_key = [p.db_field for p in path]
+            val = self
+            for ak in actual_key:
+                val = getattr(val, ak)
+            select_dict['__'.join(actual_key)] = val
         return select_dict
 
     def update(self, **kwargs):

tests/document/instance.py

@@ -484,6 +484,20 @@ class Animal(Document):
         doc.reload()
         Animal.drop_collection()
 
+    def test_reload_sharded_nested(self):
+        class SuperPhylum(EmbeddedDocument):
+            name = StringField()
+
+        class Animal(Document):
+            superphylum = EmbeddedDocumentField(SuperPhylum)
+            meta = {'shard_key': ('superphylum.name',)}
+
+        Animal.drop_collection()
+        doc = Animal(superphylum=SuperPhylum(name='Deuterostomia'))
+        doc.save()
+        doc.reload()
+        Animal.drop_collection()
+
     def test_reload_referencing(self):
         """Ensures reloading updates weakrefs correctly
         """
@@ -2715,6 +2729,32 @@ def change_shard_key():
 
         self.assertRaises(OperationError, change_shard_key)
 
+    def test_shard_key_in_embedded_document(self):
+        class Foo(EmbeddedDocument):
+            foo = StringField()
+
+        class Bar(Document):
+            meta = {
+                'shard_key': ('foo.foo',)
+            }
+            foo = EmbeddedDocumentField(Foo)
+            bar = StringField()
+
+        foo_doc = Foo(foo='hello')
+        bar_doc = Bar(foo=foo_doc, bar='world')
+        bar_doc.save()
+
+        self.assertTrue(bar_doc.id is not None)
+
+        bar_doc.bar = 'baz'
+        bar_doc.save()
+
+        def change_shard_key():
+            bar_doc.foo.foo = 'something'
+            bar_doc.save()
+
+        self.assertRaises(OperationError, change_shard_key)
+
     def test_shard_key_primary(self):
         class LogEntry(Document):
             machine = StringField(primary_key=True)