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..ac4764d3 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): @@ -1443,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) 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):