Handle asyncio.CancelledError (now a BaseException)

Author: 1st1

asyncpg/connect_utils.py

@@ -541,7 +541,7 @@ async def _connect_addr(*, addr, loop, timeout, params, config,
         if timeout <= 0:
             raise asyncio.TimeoutError
         await asyncio.wait_for(connected, loop=loop, timeout=timeout)
-    except Exception:
+    except (Exception, asyncio.CancelledError):
         tr.close()
         raise
 
@@ -614,7 +614,7 @@ async def _negotiate_ssl_connection(host, port, conn_factory, *, loop, ssl,
 
     try:
         return await conn_factory(sock=sock)  # Must come after tr.close()
-    except Exception:
+    except (Exception, asyncio.CancelledError):
         sock.close()
         raise
 

asyncpg/connection.py

@@ -1080,7 +1080,7 @@ async def close(self, *, timeout=None):
         try:
             if not self.is_closed():
                 await self._protocol.close(timeout)
-        except Exception:
+        except (Exception, asyncio.CancelledError):
             # If we fail to close gracefully, abort the connection.
             self._abort()
             raise
@@ -1213,7 +1213,7 @@ async def _cancel(self, waiter):
             # the CancelledError, and don't want the loop to warn about
             # an unretrieved exception.
             pass
-        except Exception as ex:
+        except (Exception, asyncio.CancelledError) as ex:
             if not waiter.done():
                 waiter.set_exception(ex)
         finally:

asyncpg/pool.py

@@ -146,7 +146,7 @@ async def acquire(self) -> PoolConnectionProxy:
         if self._setup is not None:
             try:
                 await self._setup(proxy)
-            except Exception as ex:
+            except (Exception, asyncio.CancelledError) as ex:
                 # If a user-defined `setup` function fails, we don't
                 # know if the connection is safe for re-use, hence
                 # we close it.  A new connection will be created
@@ -204,7 +204,7 @@ async def release(self, timeout):
                     budget -= time.monotonic() - started
 
             await self._con.reset(timeout=budget)
-        except Exception as ex:
+        except (Exception, asyncio.CancelledError) as ex:
             # If the `reset` call failed, terminate the connection.
             # A new one will be created when `acquire` is called
             # again.
@@ -480,7 +480,7 @@ async def _get_new_connection(self):
         if self._init is not None:
             try:
                 await self._init(con)
-            except Exception as ex:
+            except (Exception, asyncio.CancelledError) as ex:
                 # If a user-defined `init` function fails, we don't
                 # know if the connection is safe for re-use, hence
                 # we close it.  A new connection will be created
@@ -587,7 +587,7 @@ async def _acquire_impl():
             ch = await self._queue.get()  # type: PoolConnectionHolder
             try:
                 proxy = await ch.acquire()  # type: PoolConnectionProxy
-            except Exception:
+            except (Exception, asyncio.CancelledError):
                 self._queue.put_nowait(ch)
                 raise
             else:
@@ -679,7 +679,7 @@ async def close(self):
                 ch.close() for ch in self._holders]
             await asyncio.gather(*close_coros, loop=self._loop)
 
-        except Exception:
+        except (Exception, asyncio.CancelledError):
             self.terminate()
             raise
 

asyncpg/protocol/protocol.pyx

@@ -350,7 +350,7 @@ cdef class BaseProtocol(CoreProtocol):
                                 sink(buffer),
                                 timeout=timer.get_remaining_budget(),
                                 loop=self.loop)
-                    except Exception as ex:
+                    except (Exception, asyncio.CancelledError) as ex:
                         # Abort the COPY operation on any error in
                         # output sink.
                         self._request_cancel()
@@ -476,7 +476,7 @@ cdef class BaseProtocol(CoreProtocol):
             else:
                 raise apg_exc.InternalClientError('TimoutError was not raised')
 
-        except Exception as e:
+        except (Exception, asyncio.CancelledError) as e:
             self._write_copy_fail_msg(str(e))
             self._request_cancel()
             # Make asyncio shut up about unretrieved QueryCanceledError