PYTHON-2554 Support $merge and $out executing on secondaries

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 ReadPreference
+from pymongo.read_preferences import _AggWritePref, ReadPreference
 
 
 class _AggregationCommand(object):
@@ -70,6 +70,7 @@ def __init__(self, target, cursor_class, pipeline, options,
             options.pop('collation', None))
 
         self._max_await_time_ms = options.pop('maxAwaitTimeMS', None)
+        self._write_preference = None
 
     @property
     def _aggregation_target(self):
@@ -97,9 +98,12 @@ def _process_result(self, result, session, server, sock_info, secondary_ok):
                 result, session, server, sock_info, secondary_ok)
 
     def get_read_preference(self, session):
-        if self._performs_write:
-            return ReadPreference.PRIMARY
-        return self._target._read_preference_for(session)
+        if self._write_preference:
+            return self._write_preference
+        pref = self._target._read_preference_for(session)
+        if self._performs_write and pref != ReadPreference.PRIMARY:
+            self._write_preference = pref = _AggWritePref(pref)
+        return pref
 
     def get_cursor(self, session, server, sock_info, secondary_ok):
         # Serialize command.

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

@@ -424,6 +424,45 @@ def __call__(self, selection):
                 self.max_staleness, selection))
 
 
+class _AggWritePref:
+    """Agg $out/$merge write preference.
+
+    * If there are no readable servers, or there is any pre-5.0 server, use
+      `pref` read preference.
+    * Otherwise uses Primary().
+
+    :Parameters:
+      - `pref`: The read preference to use on MongoDB 5.0+.
+    """
+
+    __slots__ = ('pref', 'effective_pref')
+
+    def __init__(self, pref):
+        self.pref = pref
+        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."""
+        return self.effective_pref(selection)
+
+    def __repr__(self):
+        return "_AggWritePref(pref=%r)" % (self.pref,)
+
+    # Proxy other calls to the effective_pref so that _AggWritePref can be
+    # used in place of an actual read preference.
+    def __getattr__(self, name):
+        return getattr(self.effective_pref, name)
+
+
 _ALL_READ_PREFERENCES = (Primary, PrimaryPreferred,
                          Secondary, SecondaryPreferred, Nearest)
 

pymongo/topology_description.py

@@ -19,7 +19,7 @@
 
 from pymongo import common
 from pymongo.errors import ConfigurationError
-from pymongo.read_preferences import ReadPreference
+from pymongo.read_preferences import ReadPreference, _AggWritePref
 from pymongo.server_description import ServerDescription
 from pymongo.server_selectors import Selection
 from pymongo.server_type import SERVER_TYPE
@@ -263,21 +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:
-            # 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: