Merge pull request #1167 from BeardedSteve/upsert_one

Upsert one

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

@@ -2,6 +2,11 @@
 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
 =================
 - Fix `read_preference` (it had chaining issues with PyMongo 2.x and it didn't work at all with PyMongo 3.x) #1042

mongoengine/__init__.py

@@ -14,7 +14,7 @@ import errors
 __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/document.py

@@ -217,7 +217,7 @@ class Document(BaseDocument):
         Returns True if the document has been updated or False if the document
         in the database doesn't match the query.
 
-        .. note:: All unsaved changes that has been made to the document are
+        .. note:: All unsaved changes that have been made to the document are
             rejected if the method returns True.
 
         :param query: the update will be performed only if the document in the
@@ -407,7 +407,7 @@ class Document(BaseDocument):
 
     def cascade_save(self, *args, **kwargs):
         """Recursively saves any references /
-           generic references on an objects"""
+           generic references on the document"""
         _refs = kwargs.get('_refs', []) or []
 
         ReferenceField = _import_class('ReferenceField')

mongoengine/errors.py

@@ -6,7 +6,7 @@ from mongoengine.python_support import txt_type
 __all__ = ('NotRegistered', 'InvalidDocumentError', 'LookUpError',
            'DoesNotExist', 'MultipleObjectsReturned', 'InvalidQueryError',
            'OperationError', 'NotUniqueError', 'FieldDoesNotExist',
-           'ValidationError')
+           'ValidationError', 'SaveConditionError')
 
 
 class NotRegistered(Exception):

mongoengine/queryset/base.py

@@ -471,6 +471,32 @@ class BaseQuerySet(object):
                 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 @@ class QuerySetTest(unittest.TestCase):
     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()