test: Ensure proper async connection cleanup on DB test exit (pola-rs#25766)

Author: alexander-beedie

py-polars/src/polars/io/database/_executor.py

@@ -456,6 +456,16 @@ async def _sqlalchemy_async_execute(self, query: TextClause, **options: Any) ->
         """Execute a query using an async SQLAlchemy connection."""
         is_session = self._is_alchemy_session(self.cursor)
         cursor = self.cursor.begin() if is_session else self.cursor  # type: ignore[attr-defined]
+
+        # check if connection is already started (eg: user awaited `engine.connect()`);
+        # if so, use it directly without entering the context manager again
+        try:
+            if object.__getattribute__(cursor, "sync_connection") is not None:
+                result = await cursor.execute(query, **options)
+                return result
+        except AttributeError:
+            pass
+
         async with cursor as conn:  # type: ignore[union-attr]
             if is_session and not hasattr(conn, "execute"):
                 conn = conn.session

py-polars/tests/unit/io/database/test_async.py

@@ -19,6 +19,7 @@
     from collections.abc import Iterable
     from pathlib import Path
 
+
 SURREAL_MOCK_DATA: list[dict[str, Any]] = [
     {
         "id": "item:8xj31jfpdkf9gvmxdxpi",
@@ -84,58 +85,74 @@ class MockedSurrealModule(ModuleType):
 )
 def test_read_async(tmp_sqlite_db: Path) -> None:
     # confirm that we can load frame data from the core sqlalchemy async
-    # primitives: AsyncConnection, AsyncEngine, and async_sessionmaker
+    # primitives: AsyncEngine, AsyncConnection, async_sessionmaker, and AsyncSession
     from sqlalchemy.ext.asyncio import async_sessionmaker
 
-    async_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_sqlite_db}")
-    async_connection = async_engine.connect()
-    async_session = async_sessionmaker(async_engine)
-    async_session_inst = async_session()
-
-    expected_frame = pl.DataFrame(
-        {"id": [2, 1], "name": ["other", "misc"], "value": [-99.5, 100.0]}
-    )
-    async_conn: Any
-    for async_conn in (
-        async_engine,
-        async_connection,
-        async_session,
-        async_session_inst,
-    ):
-        if async_conn in (async_session, async_session_inst):
-            constraint, execute_opts = "", {}
-        else:
-            constraint = "WHERE value > :n"
-            execute_opts = {"parameters": {"n": -1000}}
-
-        df = pl.read_database(
-            query=f"""
-                SELECT id, name, value
-                FROM test_data {constraint}
-                ORDER BY id DESC
-            """,
-            connection=async_conn,
-            execute_options=execute_opts,
-        )
-        assert_frame_equal(expected_frame, df)
+    async def _test_impl() -> None:
+        async_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_sqlite_db}")
+        async_connection = await async_engine.connect()
+        try:
+            async_session = async_sessionmaker(async_engine)
+            async_session_inst = async_session()
 
+            expected_frame = pl.DataFrame(
+                {"id": [2, 1], "name": ["other", "misc"], "value": [-99.5, 100.0]}
+            )
+            async_conn: Any
+            for async_conn in (
+                async_engine,
+                async_connection,
+                async_session,
+                async_session_inst,
+            ):
+                if async_conn in (async_session, async_session_inst):
+                    constraint, execute_opts = "", {}
+                else:
+                    constraint = "WHERE value > :n"
+                    execute_opts = {"parameters": {"n": -1000}}
+
+                df = pl.read_database(
+                    query=f"""
+                        SELECT id, name, value
+                        FROM test_data {constraint}
+                        ORDER BY id DESC
+                    """,
+                    connection=async_conn,
+                    execute_options=execute_opts,
+                )
+                assert_frame_equal(expected_frame, df)
+        finally:
+            await async_session_inst.close()
+            await async_connection.close()
+            await async_engine.dispose()
 
-async def _nested_async_test(tmp_sqlite_db: Path) -> pl.DataFrame:
-    async_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_sqlite_db}")
-    return pl.read_database(
-        query="SELECT id, name FROM test_data ORDER BY id",
-        connection=async_engine.connect(),
-    )
+    asyncio.run(_test_impl())
 
 
 @pytest.mark.skipif(
     parse_version(sqlalchemy.__version__) < (2, 0),
     reason="SQLAlchemy 2.0+ required for async tests",
 )
-def test_read_async_nested(tmp_sqlite_db: Path) -> None:
-    # This tests validates that we can handle nested async calls
+@pytest.mark.parametrize("started", [True, False])
+def test_read_async_nested(tmp_sqlite_db: Path, started: bool) -> None:
+    # validate that we can handle nested async calls; check
+    # this works with connections that are started/unstarted
+    async def _test_impl() -> pl.DataFrame:
+        async_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_sqlite_db}")
+        async_connection = async_engine.connect()
+        if started:
+            async_connection = await async_connection
+        try:
+            return pl.read_database(
+                query="SELECT id, name FROM test_data ORDER BY id",
+                connection=async_connection,
+            )
+        finally:
+            await async_connection.close()
+            await async_engine.dispose()
+
     expected_frame = pl.DataFrame({"id": [1, 2], "name": ["misc", "other"]})
-    df = asyncio.run(_nested_async_test(tmp_sqlite_db))
+    df = asyncio.run(_test_impl())
     assert_frame_equal(expected_frame, df)
 
 
@@ -187,7 +204,7 @@ def test_surrealdb_fetchall(batch_size: int | None) -> None:
 
 
 def test_async_nested_captured_loop_21263() -> None:
-    # Tests awaiting a future that has "captured" the original event loop from
+    # tests awaiting a future that has "captured" the original event loop from
     # within a `_run_async` context.
     async def test_impl() -> None:
         loop = asyncio.get_running_loop()