PYTHON-1674 Support retryable reads

Add retryReads URI option that defaults to True.
Supported read operations will be retried once after transient
network, election, and shutdown errors on MongoDB 3.6+.
Supported operations are:
listCollections, listIndexes, and listDatabases
distinct
count, estimated_document_count, count_documents
aggregate (not including $out)
find (only for the initial find command, getMore commands are not
retried).
ChangeStreams: watch (initial aggregate command).
GridFS read APIs.

Test changes:
Add retryable reads spec test runner.
Disable retryable reads in network error tests.

Author: ShaneHarvey

doc/changelog.rst

@@ -33,6 +33,13 @@ Version 3.9 adds support for MongoDB 4.2. Highlights include:
 - The ``retryWrites`` URI option now defaults to ``True``. Supported write
   operations that fail with a retryable error will automatically be retried one
   time, with at-most-once semantics.
+- Support for retryable reads and the ``retryReads`` URI option which is
+  enabled by default. See the :class:`~pymongo.mongo_client.MongoClient`
+  documentation for details.
+
+  Now that supported operations are retried automatically and transparently,
+  users should consider adjusting any custom retry logic to prevent
+  an application from inadvertently retrying for too long.
 
 .. _URI options specification: https://github.com/mongodb/specifications/blob/master/source/uri-options/uri-options.rst
 

pymongo/change_stream.py

@@ -118,8 +118,8 @@ def _run_aggregation_cmd(self, session, explicit_session):
         """
         read_preference = self._target._read_preference_for(session)
         client = self._database.client
-        with client._socket_for_reads(
-                read_preference, session) as (sock_info, slave_ok):
+
+        def _cmd(session, server, sock_info, slave_ok):
             pipeline = self._full_pipeline()
             cmd = SON([("aggregate", self._aggregation_target),
                        ("pipeline", pipeline),
@@ -160,6 +160,8 @@ def _run_aggregation_cmd(self, session, explicit_session):
                 max_await_time_ms=self._max_await_time_ms,
                 session=session, explicit_session=explicit_session)
 
+        return client._retryable_read(_cmd, read_preference, session)
+
     def _create_cursor(self):
         with self._database.client._tmp_session(self._session, close=False) as s:
             return self._run_aggregation_cmd(

pymongo/client_options.py

@@ -164,6 +164,7 @@ def __init__(self, username, password, database, options):
         self.__heartbeat_frequency = options.get(
             'heartbeatfrequencyms', common.HEARTBEAT_FREQUENCY)
         self.__retry_writes = options.get('retrywrites', common.RETRY_WRITES)
+        self.__retry_reads = options.get('retryreads', common.RETRY_READS)
         self.__server_selector = options.get(
             'server_selector', any_server_selector)
 
@@ -235,3 +236,8 @@ def read_concern(self):
     def retry_writes(self):
         """If this instance should retry supported write operations."""
         return self.__retry_writes
+
+    @property
+    def retry_reads(self):
+        """If this instance should retry supported read operations."""
+        return self.__retry_reads

pymongo/collection.py

@@ -188,12 +188,6 @@ def _socket_for_reads(self, session):
         return self.__database.client._socket_for_reads(
             self._read_preference_for(session), session)
 
-    def _socket_for_primary_reads(self, session):
-        read_pref = ((session and session._txn_read_preference())
-                     or ReadPreference.PRIMARY)
-        return self.__database.client._socket_for_reads(
-            read_pref, session), read_pref
-
     def _socket_for_writes(self, session):
         return self.__database.client._socket_for_writes(session)
 
@@ -1572,7 +1566,7 @@ def parallel_scan(self, num_cursors, session=None, **kwargs):
 
     def _count(self, cmd, collation=None, session=None):
         """Internal count helper."""
-        with self._socket_for_reads(session) as (sock_info, slave_ok):
+        def _cmd(session, server, sock_info, slave_ok):
             res = self._command(
                 sock_info,
                 cmd,
@@ -1582,9 +1576,12 @@ def _count(self, cmd, collation=None, session=None):
                 read_concern=self.read_concern,
                 collation=collation,
                 session=session)
-        if res.get("errmsg", "") == "ns missing":
-            return 0
-        return int(res["n"])
+            if res.get("errmsg", "") == "ns missing":
+                return 0
+            return int(res["n"])
+
+        return self.__database.client._retryable_read(
+            _cmd, self._read_preference_for(session), session)
 
     def _aggregate_one_result(
             self, sock_info, slave_ok, cmd, collation=None, session=None):
@@ -1693,12 +1690,16 @@ def count_documents(self, filter, session=None, **kwargs):
             kwargs["hint"] = helpers._index_document(kwargs["hint"])
         collation = validate_collation_or_none(kwargs.pop('collation', None))
         cmd.update(kwargs)
-        with self._socket_for_reads(session) as (sock_info, slave_ok):
+
+        def _cmd(session, server, sock_info, slave_ok):
             result = self._aggregate_one_result(
                 sock_info, slave_ok, cmd, collation, session)
-        if not result:
-            return 0
-        return result['n']
+            if not result:
+                return 0
+            return result['n']
+
+        return self.__database.client._retryable_read(
+            _cmd, self._read_preference_for(session), session)
 
     def count(self, filter=None, session=None, **kwargs):
         """**DEPRECATED** - Get the number of documents in this collection.
