PYTHON-2170 Add support for 4.4 hedged reads

ShaneHarvey · ShaneHarvey · commit 6e39ae04c331 · 2020-05-07T17:26:00.000-07:00
diff --git a/pymongo/message.py b/pymongo/message.py
@@ -95,21 +95,18 @@ def _randint():
 def _maybe_add_read_preference(spec, read_preference):
     """Add $readPreference to spec when appropriate."""
     mode = read_preference.mode
-    tag_sets = read_preference.tag_sets
-    max_staleness = read_preference.max_staleness
+    document = read_preference.document
     # Only add $readPreference if it's something other than primary to avoid
     # problems with mongos versions that don't support read preferences. Also,
     # for maximum backwards compatibility, don't add $readPreference for
     # secondaryPreferred unless tags or maxStalenessSeconds are in use (setting
     # the slaveOkay bit has the same effect).
     if mode and (
-        mode != ReadPreference.SECONDARY_PREFERRED.mode
-        or tag_sets != [{}]
-        or max_staleness != -1):
-
+            mode != ReadPreference.SECONDARY_PREFERRED.mode or
+            len(document) > 1):
         if "$query" not in spec:
             spec = SON([("$query", spec)])
-        spec["$readPreference"] = read_preference.document
+        spec["$readPreference"] = document
     return spec
 
 
diff --git a/pymongo/read_preferences.py b/pymongo/read_preferences.py
@@ -81,17 +81,30 @@ def _validate_max_staleness(max_staleness):
     return max_staleness
 
 
+def _validate_hedge(hedge):
+    """Validate hedge."""
+    if hedge is None:
+        return None
+
+    if not isinstance(hedge, dict):
+        raise TypeError("hedge must be a dictionary, not %r" % (hedge,))
+
+    return hedge
+
+
 class _ServerMode(object):
     """Base class for all read preferences.
     """
 
-    __slots__ = ("__mongos_mode", "__mode", "__tag_sets", "__max_staleness")
+    __slots__ = ("__mongos_mode", "__mode", "__tag_sets", "__max_staleness",
+                 "__hedge")
 
-    def __init__(self, mode, tag_sets=None, max_staleness=-1):
+    def __init__(self, mode, tag_sets=None, max_staleness=-1, hedge=None):
         self.__mongos_mode = _MONGOS_MODES[mode]
         self.__mode = mode
         self.__tag_sets = _validate_tag_sets(tag_sets)
         self.__max_staleness = _validate_max_staleness(max_staleness)
+        self.__hedge = _validate_hedge(hedge)
 
     @property
     def name(self):
@@ -114,6 +127,8 @@ def document(self):
             doc['tags'] = self.__tag_sets
         if self.__max_staleness != -1:
             doc['maxStalenessSeconds'] = self.__max_staleness
+        if self.__hedge not in (None, {}):
+            doc['hedge'] = self.__hedge
         return doc
 
     @property
@@ -144,6 +159,30 @@ def max_staleness(self):
         no longer be selected for operations, or -1 for no maximum."""
         return self.__max_staleness
 
+    @property
+    def hedge(self):
+        """The read preference ``hedge`` parameter.
+
+        A dictionary that configures how the server will perform hedged reads.
+        It consists of the following keys:
+
+        - ``enabled``: Enables or disables hedged reads in sharded clusters.
+
+        Hedged reads are automatically enabled in MongoDB 4.4+ when using a
+        ``nearest`` read preference. To explicitly enable hedged reads, set
+        the ``enabled`` key  to ``true``::
+
+            >>> Nearest(hedge={'enabled': True})
+
+        To explicitly disable hedged reads, set the ``enabled`` key  to
+        ``False``::
+
+            >>> Nearest(hedge={'enabled': False})
+
+        .. versionadded:: 3.11
+        """
+        return self.__hedge
+
     @property
     def min_wire_version(self):
         """The wire protocol version the server must support.
@@ -158,14 +197,15 @@ def min_wire_version(self):
         return 0 if self.__max_staleness == -1 else 5
 
     def __repr__(self):
