Merge pull request #1042 from closeio/fix-read-preference

Fix read_preference
MongoEngine · Nov 24, 2015 · 9050869 · 9050869
2 parents 6868f66 + 54975de
commit 9050869
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 3 deletions.
diff --git 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
@@ -930,6 +930,7 @@ def read_preference(self, read_preference):
         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 @@ def _cursor_args(self):
     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
@@ -4165,7 +4165,11 @@ class MyDoc(Document):
 
     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 Bar(Document):
             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):