Add comments to catching BaseException, speed up tests

Author: NoahStapp

pymongo/asynchronous/change_stream.py

@@ -391,6 +391,7 @@ async def try_next(self) -> Optional[_DocumentType]:
             if not _resumable(exc) and not exc.timeout:
                 await self.close()
             raise
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             await self.close()
             raise

pymongo/asynchronous/client_session.py

@@ -697,6 +697,7 @@ async def callback(session, custom_arg, custom_kwarg=None):
             )
             try:
                 ret = await callback(self)
+            # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
             except BaseException as exc:
                 if self.in_transaction:
                     await self.abort_transaction()

pymongo/asynchronous/cursor.py

@@ -1126,6 +1126,7 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
             self._killed = True
             await self.close()
             raise
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             await self.close()
             raise

pymongo/asynchronous/pool.py

@@ -559,7 +559,7 @@ async def command(
             )
         except (OperationFailure, NotPrimaryError):
             raise
-        # Catch socket.error, KeyboardInterrupt, etc. and close ourselves.
+        # Catch socket.error, KeyboardInterrupt, CancelledError, etc. and close ourselves.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -576,6 +576,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
 
         try:
             await async_sendall(self.conn, message)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -586,6 +587,7 @@ async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _O
         """
         try:
             return await receive_message(self, request_id, self.max_message_size)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -1269,6 +1271,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
 
         try:
             sock = await _configured_socket(self.address, self.opts)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             async with self.lock:
                 self.active_contexts.discard(tmp_context)
@@ -1308,6 +1311,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
                 handler.contribute_socket(conn, completed_handshake=False)
 
             await conn.authenticate()
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             async with self.lock:
                 self.active_contexts.discard(conn.cancel_context)
@@ -1369,6 +1373,7 @@ async def checkout(
             async with self.lock:
                 self.active_contexts.add(conn.cancel_context)
             yield conn
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             # Exception in caller. Ensure the connection gets returned.
             # Note that when pinned is True, the session owns the
@@ -1515,6 +1520,7 @@ async def _get_conn(
                         async with self._max_connecting_cond:
                             self._pending -= 1
                             self._max_connecting_cond.notify()
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             if conn:
                 # We checked out a socket but authentication failed.

pymongo/periodic_executor.py

@@ -98,6 +98,7 @@ async def _run(self) -> None:
                 if not await self._target():
                     self._stopped = True
                     break
+            # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
             except BaseException:
                 self._stopped = True
                 raise
@@ -230,6 +231,7 @@ def _run(self) -> None:
                 if not self._target():
                     self._stopped = True
                     break
+            # Catch KeyboardInterrupt, etc. and cleanup.
             except BaseException:
                 with self._lock:
                     self._stopped = True

pymongo/synchronous/change_stream.py

@@ -389,6 +389,7 @@ def try_next(self) -> Optional[_DocumentType]:
             if not _resumable(exc) and not exc.timeout:
                 self.close()
             raise
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             self.close()
             raise

pymongo/synchronous/client_session.py

@@ -694,6 +694,7 @@ def callback(session, custom_arg, custom_kwarg=None):
             self.start_transaction(read_concern, write_concern, read_preference, max_commit_time_ms)
             try:
                 ret = callback(self)
+            # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
             except BaseException as exc:
                 if self.in_transaction:
                     self.abort_transaction()

pymongo/synchronous/cursor.py

@@ -1124,6 +1124,7 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
             self._killed = True
             self.close()
             raise
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             self.close()
             raise

pymongo/synchronous/pool.py

@@ -559,7 +559,7 @@ def command(
             )
         except (OperationFailure, NotPrimaryError):
             raise
-        # Catch socket.error, KeyboardInterrupt, etc. and close ourselves.
+        # Catch socket.error, KeyboardInterrupt, CancelledError, etc. and close ourselves.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -576,6 +576,7 @@ def send_message(self, message: bytes, max_doc_size: int) -> None:
 
         try:
             sendall(self.conn, message)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -586,6 +587,7 @@ def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
         """
         try:
             return receive_message(self, request_id, self.max_message_size)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -1263,6 +1265,7 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
 
         try:
             sock = _configured_socket(self.address, self.opts)
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             with self.lock:
                 self.active_contexts.discard(tmp_context)