-        return "%s(tag_sets=%r, max_staleness=%r)" % (
-            self.name, self.__tag_sets, self.__max_staleness)
+        return "%s(tag_sets=%r, max_staleness=%r, hedge=%r)" % (
+            self.name, self.__tag_sets, self.__max_staleness, self.__hedge)
 
     def __eq__(self, other):
         if isinstance(other, _ServerMode):
             return (self.mode == other.mode and
                     self.tag_sets == other.tag_sets and
-                    self.max_staleness == other.max_staleness)
+                    self.max_staleness == other.max_staleness and
+                    self.hedge == other.hedge)
         return NotImplemented
 
     def __ne__(self, other):
@@ -178,14 +218,16 @@ def __getstate__(self):
         """
         return {'mode': self.__mode,
                 'tag_sets': self.__tag_sets,
-                'max_staleness': self.__max_staleness}
+                'max_staleness': self.__max_staleness,
+                'hedge': self.__hedge}
 
     def __setstate__(self, value):
         """Restore from pickling."""
         self.__mode = value['mode']
         self.__mongos_mode = _MONGOS_MODES[self.__mode]
         self.__tag_sets = _validate_tag_sets(value['tag_sets'])
         self.__max_staleness = _validate_max_staleness(value['max_staleness'])
+        self.__hedge = _validate_hedge(value['hedge'])
 
 
 class Primary(_ServerMode):
@@ -234,14 +276,17 @@ class PrimaryPreferred(_ServerMode):
         replication before it will no longer be selected for operations.
         Default -1, meaning no maximum. If it is set, it must be at least
         90 seconds.
+      - `hedge`: The :attr:`~hedge` to use if the primary is not available.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
     """
 
     __slots__ = ()
 
-    def __init__(self, tag_sets=None, max_staleness=-1):
-        super(PrimaryPreferred, self).__init__(_PRIMARY_PREFERRED,
-                                               tag_sets,
-                                               max_staleness)
+    def __init__(self, tag_sets=None, max_staleness=-1, hedge=None):
+        super(PrimaryPreferred, self).__init__(
+            _PRIMARY_PREFERRED, tag_sets, max_staleness, hedge)
 
     def __call__(self, selection):
         """Apply this read preference to Selection."""
@@ -271,12 +316,17 @@ class Secondary(_ServerMode):
         replication before it will no longer be selected for operations.
         Default -1, meaning no maximum. If it is set, it must be at least
         90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
     """
 
     __slots__ = ()
 
-    def __init__(self, tag_sets=None, max_staleness=-1):
-        super(Secondary, self).__init__(_SECONDARY, tag_sets, max_staleness)
+    def __init__(self, tag_sets=None, max_staleness=-1, hedge=None):
+        super(Secondary, self).__init__(
+            _SECONDARY, tag_sets, max_staleness, hedge)
 
     def __call__(self, selection):
         """Apply this read preference to Selection."""
@@ -303,14 +353,17 @@ class SecondaryPreferred(_ServerMode):
         replication before it will no longer be selected for operations.
         Default -1, meaning no maximum. If it is set, it must be at least
         90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
     """
 
     __slots__ = ()
 
-    def __init__(self, tag_sets=None, max_staleness=-1):
-        super(SecondaryPreferred, self).__init__(_SECONDARY_PREFERRED,
-                                                 tag_sets,
-                                                 max_staleness)
+    def __init__(self, tag_sets=None, max_staleness=-1, hedge=None):
+        super(SecondaryPreferred, self).__init__(
+            _SECONDARY_PREFERRED, tag_sets, max_staleness, hedge)
 
     def __call__(self, selection):
         """Apply this read preference to Selection."""
