Merge pull request MongoEngine#1167 from BeardedSteve/upsert_one

Upsert one

Author: Omer Katz

AUTHORS

@@ -231,3 +231,4 @@ that much better:
  * Lars Butler (https://github.com/larsbutler)
  * George Macon (https://github.com/gmacon)
  * Ashley Whetter (https://github.com/AWhetter)
+ * Steven Rossiter (https://github.com/BeardedSteve)

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 Changes in 0.10.4 - DEV
 =======================
 - SaveConditionError is now importable from the top level package. #1165
+- upsert_one method added. #1157
 
 Changes in 0.10.3
 =================

mongoengine/__init__.py

@@ -14,7 +14,7 @@
 __all__ = (list(document.__all__) + fields.__all__ + connection.__all__ +
            list(queryset.__all__) + signals.__all__ + list(errors.__all__))
 
-VERSION = (0, 10, 1)
+VERSION = (0, 10, 4)
 
 
 def get_version():

mongoengine/queryset/base.py

@@ -471,6 +471,32 @@ def update(self, upsert=False, multi=True, write_concern=None,
                 raise OperationError(message)
             raise OperationError(u'Update failed (%s)' % unicode(err))
 
+
+    def upsert_one(self, write_concern=None, **update):
+        """Overwrite or add the first document matched by the query.
+
+        :param write_concern: Extra keyword arguments are passed down which
+            will be used as options for the resultant
+            ``getLastError`` command.  For example,
+            ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
+            wait until at least two servers have recorded the write and
+            will force an fsync on the primary server.
+        :param update: Django-style update keyword arguments
+
+        :returns the new or overwritten document
+
+        .. versionadded:: 0.10.2
+        """
+
+        atomic_update = self.update(multi=False, upsert=True, write_concern=write_concern,
+                             full_result=True,**update)
+
+        if atomic_update['updatedExisting']:
+            document = self.get()
+        else:
+            document = self._document.objects.with_id(atomic_update['upserted'])
+        return document
+
     def update_one(self, upsert=False, write_concern=None, **update):
         """Perform an atomic update on the fields of the first document
         matched by the query.

tests/queryset/queryset.py

@@ -680,12 +680,21 @@ def test_upsert(self):
     def test_upsert_one(self):
         self.Person.drop_collection()
 
-        self.Person.objects(name="Bob", age=30).update_one(upsert=True)
+        bob = self.Person.objects(name="Bob", age=30).upsert_one()
 
-        bob = self.Person.objects.first()
         self.assertEqual("Bob", bob.name)
         self.assertEqual(30, bob.age)
 
+        bob.name = "Bobby"
+        bob.save()
+
+        bobby = self.Person.objects(name="Bobby", age=30).upsert_one()
+
+        self.assertEqual("Bobby", bobby.name)
+        self.assertEqual(30, bobby.age)
+        self.assertEqual(bob.id, bobby.id)
+
+
     def test_set_on_insert(self):
         self.Person.drop_collection()