From 7bb9c7d47f18cdee5c68eceb7cbe6792a962dbae Mon Sep 17 00:00:00 2001
From: Nick Joyce <nick.joyce@realkinetic.com>
Date: Mon, 7 Jan 2013 14:49:39 +0000
Subject: [PATCH] Ensure that the update actions are grouped rather than
 serial.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This is a performance update. When multiple properties of the same
entity have been deleted and modified, 2 calls to update the entity are
made, one {"$set": … } and another {"$unset": … }. This is 2 network
interface calls which is a performance killer (even at lan speeds).

Fixes: #210
---
 mongoengine/document.py | 11 ++++++++---
 tests/test_document.py  | 30 ++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/mongoengine/document.py b/mongoengine/document.py
index 7b3afafb..2bd1d221 100644
--- a/mongoengine/document.py
+++ b/mongoengine/document.py
@@ -234,11 +234,16 @@ class Document(BaseDocument):
                     select_dict[actual_key] = doc[actual_key]
 
                 upsert = self._created
+                work = {}
+
                 if updates:
-                    collection.update(select_dict, {"$set": updates},
-                        upsert=upsert, safe=safe, **write_options)
+                    work["$set"] = updates
+
                 if removals:
-                    collection.update(select_dict, {"$unset": removals},
+                    work["$unset"] = removals
+
+                if work:
+                    collection.update(select_dict, work,
                         upsert=upsert, safe=safe, **write_options)
 
             warn_cascade = not cascade and 'cascade' not in self._meta
diff --git a/tests/test_document.py b/tests/test_document.py
index cd0ab8fb..fa6eb282 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -3515,6 +3515,36 @@ class ValidatorErrorTest(unittest.TestCase):
 
         self.assertRaises(OperationError, change_shard_key)
 
+    def test_set_unset_one_operation(self):
+        """Ensure that $set and $unset actions are performed in the same
+        operation.
+        """
+        class FooBar(Document):
+            meta = {
+                'collection': 'foobar',
+            }
+
+            foo = StringField(default=None)
+            bar = StringField(default=None)
+
+        FooBar.drop_collection()
+
+        # write an entity with a single prop
+        foo = FooBar(foo='foo')
+        foo.save()
+
+        self.assertEqual(foo.foo, 'foo')
+
+        # set foo to the default causing mongoengine to $unset foo.
+        foo.foo = None
+        foo.bar = 'bar'
+
+        foo.save()
+        foo.reload()
+
+        self.assertIsNone(foo.foo)
+        self.assertEqual(foo.bar, 'bar')
+
 
 if __name__ == '__main__':
     unittest.main()