WIP exhaust cursors should return all data in first read

Author: NoahStapp

pymongo/asynchronous/mongo_client.py

@@ -1951,7 +1951,7 @@ async def _cleanup_cursor_lock(
                 # exhausted the result set we *must* close the socket
                 # to stop the server from sending more data.
                 assert conn_mgr.conn is not None
-                conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
+                await conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
             else:
                 await self._close_cursor_now(cursor_id, address, session=session, conn_mgr=conn_mgr)
         if conn_mgr:

pymongo/asynchronous/network.py

@@ -195,7 +195,10 @@ async def command(
             reply = None
             response_doc: _DocumentOut = {"ok": 1}
         else:
-            reply = await async_receive_message(conn, request_id)
+            if "dropDatabase" in spec:
+                reply = await async_receive_message(conn, request_id, debug=True)
+            else:
+                reply = await async_receive_message(conn, request_id)
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields

pymongo/asynchronous/pool.py

@@ -239,7 +239,7 @@ async def unpin(self) -> None:
         if pool:
             await pool.checkin(self)
         else:
-            self.close_conn(ConnectionClosedReason.STALE)
+            await self.close_conn(ConnectionClosedReason.STALE)
 
     def hello_cmd(self) -> dict[str, Any]:
         # Handshake spec requires us to use OP_MSG+hello command for the
@@ -438,7 +438,7 @@ async def command(
             raise
         # Catch socket.error, KeyboardInterrupt, etc. and close ourselves.
         except BaseException as error:
-            self._raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     async def send_message(self, message: bytes, max_doc_size: int) -> None:
         """Send a raw BSON message or raise ConnectionFailure.
@@ -454,7 +454,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
         try:
             await async_sendall(self.conn.get_conn, message)
         except BaseException as error:
-            self._raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
         """Receive a raw BSON message or raise ConnectionFailure.
@@ -464,7 +464,7 @@ async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _O
         try:
             return await async_receive_message(self, request_id, self.max_message_size)
         except BaseException as error:
-            self._raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     def _raise_if_not_writable(self, unacknowledged: bool) -> None:
         """Raise NotPrimaryError on unacknowledged write if this socket is not
@@ -550,11 +550,11 @@ def validate_session(
                     "Can only use session with the AsyncMongoClient that started it"
                 )
 
-    def close_conn(self, reason: Optional[str]) -> None:
+    async def close_conn(self, reason: Optional[str]) -> None:
         """Close this connection with a reason."""
         if self.closed:
             return
-        self._close_conn()
+        await self._close_conn()
         if reason:
             if self.enabled_for_cmap:
                 assert self.listeners is not None
@@ -571,7 +571,7 @@ def close_conn(self, reason: Optional[str]) -> None:
                     error=reason,
                 )
 
-    def _close_conn(self) -> None:
+    async def _close_conn(self) -> None:
         """Close this connection."""
         if self.closed:
             return
@@ -580,7 +580,7 @@ def _close_conn(self) -> None:
         # Note: We catch exceptions to avoid spurious errors on interpreter
         # shutdown.
         try:
-            self.conn.close()
+            await self.conn.close()
         except asyncio.CancelledError:
             raise
         except Exception:  # noqa: S110
@@ -618,7 +618,7 @@ def idle_time_seconds(self) -> float:
         """Seconds since this socket was last checked into its pool."""
         return time.monotonic() - self.last_checkin_time
 
-    def _raise_connection_failure(self, error: BaseException) -> NoReturn:
+    async def _raise_connection_failure(self, error: BaseException) -> NoReturn:
         # Catch *all* exceptions from socket methods and close the socket. In
         # regular Python, socket operations only raise socket.error, even if
         # the underlying cause was a Ctrl-C: a signal raised during socket.recv
@@ -638,7 +638,7 @@ def _raise_connection_failure(self, error: BaseException) -> NoReturn:
             reason = None
         else:
             reason = ConnectionClosedReason.ERROR
-        self.close_conn(reason)
+        await self.close_conn(reason)
         # SSLError from PyOpenSSL inherits directly from Exception.
         if isinstance(error, (IOError, OSError, SSLError)):
             details = _get_timeout_details(self.opts)
@@ -864,7 +864,7 @@ async def _reset(
         # publishing the PoolClearedEvent.
         if close:
             for conn in sockets:
-                conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
+                await conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
             if self.enabled_for_cmap:
                 assert listeners is not None
                 listeners.publish_pool_closed(self.address)
@@ -895,7 +895,7 @@ async def _reset(
                         serviceId=service_id,
                     )
             for conn in sockets:
-                conn.close_conn(ConnectionClosedReason.STALE)
+                await conn.close_conn(ConnectionClosedReason.STALE)
 
     async def update_is_writable(self, is_writable: Optional[bool]) -> None:
         """Updates the is_writable attribute on all sockets currently in the
@@ -940,7 +940,7 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     and self.conns[-1].idle_time_seconds() > self.opts.max_idle_time_seconds
                 ):
                     conn = self.conns.pop()
-                    conn.close_conn(ConnectionClosedReason.IDLE)
+                    await conn.close_conn(ConnectionClosedReason.IDLE)
 
         while True:
             async with self.size_cond:
@@ -964,7 +964,7 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     # Close connection and return if the pool was reset during
                     # socket creation or while acquiring the pool lock.
                     if self.gen.get_overall() != reference_generation:
-                        conn.close_conn(ConnectionClosedReason.STALE)
+                        await conn.close_conn(ConnectionClosedReason.STALE)
                         return
                     self.conns.appendleft(conn)
                     self.active_contexts.discard(conn.cancel_context)
@@ -1052,7 +1052,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         except BaseException:
             async with self.lock:
                 self.active_contexts.discard(conn.cancel_context)
-            conn.close_conn(ConnectionClosedReason.ERROR)
+            await conn.close_conn(ConnectionClosedReason.ERROR)
             raise
 
         return conn
@@ -1246,7 +1246,7 @@ async def _get_conn(
                     except IndexError:
                         self._pending += 1
                 if conn:  # We got a socket from the pool
-                    if self._perished(conn):
+                    if await self._perished(conn):
                         conn = None
                         continue
                 else:  # We need to create a new connection
@@ -1259,7 +1259,7 @@ async def _get_conn(
         except BaseException:
             if conn:
                 # We checked out a socket but authentication failed.
-                conn.close_conn(ConnectionClosedReason.ERROR)
+                await conn.close_conn(ConnectionClosedReason.ERROR)
             async with self.size_cond:
                 self.requests -= 1
                 if incremented:
@@ -1319,7 +1319,7 @@ async def checkin(self, conn: AsyncConnection) -> None:
             await self.reset_without_pause()
         else:
             if self.closed:
-                conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
+                await conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
             elif conn.closed:
                 # CMAP requires the closed event be emitted after the check in.
                 if self.enabled_for_cmap:
@@ -1343,7 +1343,7 @@ async def checkin(self, conn: AsyncConnection) -> None:
                     # Hold the lock to ensure this section does not race with
                     # Pool.reset().
                     if self.stale_generation(conn.generation, conn.service_id):
-                        conn.close_conn(ConnectionClosedReason.STALE)
+                        await conn.close_conn(ConnectionClosedReason.STALE)
                     else:
                         conn.update_last_checkin_time()
                         conn.update_is_writable(bool(self.is_writable))
@@ -1361,7 +1361,7 @@ async def checkin(self, conn: AsyncConnection) -> None:
             self.operation_count -= 1
             self.size_cond.notify()
 
-    def _perished(self, conn: AsyncConnection) -> bool:
+    async def _perished(self, conn: AsyncConnection) -> bool:
         """Return True and close the connection if it is "perished".
 
         This side-effecty function checks if this socket has been idle for
@@ -1381,18 +1381,18 @@ def _perished(self, conn: AsyncConnection) -> bool:
             self.opts.max_idle_time_seconds is not None
             and idle_time_seconds > self.opts.max_idle_time_seconds
         ):
-            conn.close_conn(ConnectionClosedReason.IDLE)
+            await conn.close_conn(ConnectionClosedReason.IDLE)
             return True
 
         if self._check_interval_seconds is not None and (
             self._check_interval_seconds == 0 or idle_time_seconds > self._check_interval_seconds
         ):
             if conn.conn_closed():
-                conn.close_conn(ConnectionClosedReason.ERROR)
+                await conn.close_conn(ConnectionClosedReason.ERROR)
                 return True
 
         if self.stale_generation(conn.generation, conn.service_id):
-            conn.close_conn(ConnectionClosedReason.STALE)
+            await conn.close_conn(ConnectionClosedReason.STALE)
             return True
 
         return False
@@ -1436,9 +1436,9 @@ def _raise_wait_queue_timeout(self, checkout_started_time: float) -> NoReturn:
             f"maxPoolSize: {self.opts.max_pool_size}, timeout: {timeout}"
         )
 
-    def __del__(self) -> None:
-        # Avoid ResourceWarnings in Python 3
-        # Close all sockets without calling reset() or close() because it is
-        # not safe to acquire a lock in __del__.
-        for conn in self.conns:
-            conn.close_conn(None)
+    # def __del__(self) -> None:
+    #     # Avoid ResourceWarnings in Python 3
+    #     # Close all sockets without calling reset() or close() because it is
+    #     # not safe to acquire a lock in __del__.
+    #     for conn in self.conns:
+    #         conn.close_conn(None)