Fix connections counter handling BaseException (#457)

Co-authored-by: Amin Alaee <[email protected]>

Author: dswillr

.gitignore

@@ -8,3 +8,4 @@ __pycache__/
 *.egg-info/
 htmlcov/
 venv/
+.idea/

databases/core.py

@@ -251,7 +251,7 @@ async def __aenter__(self) -> "Connection":
             try:
                 if self._connection_counter == 1:
                     await self._connection.acquire()
-            except Exception as e:
+            except BaseException as e:
                 self._connection_counter -= 1
                 raise e
         return self

tests/test_databases.py

@@ -303,24 +303,22 @@ async def test_ddl_queries(database_url):
             await database.execute(query)
 
 
+@pytest.mark.parametrize("exception", [Exception, asyncio.CancelledError])
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @mysql_versions
 @async_adapter
-async def test_queries_after_error(database_url):
+async def test_queries_after_error(database_url, exception):
     """
     Test that the basic `execute()` works after a previous error.
     """
 
-    class DBException(Exception):
-        pass
-
     async with Database(database_url) as database:
         with patch.object(
             database.connection()._connection,
             "acquire",
-            new=AsyncMock(side_effect=DBException),
+            new=AsyncMock(side_effect=exception),
         ):
-            with pytest.raises(DBException):
+            with pytest.raises(exception):
                 query = notes.select()
                 await database.fetch_all(query)