diff --git a/docs/changelog.rst b/docs/changelog.rst index 3a397524..8fa5af05 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -2,6 +2,9 @@ Changelog ========= +Changes in 0.8.3 +================ +- Added full_result kwarg to update (#380) Changes in 0.8.2 ================ diff --git a/mongoengine/document.py b/mongoengine/document.py index a61ed079..a1bac195 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -353,7 +353,13 @@ class Document(BaseDocument): been saved. """ if not self.pk: - raise OperationError('attempt to update a document not yet saved') + if kwargs.get('upsert', False): + query = self.to_mongo() + if "_cls" in query: + del(query["_cls"]) + return self._qs.filter(**query).update_one(**kwargs) + else: + raise OperationError('attempt to update a document not yet saved') # Need to add shard key to query, or you get an error return self._qs.filter(**self._object_key).update_one(**kwargs) diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index 9b53df2a..4b32ab12 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -474,7 +474,8 @@ class QuerySet(object): queryset._collection.remove(queryset._query, write_concern=write_concern) - def update(self, upsert=False, multi=True, write_concern=None, **update): + def update(self, upsert=False, multi=True, write_concern=None, + full_result=False, **update): """Perform an atomic update on the fields matched by the query. :param upsert: Any existing document with that "_id" is overwritten. @@ -485,6 +486,8 @@ class QuerySet(object): ``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 full_result: Return the full result rather than just the number + updated. :param update: Django-style update keyword arguments .. versionadded:: 0.2 @@ -506,12 +509,13 @@ class QuerySet(object): update["$set"]["_cls"] = queryset._document._class_name else: update["$set"] = {"_cls": queryset._document._class_name} - try: - ret = queryset._collection.update(query, update, multi=multi, - upsert=upsert, **write_concern) - if ret is not None and 'n' in ret: - return ret['n'] + result = queryset._collection.update(query, update, multi=multi, + upsert=upsert, **write_concern) + if full_result: + return result + elif result: + return result['n'] except pymongo.errors.OperationFailure, err: if unicode(err) == u'multi not coded yet': message = u'update() method requires MongoDB 1.1.3+' diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 21df22c5..bd231e33 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -536,6 +536,23 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(club.members['John']['gender'], "F") self.assertEqual(club.members['John']['age'], 14) + def test_update_results(self): + self.Person.drop_collection() + + result = self.Person(name="Bob", age=25).update(upsert=True, full_result=True) + self.assertIsInstance(result, dict) + self.assertTrue("upserted" in result) + self.assertFalse(result["updatedExisting"]) + + bob = self.Person.objects.first() + result = bob.update(set__age=30, full_result=True) + self.assertIsInstance(result, dict) + self.assertTrue(result["updatedExisting"]) + + self.Person(name="Bob", age=20).save() + result = self.Person.objects(name="Bob").update(set__name="bobby", multi=True) + self.assertEqual(result, 2) + def test_upsert(self): self.Person.drop_collection()