diff --git a/AUTHORS b/AUTHORS index fbca84e4..dd04aee1 100644 --- a/AUTHORS +++ b/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) diff --git a/docs/changelog.rst b/docs/changelog.rst index c9285416..7776a99b 100644 --- a/docs/changelog.rst +++ b/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 ================= diff --git a/mongoengine/__init__.py b/mongoengine/__init__.py index eefaf08d..14a7c1ce 100644 --- a/mongoengine/__init__.py +++ b/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(): diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index ac4764d3..ce4174b9 100644 --- a/mongoengine/queryset/base.py +++ b/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. diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 8726801e..87c14757 100644 --- a/tests/queryset/queryset.py +++ b/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()