From a7aead5138f8792e4030328e23531ee7d4024744 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Wed, 24 Jun 2015 22:34:38 -0700 Subject: [PATCH 1/2] re-create the cursor object whenever we apply read_preference --- docs/changelog.rst | 1 + mongoengine/queryset/base.py | 1 + tests/queryset/queryset.py | 48 +++++++++++++++++++++++++++++++++++- 3 files changed, 49 insertions(+), 1 deletion(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 37720431..1ea4d19f 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -7,6 +7,7 @@ Changes in 0.10.2 - Allow shard key to point to a field in an embedded document. #551 - Allow arbirary metadata in fields. #1129 - ReferenceFields now support abstract document types. #837 +- Fix `read_preference` (it had chaining issues with PyMongo 2.x and it didn't work at all with PyMongo 3.x) #1042 Changes in 0.10.1 ======================= diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 0867e0e2..8880c2e3 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -930,6 +930,7 @@ class BaseQuerySet(object): validate_read_preference('read_preference', read_preference) queryset = self.clone() queryset._read_preference = read_preference + queryset._cursor_obj = None # we need to re-create the cursor object whenever we apply read_preference return queryset def scalar(self, *fields): diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 3eff7cea..8726801e 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -4165,7 +4165,11 @@ class QuerySetTest(unittest.TestCase): def test_read_preference(self): class Bar(Document): - pass + txt = StringField() + + meta = { + 'indexes': [ 'txt' ] + } Bar.drop_collection() bars = list(Bar.objects(read_preference=ReadPreference.PRIMARY)) @@ -4177,9 +4181,51 @@ class QuerySetTest(unittest.TestCase): error_class = TypeError self.assertRaises(error_class, Bar.objects, read_preference='Primary') + # read_preference as a kwarg bars = Bar.objects(read_preference=ReadPreference.SECONDARY_PREFERRED) self.assertEqual( bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) + + # read_preference as a query set method + bars = Bar.objects.read_preference(ReadPreference.SECONDARY_PREFERRED) + self.assertEqual( + bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) + + # read_preference after skip + bars = Bar.objects.skip(1) \ + .read_preference(ReadPreference.SECONDARY_PREFERRED) + self.assertEqual( + bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) + + # read_preference after limit + bars = Bar.objects.limit(1) \ + .read_preference(ReadPreference.SECONDARY_PREFERRED) + self.assertEqual( + bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) + + # read_preference after order_by + bars = Bar.objects.order_by('txt') \ + .read_preference(ReadPreference.SECONDARY_PREFERRED) + self.assertEqual( + bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) + + # read_preference after hint + bars = Bar.objects.hint([('txt', 1)]) \ + .read_preference(ReadPreference.SECONDARY_PREFERRED) + self.assertEqual( + bars._read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(bars._cursor._Cursor__read_preference, + ReadPreference.SECONDARY_PREFERRED) def test_json_simple(self): From 54975de0f387ad02a75134a4275401f684b855d7 Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Sun, 5 Jul 2015 15:59:50 -0700 Subject: [PATCH 2/2] fix read_preference for PyMongo 3+ --- mongoengine/queryset/base.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 8880c2e3..ac4764d3 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -1444,8 +1444,16 @@ class BaseQuerySet(object): def _cursor(self): if self._cursor_obj is None: - self._cursor_obj = self._collection.find(self._query, - **self._cursor_args) + # In PyMongo 3+, we define the read preference on a collection + # level, not a cursor level. Thus, we need to get a cloned + # collection object using `with_options` first. + if IS_PYMONGO_3 and self._read_preference is not None: + self._cursor_obj = self._collection\ + .with_options(read_preference=self._read_preference)\ + .find(self._query, **self._cursor_args) + else: + self._cursor_obj = self._collection.find(self._query, + **self._cursor_args) # Apply where clauses to cursor if self._where_clause: where_clause = self._sub_js_fields(self._where_clause)