PYTHON-2390 - Retryable reads use the same implicit session

Author: NoahStapp

pymongo/asynchronous/mongo_client.py

@@ -2048,17 +2048,18 @@ async def _retryable_read(
         retryable = bool(
             retryable and self.options.retry_reads and not (session and session.in_transaction)
         )
-        return await self._retry_internal(
-            func,
-            session,
-            None,
-            operation,
-            is_read=True,
-            address=address,
-            read_pref=read_pref,
-            retryable=retryable,
-            operation_id=operation_id,
-        )
+        async with self._tmp_session(session) as s:
+            return await self._retry_internal(
+                func,
+                s,
+                None,
+                operation,
+                is_read=True,
+                address=address,
+                read_pref=read_pref,
+                retryable=retryable,
+                operation_id=operation_id,
+            )
 
     async def _retryable_write(
         self,

pymongo/synchronous/mongo_client.py

@@ -2044,17 +2044,18 @@ def _retryable_read(
         retryable = bool(
             retryable and self.options.retry_reads and not (session and session.in_transaction)
         )
-        return self._retry_internal(
-            func,
-            session,
-            None,
-            operation,
-            is_read=True,
-            address=address,
-            read_pref=read_pref,
-            retryable=retryable,
-            operation_id=operation_id,
-        )
+        with self._tmp_session(session) as s:
+            return self._retry_internal(
+                func,
+                s,
+                None,
+                operation,
+                is_read=True,
+                address=address,
+                read_pref=read_pref,
+                retryable=retryable,
+                operation_id=operation_id,
+            )
 
     def _retryable_write(
         self,

test/asynchronous/test_retryable_reads.py

@@ -218,6 +218,38 @@ async def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are
         #  Assert that both events occurred on the same mongos.
         assert listener.succeeded_events[0].connection_id == listener.failed_events[0].connection_id
 
+    @async_client_context.require_failCommand_fail_point
+    async def test_retryable_reads_are_retried_on_the_same_implicit_session(self):
+        fail_command = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": 1},
+            "data": {"failCommands": ["count"], "errorCode": 6},
+        }
+
+        listener = OvertCommandListener()
+        client = await self.async_rs_or_single_client(
+            directConnection=False,
+            event_listeners=[listener],
+            retryReads=True,
+        )
+
+        await async_set_fail_point(client, fail_command)
+
+        await client.t.t.estimated_document_count()
+
+        # Disable failpoint.
+        fail_command["mode"] = "off"
+        await async_set_fail_point(client, fail_command)
+
+        #  Assert that both events occurred on the same session.
+        lsids = [
+            event.command["lsid"]
+            for event in listener.started_events
+            if event.command_name == "count"
+        ]
+        self.assertEqual(len(lsids), 2)
+        self.assertEqual(lsids[0], lsids[1])
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_retryable_reads.py

@@ -216,6 +216,38 @@ def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are_avail
         #  Assert that both events occurred on the same mongos.
         assert listener.succeeded_events[0].connection_id == listener.failed_events[0].connection_id
 
+    @client_context.require_failCommand_fail_point
+    def test_retryable_reads_are_retried_on_the_same_implicit_session(self):
+        fail_command = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": 1},
+            "data": {"failCommands": ["count"], "errorCode": 6},
+        }
+
+        listener = OvertCommandListener()
+        client = self.rs_or_single_client(
+            directConnection=False,
+            event_listeners=[listener],
+            retryReads=True,
+        )
+
+        set_fail_point(client, fail_command)
+
+        client.t.t.estimated_document_count()
+
+        # Disable failpoint.
+        fail_command["mode"] = "off"
+        set_fail_point(client, fail_command)
+
+        #  Assert that both events occurred on the same session.
+        lsids = [
+            event.command["lsid"]
+            for event in listener.started_events
+            if event.command_name == "count"
+        ]
+        self.assertEqual(len(lsids), 2)
+        self.assertEqual(lsids[0], lsids[1])
+
 
 if __name__ == "__main__":
     unittest.main()