From 4d5c6d11ab8665f77aef8ac6e9b2e58843188b75 Mon Sep 17 00:00:00 2001 From: Calgary Michael Date: Fri, 2 Feb 2018 22:04:30 -0600 Subject: [PATCH 1/3] removed deprecated warning for 'update' method --- mongoengine/context_managers.py | 11 ++++++++++- mongoengine/queryset/base.py | 17 ++++++++++------- tests/queryset/queryset.py | 11 ++++++----- 3 files changed, 26 insertions(+), 13 deletions(-) diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index c477575e..e6295570 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -1,9 +1,11 @@ from mongoengine.common import _import_class from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db +from pymongo.write_concern import WriteConcern +from contextlib import contextmanager __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): + yield collection.with_options(write_concern=WriteConcern( + **dict(collection.write_concern.document.items()), + **write_concerns)) diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 6f9c372c..c6e2137c 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 switch_db, set_write_concern from mongoengine.errors import (InvalidQueryError, LookUpError, NotUniqueError, OperationError) from mongoengine.python_support import IS_PYMONGO_3 @@ -510,12 +510,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: @@ -544,10 +547,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 43800fff..848fe35d 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 * @@ -656,14 +657,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( From fa38bfd4e8884fe2d6c641fa185c90641c5a24c8 Mon Sep 17 00:00:00 2001 From: Calgary Michael Date: Fri, 2 Feb 2018 22:30:06 -0600 Subject: [PATCH 2/3] made set_write_concern python2.7 compatible --- mongoengine/context_managers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index e6295570..cc35cdd7 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -221,6 +221,6 @@ class query_counter(object): @contextmanager def set_write_concern(collection, write_concerns): - yield collection.with_options(write_concern=WriteConcern( - **dict(collection.write_concern.document.items()), - **write_concerns)) + old_concerns = dict(collection.write_concern.document.items()) + combined_concerns = old_concerns.update(write_concerns) + yield collection.with_options(write_concern=WriteConcern(**combined_concerns)) From 6835c15d9b0a0379036417e642516f7cdef76840 Mon Sep 17 00:00:00 2001 From: Calgary Michael Date: Fri, 2 Feb 2018 22:41:07 -0600 Subject: [PATCH 3/3] fixing bug in previous commit --- mongoengine/context_managers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index cc35cdd7..c6d0d40f 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -221,6 +221,6 @@ class query_counter(object): @contextmanager def set_write_concern(collection, write_concerns): - old_concerns = dict(collection.write_concern.document.items()) - combined_concerns = old_concerns.update(write_concerns) + combined_concerns = dict(collection.write_concern.document.items()) + combined_concerns.update(write_concerns) yield collection.with_options(write_concern=WriteConcern(**combined_concerns))