diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index 0f6c8698..0c58b57c 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -264,4 +264,6 @@ def set_read_write_concern(collection, write_concerns, read_concern): combined_write_concerns = dict(collection.write_concern.document.items()) combined_write_concerns.update(write_concerns) - yield collection.with_options(write_concern=WriteConcern(**combined_write_concerns), read_concern=read_concern) + yield collection.with_options( + write_concern=WriteConcern(**combined_write_concerns), read_concern=read_concern + ) diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 0b76235c..0743429c 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -19,7 +19,11 @@ 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 set_write_concern, set_read_write_concern, switch_db +from mongoengine.context_managers import ( + set_write_concern, + set_read_write_concern, + switch_db, +) from mongoengine.errors import ( BulkWriteError, InvalidQueryError, @@ -492,7 +496,13 @@ class BaseQuerySet(object): return result.deleted_count def update( - self, upsert=False, multi=True, write_concern=None, read_concern=None, full_result=False, **update + self, + upsert=False, + multi=True, + write_concern=None, + read_concern=None, + full_result=False, + **update ): """Perform an atomic update on the fields matched by the query. @@ -531,7 +541,9 @@ class BaseQuerySet(object): else: update["$set"] = {"_cls": queryset._document._class_name} try: - with set_read_write_concern(queryset._collection, write_concern, read_concern) as collection: + with set_read_write_concern( + queryset._collection, write_concern, read_concern + ) as collection: update_func = collection.update_one if multi: update_func = collection.update_many @@ -1214,7 +1226,7 @@ class BaseQuerySet(object): queryset._read_concern = read_concern queryset._cursor_obj = None # we need to re-create the cursor object whenever we apply read_concern return queryset - + def scalar(self, *fields): """Instead of returning Document instances, return either a specific value or a tuple of values in order. diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py index c238752d..708033f4 100644 --- a/tests/queryset/test_queryset.py +++ b/tests/queryset/test_queryset.py @@ -4671,41 +4671,33 @@ class TestQueryset(unittest.TestCase): bars = list(Bar.objects.read_concern(None)) assert bars == [bar] - bars = Bar.objects.read_concern(ReadConcern(level='local')) - assert bars._read_concern == ReadConcern(level='local') - assert ( - bars._cursor.collection.read_concern - == ReadConcern(level='local') - ) + bars = Bar.objects.read_concern(ReadConcern(level="local")) + assert bars._read_concern == ReadConcern(level="local") + assert bars._cursor.collection.read_concern == ReadConcern(level="local") # Make sure that `.read_concern(...)` does accept string values. with pytest.raises(TypeError): - Bar.objects.read_concern('local') + Bar.objects.read_concern("local") def assert_read_concern(qs, expected_read_concern): assert qs._read_concern == expected_read_concern assert qs._cursor.collection.read_concern == expected_read_concern # Make sure read concern is respected after a `.skip(...)`. - bars = Bar.objects.skip(1).read_concern(ReadConcern('majority')) - assert_read_concern(bars, ReadConcern('majority')) + bars = Bar.objects.skip(1).read_concern(ReadConcern("majority")) + assert_read_concern(bars, ReadConcern("majority")) # Make sure read concern is respected after a `.limit(...)`. - bars = Bar.objects.limit(1).read_concern(ReadConcern('majority')) - assert_read_concern(bars, ReadConcern('majority')) + bars = Bar.objects.limit(1).read_concern(ReadConcern("majority")) + assert_read_concern(bars, ReadConcern("majority")) # Make sure read concern is respected after an `.order_by(...)`. - bars = Bar.objects.order_by("txt").read_concern( - ReadConcern('majority') - ) - assert_read_concern(bars, ReadConcern('majority')) + bars = Bar.objects.order_by("txt").read_concern(ReadConcern("majority")) + assert_read_concern(bars, ReadConcern("majority")) # Make sure read concern is respected after a `.hint(...)`. - bars = Bar.objects.hint([("txt", 1)]).read_concern( - ReadConcern('majority') - ) - assert_read_concern(bars, ReadConcern('majority')) - + bars = Bar.objects.hint([("txt", 1)]).read_concern(ReadConcern("majority")) + assert_read_concern(bars, ReadConcern("majority")) def test_json_simple(self): class Embedded(EmbeddedDocument):