fix(db): latch init_db per database and time-bound the alembic runner (#342)

Bootstrap and the API lifespan both ran init_db in one process seconds
apart; the second call re-entered the alembic-in-thread machinery
(nested asyncio.run + NullPool engine + greenlet bridge in a reused
worker thread) for zero benefit and hung two consecutive NAS boots
there, blocking the API bind forever with zero SQL activity. init_db
now latches per database URL (drop_db resets it; a different DB always
runs fully), and the alembic worker is bounded at 300s -- a wedged
thread fails startup loudly with a pinpointed error so the container
restarts into a clean retry instead of hanging silently.

Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
2026-07-08 23:39:39 +02:00
committed by GitHub
co-authored by Renn F
parent 5e7c498d00
commit 5886336259
2 changed files with 135 additions and 1 deletions
+43 -1
View File
@@ -208,6 +208,15 @@ async def _db_has_alembic_version(conn: AsyncConnection) -> bool:
return bool(result.scalar()) return bool(result.scalar())
# Hard ceiling on the alembic worker thread. Its env.py nests asyncio.run +
# a fresh NullPool engine + a greenlet bridge inside a (possibly reused)
# executor thread — a hang there previously blocked the API bind forever
# (2026-07-08 NAS outage: two consecutive boots stuck after the alembic
# context lines with zero SQL activity). A timeout can't kill the thread,
# but failing loud lets the container restart into a clean retry.
_ALEMBIC_TIMEOUT_SECONDS = 300
async def run_migrations() -> None: async def run_migrations() -> None:
""" """
Apply Alembic migrations up to head. Apply Alembic migrations up to head.
@@ -242,9 +251,31 @@ async def run_migrations() -> None:
revision=initial_revision, revision=initial_revision,
) )
command.stamp(cfg, initial_revision) command.stamp(cfg, initial_revision)
logger.info("Alembic upgrade starting")
command.upgrade(cfg, "head") command.upgrade(cfg, "head")
logger.info("Alembic upgrade finished")
await asyncio.to_thread(_run_alembic) try:
await asyncio.wait_for(
asyncio.to_thread(_run_alembic), timeout=_ALEMBIC_TIMEOUT_SECONDS
)
except TimeoutError as e:
raise RuntimeError(
f"alembic migration runner exceeded {_ALEMBIC_TIMEOUT_SECONDS}s — "
"worker thread wedged (nested asyncio.run in alembic/env.py); "
"failing startup loudly instead of hanging the API bind"
) from e
class _InitState:
"""Per-process init_db latch, keyed by database URL (see init_db docstring).
URL-keyed so a process that initializes a DIFFERENT database (test
fixtures build throwaway DBs) always runs in full; only a repeat call
for the same database no-ops.
"""
completed_url: str | None = None
async def init_db() -> None: async def init_db() -> None:
@@ -264,7 +295,16 @@ async def init_db() -> None:
to gap-fill any ORM table a migration didn't create. to gap-fill any ORM table a migration didn't create.
`create_all` cannot ALTER an existing table, so an ORM column `create_all` cannot ALTER an existing table, so an ORM column
added without a migration needs a fresh rebuild to appear. added without a migration needs a fresh rebuild to appear.
Idempotent per process: bootstrap and the API lifespan both call this in
the same interpreter seconds apart; the second call re-entered the fragile
alembic-in-thread machinery for zero benefit and hung the 2026-07-08 NAS
boot twice. A completed run latches, later calls no-op. drop_db resets the
latch so tests rebuilding the schema keep working.
""" """
if _InitState.completed_url == settings.database_url:
logger.info("init_db already completed in this process — skipping")
return
engine = get_engine() engine = get_engine()
async with engine.begin() as conn: async with engine.begin() as conn:
# pgvector must exist before tables that use the vector type # pgvector must exist before tables that use the vector type
@@ -299,6 +339,7 @@ async def init_db() -> None:
# subsequent request to introspect the current (post-migration) schema. # subsequent request to introspect the current (post-migration) schema.
await engine.dispose() await engine.dispose()
logger.info("DB engine pool disposed to refresh asyncpg type cache") logger.info("DB engine pool disposed to refresh asyncpg type cache")
_InitState.completed_url = settings.database_url
async def drop_db() -> None: async def drop_db() -> None:
@@ -310,6 +351,7 @@ async def drop_db() -> None:
engine = get_engine() engine = get_engine()
async with engine.begin() as conn: async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all) await conn.run_sync(Base.metadata.drop_all)
_InitState.completed_url = None
async def close_db() -> None: async def close_db() -> None:
+92
View File
@@ -8,6 +8,7 @@ and drop/close.
from __future__ import annotations from __future__ import annotations
import time
from typing import TYPE_CHECKING, cast from typing import TYPE_CHECKING, cast
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
@@ -17,6 +18,7 @@ from roboco.db.base import (
_db_has_alembic_version, _db_has_alembic_version,
_db_has_tables, _db_has_tables,
_DbHolder, _DbHolder,
_InitState,
close_db, close_db,
drop_db, drop_db,
get_db, get_db,
@@ -39,9 +41,11 @@ def _reset_holder() -> Generator[None]:
"""Snapshot/restore the singleton so tests don't poison the live engine.""" """Snapshot/restore the singleton so tests don't poison the live engine."""
saved_engine = _DbHolder.engine saved_engine = _DbHolder.engine
saved_factory = _DbHolder.session_factory saved_factory = _DbHolder.session_factory
_InitState.completed_url = None
yield yield
_DbHolder.engine = saved_engine _DbHolder.engine = saved_engine
_DbHolder.session_factory = saved_factory _DbHolder.session_factory = saved_factory
_InitState.completed_url = None
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -454,6 +458,94 @@ async def test_init_db_fresh_db_runs_migrations() -> None:
fake_engine.dispose.assert_awaited_once() fake_engine.dispose.assert_awaited_once()
def _fake_engine_for_init() -> tuple[MagicMock, MagicMock]:
fake_conn = MagicMock()
fake_conn.execute = AsyncMock()
fake_conn.run_sync = AsyncMock()
class _ConnCm:
async def __aenter__(self) -> object:
return fake_conn
async def __aexit__(self, *_args: object) -> None:
return None
fake_engine = MagicMock()
fake_engine.begin = MagicMock(return_value=_ConnCm())
fake_engine.connect = MagicMock(return_value=_ConnCm())
fake_engine.dispose = AsyncMock()
return fake_engine, fake_conn
@pytest.mark.asyncio
async def test_init_db_second_call_same_db_is_noop() -> None:
"""Bootstrap and the API lifespan both call init_db in one process; the
second call must not re-enter the alembic machinery (2026-07-08 NAS hang)."""
fake_engine, _ = _fake_engine_for_init()
with (
patch("roboco.db.base.get_engine", return_value=fake_engine),
patch("roboco.db.base._db_has_tables", new=AsyncMock(return_value=True)),
patch("roboco.db.base.run_migrations", new=AsyncMock()) as rm,
):
await init_db()
await init_db()
rm.assert_awaited_once()
fake_engine.dispose.assert_awaited_once()
@pytest.mark.asyncio
async def test_init_db_reruns_for_a_different_database_url() -> None:
"""The latch is URL-keyed: a process initializing a different DB runs fully."""
_InitState.completed_url = "postgresql+asyncpg://other-host/other-db"
fake_engine, _ = _fake_engine_for_init()
with (
patch("roboco.db.base.get_engine", return_value=fake_engine),
patch("roboco.db.base._db_has_tables", new=AsyncMock(return_value=True)),
patch("roboco.db.base.run_migrations", new=AsyncMock()) as rm,
):
await init_db()
rm.assert_awaited_once()
@pytest.mark.asyncio
async def test_drop_db_resets_the_init_latch() -> None:
"""drop_db clears the latch so a rebuild in the same process runs fully."""
fake_engine, _ = _fake_engine_for_init()
with (
patch("roboco.db.base.get_engine", return_value=fake_engine),
patch("roboco.db.base._db_has_tables", new=AsyncMock(return_value=True)),
patch("roboco.db.base.run_migrations", new=AsyncMock()) as rm,
):
await init_db()
await drop_db()
await init_db()
expected_full_runs = 2
assert rm.await_count == expected_full_runs
@pytest.mark.asyncio
async def test_run_migrations_times_out_loudly_on_wedged_worker() -> None:
"""A wedged alembic worker thread fails startup with a clear error instead
of hanging the API bind forever (the 2026-07-08 boot-hang shape)."""
fake_engine, _ = _fake_engine_for_init()
fake_command = MagicMock()
fake_command.upgrade = MagicMock(side_effect=lambda *_a, **_k: time.sleep(0.5))
with (
patch("roboco.db.base.get_engine", return_value=fake_engine),
patch(
"roboco.db.base._db_has_alembic_version",
new=AsyncMock(return_value=True),
),
patch("roboco.db.base.command", fake_command),
patch("roboco.db.base._ALEMBIC_TIMEOUT_SECONDS", 0.05),
pytest.raises(RuntimeError, match="alembic migration runner exceeded"),
):
await run_migrations()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# drop_db / close_db # drop_db / close_db
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------