@@ -2149,8 +2150,10 @@ def list_indexes(self, session=None):
         codec_options = CodecOptions(SON)
         coll = self.with_options(codec_options=codec_options,
                                  read_preference=ReadPreference.PRIMARY)
-        sock_ctx, read_pref = self._socket_for_primary_reads(session)
-        with sock_ctx as (sock_info, slave_ok):
+        read_pref = ((session and session._txn_read_preference())
+                     or ReadPreference.PRIMARY)
+
+        def _cmd(session, server, sock_info, slave_ok):
             cmd = SON([("listIndexes", self.__name), ("cursor", {})])
             if sock_info.max_wire_version > 2:
                 with self.__database.client._tmp_session(session, False) as s:
@@ -2179,6 +2182,9 @@ def list_indexes(self, session=None):
                 # will never be a getMore call.
                 return CommandCursor(coll, cursor, sock_info.address)
 
+        return self.__database.client._retryable_read(
+            _cmd, read_pref, session)
+
     def index_information(self, session=None):
         """Get information on this collection's indexes.
 
@@ -2275,10 +2281,11 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
                 "useCursor", kwargs.pop("useCursor"))
         batch_size = common.validate_non_negative_integer_or_none(
             "batchSize", kwargs.pop("batchSize", None))
+
+        dollar_out = pipeline and '$out' in pipeline[-1]
         # If the server does not support the "cursor" option we
         # ignore useCursor and batchSize.
-        with self._socket_for_reads(session) as (sock_info, slave_ok):
-            dollar_out = pipeline and '$out' in pipeline[-1]
+        def _cmd(session, server, sock_info, slave_ok):
             if use_cursor:
                 if "cursor" not in kwargs:
                     kwargs["cursor"] = {}
@@ -2336,6 +2343,10 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
                 max_await_time_ms=max_await_time_ms,
                 session=session, explicit_session=explicit_session)
 
+        return self.__database.client._retryable_read(
+            _cmd, self._read_preference_for(session), session,
+            retryable=not dollar_out)
+
     def aggregate(self, pipeline, session=None, **kwargs):
         """Perform an aggregation using the aggregation framework on this
         collection.
@@ -2681,12 +2692,53 @@ def distinct(self, key, filter=None, session=None, **kwargs):
             kwargs["query"] = filter
         collation = validate_collation_or_none(kwargs.pop('collation', None))
         cmd.update(kwargs)
