From f27debe7f974449f34b0acc9a0c8a8de7f1be437 Mon Sep 17 00:00:00 2001
From: Dmitry Balabanov <dimonb@gmail.com>
Date: Thu, 30 Aug 2012 12:40:44 +0400
Subject: [PATCH 1/3] Respect sharding key when delete object from collection

---
 mongoengine/document.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/mongoengine/document.py b/mongoengine/document.py
index f8bf769d..1af7d017 100644
--- a/mongoengine/document.py
+++ b/mongoengine/document.py
@@ -265,6 +265,16 @@ class Document(BaseDocument):
                 ref.save(**kwargs)
                 ref._changed_fields = []
 
+    @property
+    def _object_key(self):
+        """Dict to identify object in collection
+        """
+        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)
+        return select_dict
+
     def update(self, **kwargs):
         """Performs an update on the :class:`~mongoengine.Document`
         A convenience wrapper to :meth:`~mongoengine.QuerySet.update`.
@@ -276,11 +286,7 @@ class Document(BaseDocument):
             raise OperationError('attempt to update a document not yet saved')
 
         # Need to add shard key to query, or you get an error
-        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)
-        return self.__class__.objects(**select_dict).update_one(**kwargs)
+        return self.__class__.objects(**self._object_key).update_one(**kwargs)
 
     def delete(self, safe=False):
         """Delete the :class:`~mongoengine.Document` from the database. This
@@ -291,7 +297,7 @@ class Document(BaseDocument):
         signals.pre_delete.send(self.__class__, document=self)
 
         try:
-            self.__class__.objects(pk=self.pk).delete(safe=safe)
+            self.__class__.objects(**self._object_key).delete(safe=safe)
         except pymongo.errors.OperationFailure, err:
             message = u'Could not delete document (%s)' % err.message
             raise OperationError(message)

From d79ae30f31388ebff9b10350f8291df7bdd02313 Mon Sep 17 00:00:00 2001
From: Dmitry Balabanov <dimonb@gmail.com>
Date: Thu, 30 Aug 2012 14:20:53 +0400
Subject: [PATCH 2/3] fix object reload with shard_key in meta

---
 mongoengine/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mongoengine/base.py b/mongoengine/base.py
index 6fb26cb7..0a9e10c5 100644
--- a/mongoengine/base.py
+++ b/mongoengine/base.py
@@ -840,7 +840,7 @@ class BaseDocument(object):
                 if hasattr(self, '_changed_fields'):
                     self._mark_as_changed(name)
 
-        if not self._created and name in self._meta.get('shard_key', tuple()):
+        if not self._created and name in self._meta.get('shard_key', tuple()) and self._data[name] != value:
             from queryset import OperationError
             raise OperationError("Shard Keys are immutable. Tried to update %s" % name)
 

From ab60fd0490dc6b071973a344b7f4ddb23f7d2e84 Mon Sep 17 00:00:00 2001
From: Dmitry Balabanov <dimonb@gmail.com>
Date: Thu, 30 Aug 2012 14:37:11 +0400
Subject: [PATCH 3/3] sharded collection document reload testcase

---
 tests/test_document.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/test_document.py b/tests/test_document.py
index 6915caff..a551040b 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -1226,6 +1226,17 @@ class DocumentTest(unittest.TestCase):
         self.assertEqual(person.name, "Mr Test User")
         self.assertEqual(person.age, 21)
 
+    def test_reload_sharded(self):
+        class Animal(Document):
+            superphylum = StringField()
+            meta = {'shard_key': ('superphylum',)}
+
+        Animal.drop_collection()
+        doc = Animal(superphylum = 'Deuterostomia')
+        doc.save()
+        doc.reload()
+        Animal.drop_collection()
+
     def test_reload_referencing(self):
         """Ensures reloading updates weakrefs correctly
         """