diff --git a/src/basic_memory/alembic/env.py b/src/basic_memory/alembic/env.py index 4b9bb4b2e..8649da590 100644 --- a/src/basic_memory/alembic/env.py +++ b/src/basic_memory/alembic/env.py @@ -13,7 +13,7 @@ from alembic import context from basic_memory.config import ConfigManager -from basic_memory.migration_loop import is_running_loop_error, running_on_uvloop +from basic_memory.migration_loop import running_on_uvloop # Allow nested event loops (needed for pytest-asyncio and other async contexts). # nest_asyncio cannot patch a uvloop loop or Python 3.14+; in those cases we skip @@ -167,17 +167,30 @@ def run_in_thread(): future.result() # Wait for completion and re-raise any exceptions -def _run_async_engine_migrations(connectable) -> None: - """Run async-engine migrations with a running-loop fallback.""" +def _loop_is_running() -> bool: + """Check if an event loop is currently running.""" try: + asyncio.get_running_loop() + except RuntimeError: + # Trigger: no running loop (get_running_loop raises RuntimeError) + # Outcome: return False to indicate no loop is running + return False + return True + + +def _run_async_engine_migrations(connectable) -> None: + """Run async migrations, adapting to whether an event loop is already running.""" + + # Trigger: caller may be inside a running event loop (e.g. pytest-asyncio, uvicorn). + # Why: asyncio.run() raises RuntimeError when nested inside a running loop, so we + # detect the condition up front via _loop_is_running() instead of catching the error. + # Outcome: running-loop context -> thread fallback; no loop -> asyncio.run() directly. + if _loop_is_running(): + # Can't nest asyncio.run() inside a running loop -> offload to a thread. + _run_async_migrations_in_thread(connectable) + else: + # No running loop: asyncio.run() is safe; let any unexpected errors propagate. _run_async_migrations_with_asyncio_run(connectable) - except RuntimeError as e: - if is_running_loop_error(e): - # We're in a running event loop (likely uvloop or Python 3.14+ tests). - # Switch to a dedicated thread so Alembic can finish without nesting loops. - _run_async_migrations_in_thread(connectable) - else: - raise def run_migrations_online() -> None: diff --git a/src/basic_memory/migration_loop.py b/src/basic_memory/migration_loop.py index 38e01a83e..d9179b1d9 100644 --- a/src/basic_memory/migration_loop.py +++ b/src/basic_memory/migration_loop.py @@ -28,18 +28,3 @@ def running_on_uvloop() -> bool: except ImportError: return False return isinstance(asyncio.get_event_loop_policy(), uvloop.EventLoopPolicy) - - -def is_running_loop_error(exc: BaseException) -> bool: - """Whether an error means ``asyncio.run`` hit an already-running event loop. - - Two spellings reach the migration fallback: the stdlib "cannot be called - from a running event loop", and "this event loop is already running" when - nest_asyncio mis-patched a uvloop loop. Both mean the same thing — retry the - migration in a dedicated thread rather than re-raising. - """ - msg = str(exc) - return ( - "cannot be called from a running event loop" in msg - or "this event loop is already running" in msg - ) diff --git a/tests/test_alembic_env.py b/tests/test_alembic_env.py index 2dd3a6ae1..64e20ec65 100644 --- a/tests/test_alembic_env.py +++ b/tests/test_alembic_env.py @@ -79,34 +79,63 @@ def raising_asyncio_run(coro): assert fake_coro.closed is True -def test_running_loop_error_uses_thread_fallback(monkeypatch, tmp_path): - """Async-engine helper should switch to the thread fallback for running-loop errors.""" +def test_running_loop_uses_thread_path(monkeypatch, tmp_path): + """When a loop is already running, _run_async_engine_migrations must use the thread path.""" + env_module = load_alembic_env_module(monkeypatch, tmp_path) connectable = object() - fallback_calls: list[object] = [] + thread_calls: list[object] = [] + asyncio_run_calls: list[object] = [] + + monkeypatch.setattr(env_module, "_loop_is_running", lambda: True) + monkeypatch.setattr( + env_module, "_run_async_migrations_in_thread", lambda c: thread_calls.append(c) + ) + monkeypatch.setattr( + env_module, + "_run_async_migrations_with_asyncio_run", + lambda c: asyncio_run_calls.append(c), + ) - def raising_run(connectable): - raise RuntimeError("asyncio.run() cannot be called from a running event loop") + env_module._run_async_engine_migrations(connectable) - def record_fallback(target): - fallback_calls.append(target) + assert thread_calls == [connectable], "thread fallback should have been called" + assert asyncio_run_calls == [], "asyncio.run path must not be called when loop is running" - monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run) - monkeypatch.setattr(env_module, "_run_async_migrations_in_thread", record_fallback) + +def test_no_running_loop_uses_asyncio_run_path(monkeypatch, tmp_path): + """When no loop is running, _run_async_engine_migrations must use asyncio.run directly.""" + env_module = load_alembic_env_module(monkeypatch, tmp_path) + connectable = object() + asyncio_run_calls: list[object] = [] + thread_calls: list[object] = [] + + monkeypatch.setattr(env_module, "_loop_is_running", lambda: False) + monkeypatch.setattr( + env_module, + "_run_async_migrations_with_asyncio_run", + lambda c: asyncio_run_calls.append(c), + ) + monkeypatch.setattr( + env_module, "_run_async_migrations_in_thread", lambda c: thread_calls.append(c) + ) env_module._run_async_engine_migrations(connectable) - assert fallback_calls == [connectable] + assert asyncio_run_calls == [connectable], "asyncio.run path should have been called" + assert thread_calls == [], "thread fallback must not be called when no loop is running" -def test_non_loop_runtime_error_is_re_raised(monkeypatch, tmp_path): - """Unexpected RuntimeError values should not be swallowed by the fallback path.""" +def test_runtime_error_from_asyncio_run_propagates(monkeypatch, tmp_path): + """RuntimeErrors from _run_async_migrations_with_asyncio_run should propagate unchanged.""" env_module = load_alembic_env_module(monkeypatch, tmp_path) + monkeypatch.setattr(env_module, "_loop_is_running", lambda: False) + def raising_run(connectable): - raise RuntimeError("different runtime failure") + raise RuntimeError("unexpected failure") monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run) - with pytest.raises(RuntimeError, match="different runtime failure"): + with pytest.raises(RuntimeError, match="unexpected failure"): env_module._run_async_engine_migrations(object()) diff --git a/tests/test_migration_loop.py b/tests/test_migration_loop.py index 668def3c4..98d4c77d2 100644 --- a/tests/test_migration_loop.py +++ b/tests/test_migration_loop.py @@ -14,22 +14,6 @@ from basic_memory import migration_loop -def test_is_running_loop_error_matches_stdlib_message(): - err = RuntimeError("asyncio.run() cannot be called from a running event loop") - assert migration_loop.is_running_loop_error(err) is True - - -def test_is_running_loop_error_matches_nest_asyncio_uvloop_message(): - # nest_asyncio mis-patched a uvloop loop -> different wording, same meaning. - err = RuntimeError("this event loop is already running") - assert migration_loop.is_running_loop_error(err) is True - - -def test_is_running_loop_error_rejects_unrelated_runtime_error(): - err = RuntimeError("migration failed: column already exists") - assert migration_loop.is_running_loop_error(err) is False - - @pytest.mark.skipif(sys.platform == "win32", reason="uvloop is not available on Windows") def test_running_on_uvloop_true_when_policy_is_uvloop(monkeypatch): import uvloop