-        with self._socket_for_reads(session) as (sock_info, slave_ok):
-            return self._command(sock_info, cmd, slave_ok,
-                                 read_concern=self.read_concern,
-                                 collation=collation,
-                                 session=session,
-                                 user_fields={"values": 1})["values"]
+        def _cmd(session, server, sock_info, slave_ok):
+            return self._command(
+                sock_info, cmd, slave_ok, read_concern=self.read_concern,
+                collation=collation, session=session,
+                user_fields={"values": 1})["values"]
+
+        return self.__database.client._retryable_read(
+            _cmd, self._read_preference_for(session), session)
+
+    def _map_reduce(self, map, reduce, out, session, read_pref, **kwargs):
+        """Internal mapReduce helper."""
+        cmd = SON([("mapReduce", self.__name),
+                   ("map", map),
+                   ("reduce", reduce),
+                   ("out", out)])
+        collation = validate_collation_or_none(kwargs.pop('collation', None))
+        cmd.update(kwargs)
+
+        inline = 'inline' in out
+
+        if inline:
+            user_fields = {'results': 1}
+        else:
+            user_fields = None
+
+        read_pref = ((session and session._txn_read_preference())
+                     or read_pref)
+
+        with self.__database.client._socket_for_reads(read_pref, session) as (
+                sock_info, slave_ok):
+            if (sock_info.max_wire_version >= 4 and
+                    ('readConcern' not in cmd) and
+                    inline):
+                read_concern = self.read_concern
+            else:
+                read_concern = None
+            if 'writeConcern' not in cmd and not inline:
+                write_concern = self._write_concern_for(session)
+            else:
+                write_concern = None
+
+            return self._command(
+                sock_info, cmd, slave_ok, read_pref,
+                read_concern=read_concern,
+                write_concern=write_concern,
+                collation=collation, session=session,
+                user_fields=user_fields)
 
     def map_reduce(self, map, reduce, out, full_response=False, session=None,
                    **kwargs):
