pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'calgary/issue-1491' into update_pymongo3

Browse Source
This commit is contained in:
Erdenezul Batmunkh 2018-05-20 12:38:22 +08:00
commit 0da8fb379d
3 changed files with 26 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mongoengine/context_managers.py
View File

@ -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):
combined_concerns = dict(collection.write_concern.document.items())
combined_concerns.update(write_concerns)
yield collection.with_options(write_concern=WriteConcern(**combined_concerns))

17
mongoengine/queryset/base.py
View File

@ -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
@ -511,12 +511,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:
@ -545,10 +548,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):

11
tests/queryset/queryset.py
View File

@ -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(