PYTHON-2554 Update for new 5.0 spec rules

Author: ShaneHarvey

pymongo/aggregation.py

@@ -19,7 +19,7 @@
 from pymongo import common
 from pymongo.collation import validate_collation_or_none
 from pymongo.errors import ConfigurationError
-from pymongo.read_preferences import _AggWritePref
+from pymongo.read_preferences import _AggWritePref, ReadPreference
 
 
 class _AggregationCommand(object):
@@ -101,7 +101,7 @@ def get_read_preference(self, session):
         if self._write_preference:
             return self._write_preference
         pref = self._target._read_preference_for(session)
-        if self._performs_write:
+        if self._performs_write and pref != ReadPreference.PRIMARY:
             self._write_preference = pref = _AggWritePref(pref)
         return pref
 

pymongo/mongo_client.py

@@ -1161,7 +1161,7 @@ def _secondaryok_for_server(self, read_preference, server, session):
 
         with self._get_socket(server, session) as sock_info:
             secondary_ok = (single and not sock_info.is_mongos) or (
-                read_preference != ReadPreference.PRIMARY)
+                read_preference.mode != ReadPreference.PRIMARY.mode)
             yield sock_info, secondary_ok
 
     @contextlib.contextmanager

pymongo/read_preferences.py

@@ -438,19 +438,20 @@ class _AggWritePref:
 
     def __init__(self, pref):
         self.pref = pref
-        self.effective_pref = None
+        self.effective_pref = ReadPreference.PRIMARY
+
+    def selection_hook(self, topology_description):
+        common_wv = topology_description.common_wire_version
+        if (topology_description.has_readable_server(
+                ReadPreference.PRIMARY_PREFERRED) and
+            common_wv and common_wv < 13):
+            self.effective_pref = ReadPreference.PRIMARY
+        else:
+            self.effective_pref = self.pref
 
     def __call__(self, selection):
         """Apply this read preference to a Selection."""
-        has_50 = any(s for s in selection.server_descriptions
-                     if s.max_wire_version >= 13)
-        if not has_50:
-            # Fallback to using the primary:
-            self.effective_pref = ReadPreference.PRIMARY
-            return selection.primary_selection
-        # Apply the original selection:
-        self.effective_pref = self.pref
-        return self.pref(selection)
+        return self.effective_pref(selection)
 
     def __repr__(self):
         return "_AggWritePref(pref=%r)" % (self.pref,)

pymongo/topology_description.py

@@ -263,24 +263,24 @@ def apply_selector(self, selector, address=None, custom_selector=None):
                                              selector.min_wire_version,
                                              common_wv))
 
+        if isinstance(selector, _AggWritePref):
+            selector.selection_hook(self)
+
         if self.topology_type == TOPOLOGY_TYPE.Unknown:
             return []
         elif self.topology_type in (TOPOLOGY_TYPE.Single,
                                     TOPOLOGY_TYPE.LoadBalanced):
             # Ignore selectors for standalone and load balancer mode.
             return self.known_servers
-        elif address:
+        if address:
             # Ignore selectors when explicit address is requested.
             description = self.server_descriptions().get(address)
             return [description] if description else []
-        elif (self.topology_type == TOPOLOGY_TYPE.Sharded and
-              not isinstance(selector, _AggWritePref)):
-            # TODO: adjust _AggWritePref for sharded clusters once we select
-            # a server.
-            # Ignore read preference.
-            selection = Selection.from_topology_description(self)
-        else:
-            selection = selector(Selection.from_topology_description(self))
+
+        selection = Selection.from_topology_description(self)
+        # Ignore read preference for sharded clusters.
+        if self.topology_type != TOPOLOGY_TYPE.Sharded:
+            selection = selector(selection)
 
         # Apply custom selector followed by localThresholdMS.
         if custom_selector is not None and selection:

test/test_read_preferences.py

@@ -433,7 +433,9 @@ def test_aggregate(self):
                                [{'$project': {'_id': 1}}])
 
     def test_aggregate_write(self):
-        self._test_coll_helper(False, self.c.pymongo_test.test,
+        # 5.0 servers support $out on secondaries.
+        secondary_ok = client_context.version.at_least(5, 0)
+        self._test_coll_helper(secondary_ok, self.c.pymongo_test.test,
                                'aggregate',
                                [{'$project': {'_id': 1}}, {'$out': "agg_write_test"}])