mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
+43
-1
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user