@@ -1302,6 +1305,7 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
                 handler.contribute_socket(conn, completed_handshake=False)
 
             conn.authenticate()
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             with self.lock:
                 self.active_contexts.discard(conn.cancel_context)
@@ -1363,6 +1367,7 @@ def checkout(
             with self.lock:
                 self.active_contexts.add(conn.cancel_context)
             yield conn
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             # Exception in caller. Ensure the connection gets returned.
             # Note that when pinned is True, the session owns the
@@ -1509,6 +1514,7 @@ def _get_conn(
                         with self._max_connecting_cond:
                             self._pending -= 1
                             self._max_connecting_cond.notify()
+        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException:
             if conn:
                 # We checked out a socket but authentication failed.

test/asynchronous/test_async_cancellation.py

@@ -26,15 +26,14 @@
 
 class TestAsyncCancellation(AsyncIntegrationTest):
     async def test_async_cancellation_closes_connection(self):
-        client = await self.async_rs_or_single_client(maxPoolSize=1)
-        pool = await async_get_pool(client)
-        await connected(client)
+        pool = await async_get_pool(self.client)
+        await connected(self.client)
         conn = one(pool.conns)
-        await client.db.test.insert_one({"x": 1})
-        self.addAsyncCleanup(client.db.test.drop)
+        await self.client.db.test.insert_one({"x": 1})
+        self.addAsyncCleanup(self.client.db.test.drop)
 
         async def task():
-            await client.db.test.find_one({"$where": delay(0.2)})
+            await self.client.db.test.find_one({"$where": delay(0.2)})
 
         task = asyncio.create_task(task())
 
@@ -48,15 +47,14 @@ async def task():
 
     @async_client_context.require_transactions
     async def test_async_cancellation_aborts_transaction(self):
-        client = await self.async_rs_or_single_client()
-        await connected(client)
-        await client.db.test.insert_one({"x": 1})
-        self.addAsyncCleanup(client.db.test.drop)
+        await connected(self.client)
+        await self.client.db.test.insert_one({"x": 1})
+        self.addAsyncCleanup(self.client.db.test.drop)
 
-        session = client.start_session()
+        session = self.client.start_session()
 
         async def callback(session):
-            await client.db.test.find_one({"$where": delay(0.2)}, session=session)
+            await self.client.db.test.find_one({"$where": delay(0.2)}, session=session)
 
         async def task():
             await session.with_transaction(callback)
@@ -73,13 +71,12 @@ async def task():
 
     @async_client_context.require_failCommand_blockConnection
     async def test_async_cancellation_closes_cursor(self):
-        client = await self.async_rs_or_single_client()
-        await connected(client)
+        await connected(self.client)
         for _ in range(2):
-            await client.db.test.insert_one({"x": 1})
-        self.addAsyncCleanup(client.db.test.drop)
+            await self.client.db.test.insert_one({"x": 1})
+        self.addAsyncCleanup(self.client.db.test.drop)
 
-        cursor = client.db.test.find({}, batch_size=1)
+        cursor = self.client.db.test.find({}, batch_size=1)
         await cursor.next()
 
         # Make sure getMore commands block
@@ -106,11 +103,10 @@ async def task():
     @async_client_context.require_change_streams
     @async_client_context.require_failCommand_blockConnection
     async def test_async_cancellation_closes_change_stream(self):
-        client = await self.async_rs_or_single_client()
-        await connected(client)
-        self.addAsyncCleanup(client.db.test.drop)
+        await connected(self.client)
+        self.addAsyncCleanup(self.client.db.test.drop)
 
-        change_stream = await client.db.test.watch(batch_size=2)
+        change_stream = await self.client.db.test.watch(batch_size=2)
 
         # Make sure getMore commands block
         fail_command = {
@@ -122,7 +118,7 @@ async def test_async_cancellation_closes_change_stream(self):
         async def task():
             async with self.fail_point(fail_command):
                 for _ in range(2):
-                    await client.db.test.insert_one({"x": 1})
+                    await self.client.db.test.insert_one({"x": 1})
                 await change_stream.next()
 
         task = asyncio.create_task(task())