@@ -342,12 +395,17 @@ class Nearest(_ServerMode):
         replication before it will no longer be selected for operations.
         Default -1, meaning no maximum. If it is set, it must be at least
         90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
     """
 
     __slots__ = ()
 
-    def __init__(self, tag_sets=None, max_staleness=-1):
-        super(Nearest, self).__init__(_NEAREST, tag_sets, max_staleness)
+    def __init__(self, tag_sets=None, max_staleness=-1, hedge=None):
+        super(Nearest, self).__init__(
+            _NEAREST, tag_sets, max_staleness, hedge)
 
     def __call__(self, selection):
         """Apply this read preference to Selection."""
diff --git a/test/test_read_preferences.py b/test/test_read_preferences.py
@@ -46,6 +46,7 @@
 from test.utils import (connected,
                         ignore_deprecations,
                         one,
+                        OvertCommandListener,
                         rs_client,
                         single_client,
                         wait_until)
@@ -567,6 +568,73 @@ def test_read_preference_document(self):
         with self.assertRaises(ValueError):
             Nearest(max_staleness=-2)
 
+    def test_read_preference_document_hedge(self):
+        cases = {
+            'primaryPreferred': PrimaryPreferred,
+            'secondary': Secondary,
+            'secondaryPreferred': SecondaryPreferred,
+            'nearest': Nearest,
+        }
+        for mode, cls in cases.items():
+            with self.assertRaises(TypeError):
+                cls(hedge=[])
+
+            pref = cls(hedge={})
+            self.assertEqual(pref.document, {'mode': mode})
+            out = _maybe_add_read_preference({}, pref)
+            if cls == SecondaryPreferred:
+                # SecondaryPreferred without hedge doesn't add $readPreference.
+                self.assertEqual(out, {})
+            else:
+                self.assertEqual(
+                    out,
+                    SON([("$query", {}), ("$readPreference", pref.document)]))
+
+            hedge = {'enabled': True}
+            pref = cls(hedge=hedge)
+            self.assertEqual(pref.document, {'mode': mode, 'hedge': hedge})
+            out = _maybe_add_read_preference({}, pref)
+            self.assertEqual(
+                out, SON([("$query", {}), ("$readPreference", pref.document)]))
+
+            hedge = {'enabled': False}
+            pref = cls(hedge=hedge)
+            self.assertEqual(pref.document, {'mode': mode, 'hedge': hedge})
+            out = _maybe_add_read_preference({}, pref)
+            self.assertEqual(
+                out, SON([("$query", {}), ("$readPreference", pref.document)]))
+
+            hedge = {'enabled': False, 'extra': 'option'}
+            pref = cls(hedge=hedge)
+            self.assertEqual(pref.document, {'mode': mode, 'hedge': hedge})
+            out = _maybe_add_read_preference({}, pref)
+            self.assertEqual(
+                out, SON([("$query", {}), ("$readPreference", pref.document)]))
+
+    # Require OP_MSG so that $readPreference is visible in the command event.
+    @client_context.require_version_min(3, 6)
+    def test_send_hedge(self):
+        cases = {
+            'primaryPreferred': PrimaryPreferred,
+            'secondary': Secondary,
+            'secondaryPreferred': SecondaryPreferred,
+            'nearest': Nearest,
+        }
+        listener = OvertCommandListener()
+        client = rs_client(event_listeners=[listener])
+        self.addCleanup(client.close)
+        client.admin.command('ping')
+        for mode, cls in cases.items():
+            pref = cls(hedge={'enabled': True})
+            coll = client.test.get_collection('test', read_preference=pref)
+            listener.reset()
+            coll.find_one()
+            started = listener.results['started']
+            self.assertEqual(len(started), 1, started)
+            cmd = started[0].command
+            self.assertIn('$readPreference', cmd)
+            self.assertEqual(cmd['$readPreference'], pref.document)
+
     def test_maybe_add_read_preference(self):
 
         # Primary doesn't add $readPreference
diff --git a/test/test_topology.py b/test/test_topology.py
@@ -788,12 +788,13 @@ def test_no_secondary(self):
 
         self.assertMessage(
             'No replica set members match selector'
-            ' "Secondary(tag_sets=None, max_staleness=-1)"',
+            ' "Secondary(tag_sets=None, max_staleness=-1, hedge=None)"',
             t, ReadPreference.SECONDARY)
 
         self.assertMessage(
             "No replica set members match selector"
-            " \"Secondary(tag_sets=[{'dc': 'ny'}], max_staleness=-1)\"",
+            " \"Secondary(tag_sets=[{'dc': 'ny'}], max_staleness=-1, "
+            "hedge=None)\"",
             t, Secondary(tag_sets=[{'dc': 'ny'}]))
 
     def test_bad_replica_set_name(self):
