removed deprecated warning for 'update' method
This commit is contained in:
		| @@ -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)) | ||||
|   | ||||
| @@ -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): | ||||
|   | ||||
| @@ -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( | ||||
|   | ||||
		Reference in New Issue
	
	Block a user