From f265915aa227f941fa07fae2df85155be9e0f3d1 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Thu, 8 Nov 2012 16:35:20 +0000
Subject: [PATCH] Updated inheritable objects created by upsert now contain
 _cls (MongoEngine/mongoengine#118)

---
 docs/changelog.rst               |  1 +
 mongoengine/queryset/queryset.py | 12 ++++++++++--
 tests/test_queryset.py           | 19 +++++++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index 778a047f..e8d3d574 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -4,6 +4,7 @@ Changelog
 
 Changes in 0.8
 ==============
+- Updated inheritable objects created by upsert now contain _cls (MongoEngine/mongoengine#118)
 - Added support for creating documents with embedded documents in a single operation (MongoEngine/mongoengine#6)
 - Added to_json and from_json to Document (MongoEngine/mongoengine#1)
 - Added to_json and from_json to QuerySet (MongoEngine/mongoengine#131)
diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py
index 3c44f012..bfd15a8c 100644
--- a/mongoengine/queryset/queryset.py
+++ b/mongoengine/queryset/queryset.py
@@ -824,8 +824,16 @@ class QuerySet(object):
         if not write_options:
             write_options = {}
 
-        update = transform.update(self._document, **update)
         query = self._query
+        update = transform.update(self._document, **update)
+
+        # If doing an atomic upsert on an inheritable class
+        # then ensure we add _cls to the update operation
+        if upsert and '_cls' in query:
+            if '$set' in update:
+                update["$set"]["_cls"] = self._document._class_name
+            else:
+                update["$set"] = {"_cls": self._document._class_name}
 
         try:
             ret = self._collection.update(query, update, multi=multi,
@@ -852,7 +860,7 @@ class QuerySet(object):
 
         .. versionadded:: 0.2
         """
-        return self.update(safe_update=True, upsert=False, multi=False,
+        return self.update(safe_update=True, upsert=upsert, multi=False,
                            write_options=None, **update)
 
     def __iter__(self):
diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index 9dfe9a27..a86920e9 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -3741,6 +3741,25 @@ class QueryFieldListTest(unittest.TestCase):
 
         self.assertEqual(doc_objects, Doc.objects.from_json(json_data))
 
+    def test_upsert_includes_cls(self):
+        """Upserts should include _cls information for inheritable classes
+        """
+
+        class Test(Document):
+            test = StringField()
+
+        Test.drop_collection()
+        Test.objects(test='foo').update_one(upsert=True, set__test='foo')
+        self.assertFalse('_cls' in Test._collection.find_one())
+
+        class Test(Document):
+            meta = {'allow_inheritance': True}
+            test = StringField()
+
+        Test.drop_collection()
+
+        Test.objects(test='foo').update_one(upsert=True, set__test='foo')
+        self.assertTrue('_cls' in Test._collection.find_one())
 
 if __name__ == '__main__':
     unittest.main()