@@ -2747,36 +2799,8 @@ def map_reduce(self, map, reduce, out, full_response=False, session=None,
             raise TypeError("'out' must be an instance of "
                             "%s or a mapping" % (string_type.__name__,))
 
-        cmd = SON([("mapreduce", self.__name),
-                   ("map", map),
-                   ("reduce", reduce),
-                   ("out", out)])
-        collation = validate_collation_or_none(kwargs.pop('collation', None))
-        cmd.update(kwargs)
-
-        inline = 'inline' in cmd['out']
-        sock_ctx, read_pref = self._socket_for_primary_reads(session)
-        with sock_ctx as (sock_info, slave_ok):
-            if (sock_info.max_wire_version >= 4 and 'readConcern' not in cmd and
-                    inline):
-                read_concern = self.read_concern
-            else:
-                read_concern = None
-            if 'writeConcern' not in cmd and not inline:
-                write_concern = self._write_concern_for(session)
-            else:
-                write_concern = None
-            if inline:
-                user_fields = {'results': 1}
-            else:
-                user_fields = None
-
-            response = self._command(
-                sock_info, cmd, slave_ok, read_pref,
-                read_concern=read_concern,
-                write_concern=write_concern,
-                collation=collation, session=session,
-                user_fields=user_fields)
+        response = self._map_reduce(map, reduce, out, session,
+                                    ReadPreference.PRIMARY, **kwargs)
 
         if full_response or not response.get('result'):
             return response
@@ -2822,23 +2846,8 @@ def inline_map_reduce(self, map, reduce, full_response=False, session=None,
            Added the `collation` option.
 
         """
-        cmd = SON([("mapreduce", self.__name),
-                   ("map", map),
-                   ("reduce", reduce),
-                   ("out", {"inline": 1})])
-        user_fields = {'results': 1}
-        collation = validate_collation_or_none(kwargs.pop('collation', None))
-        cmd.update(kwargs)
-        with self._socket_for_reads(session) as (sock_info, slave_ok):
-            if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
-                res = self._command(sock_info, cmd, slave_ok,
-                                    read_concern=self.read_concern,
-                                    collation=collation, session=session,
-                                    user_fields=user_fields)
-            else:
-                res = self._command(sock_info, cmd, slave_ok,
-                                    collation=collation, session=session,
-                                    user_fields=user_fields)
+        res = self._map_reduce(map, reduce, {"inline": 1}, session,
+                               self.read_preference, **kwargs)
 
         if full_response:
             return res

pymongo/command_cursor.py

@@ -128,9 +128,8 @@ def kill():
 
         client = self.__collection.database.client
         try:
-            response = client._send_message_with_response(
-                operation, address=self.__address,
-                unpack_res=self._unpack_response)
+            response = client._run_operation_with_response(
+                operation, self._unpack_response, address=self.__address)
         except OperationFailure:
             kill()
             raise

pymongo/common.py

@@ -91,6 +91,9 @@
 # Default value for retryWrites.
 RETRY_WRITES = True
 
+# Default value for retryReads.
+RETRY_READS = True
+
 # mongod/s 2.6 and above return code 59 when a command doesn't exist.
 COMMAND_NOT_FOUND_CODES = (59,)
 
@@ -569,6 +572,7 @@ def validate_tzinfo(dummy, value):
     'readpreference': validate_read_preference_mode,
     'readpreferencetags': validate_read_preference_tags,
     'replicaset': validate_string_or_none,
+    'retryreads': validate_boolean_or_string,
     'retrywrites': validate_boolean_or_string,
     'serverselectiontimeoutms': validate_timeout_or_zero,
     'sockettimeoutms': validate_timeout_or_none,

pymongo/cursor.py

@@ -937,10 +937,11 @@ def __send_message(self, operation):
         Can raise ConnectionFailure.
         """
         client = self.__collection.database.client
+
         try:
-            response = client._send_message_with_response(
-                operation, exhaust=self.__exhaust, address=self.__address,
-                unpack_res=self._unpack_response)
+            response = client._run_operation_with_response(
+                operation, self._unpack_response, exhaust=self.__exhaust,
+                address=self.__address)
         except OperationFailure:
             self.__killed = True
 

pymongo/database.py

@@ -657,6 +657,22 @@ def command(self, command, value=1, check=True,
                                  check, allowable_errors, read_preference,
                                  codec_options, session=session, **kwargs)
 
+    def _retryable_read_command(self, command, value=1, check=True,
+                allowable_errors=None, read_preference=None,
+                codec_options=DEFAULT_CODEC_OPTIONS, session=None, **kwargs):
+        """Same as command but used for retryable read commands."""
+        if read_preference is None:
+            read_preference = ((session and session._txn_read_preference())
+                               or ReadPreference.PRIMARY)
+
+        def _cmd(session, server, sock_info, slave_ok):
+            return self._command(sock_info, command, slave_ok, value,
+                                 check, allowable_errors, read_preference,
+                                 codec_options, session=session, **kwargs)
+
+        return self.__client._retryable_read(
+            _cmd, read_preference, session)
+
     def _list_collections(self, sock_info, slave_okay, session,
                           read_preference, **kwargs):
         """Internal listCollections helper."""
@@ -718,12 +734,15 @@ def list_collections(self, session=None, filter=None, **kwargs):
             kwargs['filter'] = filter
         read_pref = ((session and session._txn_read_preference())
                      or ReadPreference.PRIMARY)
-        with self.__client._socket_for_reads(
-                read_pref, session) as (sock_info, slave_okay):
+
+        def _cmd(session, server, sock_info, slave_okay):
             return self._list_collections(
                 sock_info, slave_okay, session, read_preference=read_pref,
                 **kwargs)
 
+        return self.__client._retryable_read(
+            _cmd, read_pref, session)
+
     def list_collection_names(self, session=None, filter=None, **kwargs):
         """Get a list of all the collection names in this database.