From e97445e689e8daf1f829179b4e4887506603f673 Mon Sep 17 00:00:00 2001 From: Simon Lin Date: Fri, 7 Mar 2025 19:59:53 +1100 Subject: [PATCH] Revert "feat(python): Drop `nest-asyncio` in favor of custom logic (#20793)" This reverts commit 3696e53200a30f67b82a99da5978b568554908b7. --- py-polars/polars/io/database/_utils.py | 46 +++++-------------- py-polars/polars/meta/versions.py | 2 + py-polars/pyproject.toml | 3 +- py-polars/requirements-dev.txt | 1 + .../tests/unit/io/database/test_async.py | 5 +- 5 files changed, 21 insertions(+), 36 deletions(-) diff --git a/py-polars/polars/io/database/_utils.py b/py-polars/polars/io/database/_utils.py index ac6b0680ba27..8fe62928cdea 100644 --- a/py-polars/polars/io/database/_utils.py +++ b/py-polars/polars/io/database/_utils.py @@ -1,9 +1,6 @@ from __future__ import annotations -import asyncio import re -import threading -from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any from polars._utils.various import parse_version @@ -17,39 +14,20 @@ from polars._typing import SchemaDict -def _run_async( - coroutine: Coroutine[Any, Any, Any], *, timeout: float | None = None -) -> Any: - """Run asynchronous code as if it were synchronous. +def _run_async(co: Coroutine[Any, Any, Any]) -> Any: + """Run asynchronous code as if it was synchronous.""" + import asyncio - This is required for execution in Jupyter notebook environments. - """ - # Implementation taken from StackOverflow answer here: - # https://stackoverflow.com/a/78911765/2344703 + from polars._utils.unstable import issue_unstable_warning + from polars.dependencies import import_optional - try: - loop = asyncio.get_running_loop() - except RuntimeError: - # If there is no running loop, use `asyncio.run` normally - return asyncio.run(coroutine) - - def run_in_new_loop() -> Any: - new_loop = asyncio.new_event_loop() - asyncio.set_event_loop(new_loop) - try: - return new_loop.run_until_complete(coroutine) - finally: - new_loop.close() - - if threading.current_thread() is threading.main_thread(): - if not loop.is_running(): - return loop.run_until_complete(coroutine) - else: - with ThreadPoolExecutor() as pool: - future = pool.submit(run_in_new_loop) - return future.result(timeout=timeout) - else: - return asyncio.run_coroutine_threadsafe(coroutine, loop).result() + issue_unstable_warning( + "Use of asynchronous connections is currently considered unstable " + "and unexpected issues may arise; if this happens, please report them." + ) + nest_asyncio = import_optional("nest_asyncio") + nest_asyncio.apply() + return asyncio.run(co) def _read_sql_connectorx( diff --git a/py-polars/polars/meta/versions.py b/py-polars/polars/meta/versions.py index 32f3486012c1..521444d9c15c 100644 --- a/py-polars/polars/meta/versions.py +++ b/py-polars/polars/meta/versions.py @@ -30,6 +30,7 @@ def show_versions() -> None: fsspec: 2023.12.2 gevent: 24.2.1 matplotlib: 3.8.4 + nest_asyncio: 1.6.0 numpy: 1.26.4 openpyxl: 3.1.2 pandas: 2.2.2 @@ -84,6 +85,7 @@ def _get_dependency_list() -> list[str]: "google.auth", "great_tables", "matplotlib", + "nest_asyncio", "numpy", "openpyxl", "pandas", diff --git a/py-polars/pyproject.toml b/py-polars/pyproject.toml index 0c8f54b1d189..7be012b40d45 100644 --- a/py-polars/pyproject.toml +++ b/py-polars/pyproject.toml @@ -59,7 +59,7 @@ excel = ["polars[calamine,openpyxl,xlsx2csv,xlsxwriter]"] adbc = ["adbc-driver-manager[dbapi]", "adbc-driver-sqlite[dbapi]"] connectorx = ["connectorx >= 0.3.2"] sqlalchemy = ["sqlalchemy", "polars[pandas]"] -database = ["polars[adbc,connectorx,sqlalchemy]"] +database = ["polars[adbc,connectorx,sqlalchemy]", "nest-asyncio"] # Cloud fsspec = ["fsspec"] @@ -117,6 +117,7 @@ module = [ "kuzu", "matplotlib.*", "moto.server", + "nest_asyncio", "openpyxl", "polars.polars", "polars_cloud", diff --git a/py-polars/requirements-dev.txt b/py-polars/requirements-dev.txt index 9df70187a1fc..b961f89e5f3c 100644 --- a/py-polars/requirements-dev.txt +++ b/py-polars/requirements-dev.txt @@ -31,6 +31,7 @@ adbc-driver-sqlite; platform_system != 'Windows' aiosqlite connectorx kuzu +nest-asyncio # Cloud cloudpickle fsspec diff --git a/py-polars/tests/unit/io/database/test_async.py b/py-polars/tests/unit/io/database/test_async.py index 0afb570c428b..69766f2aa94f 100644 --- a/py-polars/tests/unit/io/database/test_async.py +++ b/py-polars/tests/unit/io/database/test_async.py @@ -132,7 +132,10 @@ async def _nested_async_test(tmp_sqlite_db: Path) -> pl.DataFrame: 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 + # this tests validates that we can handle nested async calls. without + # the nested asyncio handling provided by `nest_asyncio` this test + # would raise a RuntimeError + expected_frame = pl.DataFrame({"id": [1, 2], "name": ["misc", "other"]}) df = asyncio.run(_nested_async_test(tmp_sqlite_db)) assert_frame_equal(expected_frame, df)