removed deprecated warning for 'update' method
This commit is contained in:
		| @@ -1,9 +1,11 @@ | |||||||
| from mongoengine.common import _import_class | from mongoengine.common import _import_class | ||||||
| from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db | 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', | __all__ = ('switch_db', 'switch_collection', 'no_dereference', | ||||||
|            'no_sub_classes', 'query_counter') |            'no_sub_classes', 'query_counter', 'set_write_concern') | ||||||
|  |  | ||||||
|  |  | ||||||
| class switch_db(object): | class switch_db(object): | ||||||
| @@ -215,3 +217,10 @@ class query_counter(object): | |||||||
|         count = self.db.system.profile.find(ignore_query).count() - self.counter |         count = self.db.system.profile.find(ignore_query).count() - self.counter | ||||||
|         self.counter += 1 |         self.counter += 1 | ||||||
|         return count |         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.base import get_document | ||||||
| from mongoengine.common import _import_class | from mongoengine.common import _import_class | ||||||
| from mongoengine.connection import get_db | 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, | from mongoengine.errors import (InvalidQueryError, LookUpError, | ||||||
|                                 NotUniqueError, OperationError) |                                 NotUniqueError, OperationError) | ||||||
| from mongoengine.python_support import IS_PYMONGO_3 | from mongoengine.python_support import IS_PYMONGO_3 | ||||||
| @@ -510,12 +510,15 @@ class BaseQuerySet(object): | |||||||
|             else: |             else: | ||||||
|                 update['$set'] = {'_cls': queryset._document._class_name} |                 update['$set'] = {'_cls': queryset._document._class_name} | ||||||
|         try: |         try: | ||||||
|             result = queryset._collection.update(query, update, multi=multi, |             with set_write_concern(queryset._collection, write_concern) as collection: | ||||||
|                                                  upsert=upsert, **write_concern) |                 update_func = collection.update_one | ||||||
|  |                 if multi: | ||||||
|  |                     update_func = collection.update_many | ||||||
|  |                 result = update_func(query, update, upsert=upsert) | ||||||
|             if full_result: |             if full_result: | ||||||
|                 return result |                 return result | ||||||
|             elif result: |             elif result.raw_result: | ||||||
|                 return result['n'] |                 return result.raw_result['n'] | ||||||
|         except pymongo.errors.DuplicateKeyError as err: |         except pymongo.errors.DuplicateKeyError as err: | ||||||
|             raise NotUniqueError(u'Update failed (%s)' % six.text_type(err)) |             raise NotUniqueError(u'Update failed (%s)' % six.text_type(err)) | ||||||
|         except pymongo.errors.OperationFailure as err: |         except pymongo.errors.OperationFailure as err: | ||||||
| @@ -544,10 +547,10 @@ class BaseQuerySet(object): | |||||||
|                                     write_concern=write_concern, |                                     write_concern=write_concern, | ||||||
|                                     full_result=True, **update) |                                     full_result=True, **update) | ||||||
|  |  | ||||||
|         if atomic_update['updatedExisting']: |         if atomic_update.raw_result['updatedExisting']: | ||||||
|             document = self.get() |             document = self.get() | ||||||
|         else: |         else: | ||||||
|             document = self._document.objects.with_id(atomic_update['upserted']) |             document = self._document.objects.with_id(atomic_update.upserted_id) | ||||||
|         return document |         return document | ||||||
|  |  | ||||||
|     def update_one(self, upsert=False, write_concern=None, **update): |     def update_one(self, upsert=False, write_concern=None, **update): | ||||||
|   | |||||||
| @@ -9,6 +9,7 @@ from nose.plugins.skip import SkipTest | |||||||
| import pymongo | import pymongo | ||||||
| from pymongo.errors import ConfigurationError | from pymongo.errors import ConfigurationError | ||||||
| from pymongo.read_preferences import ReadPreference | from pymongo.read_preferences import ReadPreference | ||||||
|  | from pymongo.results import UpdateResult | ||||||
| import six | import six | ||||||
|  |  | ||||||
| from mongoengine import * | from mongoengine import * | ||||||
| @@ -656,14 +657,14 @@ class QuerySetTest(unittest.TestCase): | |||||||
|  |  | ||||||
|         result = self.Person(name="Bob", age=25).update( |         result = self.Person(name="Bob", age=25).update( | ||||||
|             upsert=True, full_result=True) |             upsert=True, full_result=True) | ||||||
|         self.assertTrue(isinstance(result, dict)) |         self.assertTrue(isinstance(result, UpdateResult)) | ||||||
|         self.assertTrue("upserted" in result) |         self.assertTrue("upserted" in result.raw_result) | ||||||
|         self.assertFalse(result["updatedExisting"]) |         self.assertFalse(result.raw_result["updatedExisting"]) | ||||||
|  |  | ||||||
|         bob = self.Person.objects.first() |         bob = self.Person.objects.first() | ||||||
|         result = bob.update(set__age=30, full_result=True) |         result = bob.update(set__age=30, full_result=True) | ||||||
|         self.assertTrue(isinstance(result, dict)) |         self.assertTrue(isinstance(result, UpdateResult)) | ||||||
|         self.assertTrue(result["updatedExisting"]) |         self.assertTrue(result.raw_result["updatedExisting"]) | ||||||
|  |  | ||||||
|         self.Person(name="Bob", age=20).save() |         self.Person(name="Bob", age=20).save() | ||||||
|         result = self.Person.objects(name="Bob").update( |         result = self.Person.objects(name="Bob").update( | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user