diff --git a/.travis.yml b/.travis.yml index a70c711e..29a72d1c 100644 --- a/.travis.yml +++ b/.travis.yml @@ -19,8 +19,6 @@ python: - pypy env: -- MONGODB=2.6 PYMONGO=2.7 -- MONGODB=2.6 PYMONGO=2.8 - MONGODB=2.6 PYMONGO=3.x matrix: @@ -28,14 +26,10 @@ matrix: fast_finish: true include: - - python: 2.7 - env: MONGODB=2.4 PYMONGO=2.7 - python: 2.7 env: MONGODB=2.4 PYMONGO=3.5 - python: 2.7 env: MONGODB=3.0 PYMONGO=3.x - - python: 3.5 - env: MONGODB=2.4 PYMONGO=2.7 - python: 3.5 env: MONGODB=2.4 PYMONGO=3.5 - python: 3.5 diff --git a/docs/changelog.rst b/docs/changelog.rst index 4dec42ce..47c9af9c 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -11,6 +11,7 @@ dev - Update cached fields when fields argument is given #1712 - Add a db parameter to register_connection for compatibility with connect - Use insert_one, insert_many in Document.insert #1491 +- Use new update_one, update_many on document/queryset update #1491 Changes in 0.15.0 ================= diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index c477575e..ec2e9e8b 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -1,9 +1,11 @@ +from contextlib import contextmanager +from pymongo.write_concern import WriteConcern from mongoengine.common import _import_class from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db __all__ = ('switch_db', 'switch_collection', 'no_dereference', - 'no_sub_classes', 'query_counter') + 'no_sub_classes', 'query_counter', 'set_write_concern') class switch_db(object): @@ -215,3 +217,10 @@ class query_counter(object): count = self.db.system.profile.find(ignore_query).count() - self.counter self.counter += 1 return count + + +@contextmanager +def set_write_concern(collection, write_concerns): + combined_concerns = dict(collection.write_concern.document.items()) + combined_concerns.update(write_concerns) + yield collection.with_options(write_concern=WriteConcern(**combined_concerns)) diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index eec7df18..2f2d38dd 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -18,7 +18,7 @@ from mongoengine import signals from mongoengine.base import get_document from mongoengine.common import _import_class from mongoengine.connection import get_db -from mongoengine.context_managers import switch_db +from mongoengine.context_managers import set_write_concern, switch_db from mongoengine.errors import (InvalidQueryError, LookUpError, NotUniqueError, OperationError) from mongoengine.python_support import IS_PYMONGO_3 @@ -517,12 +517,15 @@ class BaseQuerySet(object): else: update['$set'] = {'_cls': queryset._document._class_name} try: - result = queryset._collection.update(query, update, multi=multi, - upsert=upsert, **write_concern) + with set_write_concern(queryset._collection, write_concern) as collection: + update_func = collection.update_one + if multi: + update_func = collection.update_many + result = update_func(query, update, upsert=upsert) if full_result: return result - elif result: - return result['n'] + elif result.raw_result: + return result.raw_result['n'] except pymongo.errors.DuplicateKeyError as err: raise NotUniqueError(u'Update failed (%s)' % six.text_type(err)) except pymongo.errors.OperationFailure as err: @@ -551,10 +554,10 @@ class BaseQuerySet(object): write_concern=write_concern, full_result=True, **update) - if atomic_update['updatedExisting']: + if atomic_update.raw_result['updatedExisting']: document = self.get() else: - document = self._document.objects.with_id(atomic_update['upserted']) + document = self._document.objects.with_id(atomic_update.upserted_id) return document def update_one(self, upsert=False, write_concern=None, **update): diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 4f940520..5c3f179f 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -9,6 +9,7 @@ from nose.plugins.skip import SkipTest import pymongo from pymongo.errors import ConfigurationError from pymongo.read_preferences import ReadPreference +from pymongo.results import UpdateResult import six from mongoengine import * @@ -670,14 +671,14 @@ class QuerySetTest(unittest.TestCase): result = self.Person(name="Bob", age=25).update( upsert=True, full_result=True) - self.assertTrue(isinstance(result, dict)) - self.assertTrue("upserted" in result) - self.assertFalse(result["updatedExisting"]) + self.assertTrue(isinstance(result, UpdateResult)) + self.assertTrue("upserted" in result.raw_result) + self.assertFalse(result.raw_result["updatedExisting"]) bob = self.Person.objects.first() result = bob.update(set__age=30, full_result=True) - self.assertTrue(isinstance(result, dict)) - self.assertTrue(result["updatedExisting"]) + self.assertTrue(isinstance(result, UpdateResult)) + self.assertTrue(result.raw_result["updatedExisting"]) self.Person(name="Bob", age=20).save() result = self.Person.objects(name="Bob").update( diff --git a/tox.ini b/tox.ini index 9bb0c5ec..2f2b1757 100644 --- a/tox.ini +++ b/tox.ini @@ -1,13 +1,11 @@ [tox] -envlist = {py27,py35,pypy,pypy3}-{mg27,mg28,mg35,mg3x} +envlist = {py27,py35,pypy,pypy3}-{mg35,mg3x} [testenv] commands = python setup.py nosetests {posargs} deps = nose - mg27: PyMongo<2.8 - mg28: PyMongo>=2.8,<2.9 mg35: PyMongo==3.5 mg3x: PyMongo>=3.0 setenv =