Add missing import

This commit is contained in:
Agustin Barto 2020-03-17 21:58:36 -03:00
parent bc77322c2f
commit 7cc964c7d8

View File

@ -1,6 +1,7 @@
from contextlib import contextmanager
from pymongo.write_concern import WriteConcern
from pymongo.read_concern import ReadConcern
from six import iteritems
from mongoengine.common import _import_class
@ -263,10 +264,14 @@ def set_write_concern(collection, write_concerns):
@contextmanager
def set_read_write_concern(collection, write_concerns, read_concerns):
combined_write_concerns = dict(collection.write_concern.document.items())
combined_write_concerns.update(write_concerns)
if write_concerns is not None:
combined_write_concerns.update(write_concerns)
combined_read_concerns = dict(collection.read_concern.document.items())
combined_read_concerns.update(read_concerns)
if read_concerns is not None:
combined_read_concerns.update(read_concerns)
yield collection.with_options(
write_concern=WriteConcern(**combined_write_concerns),