82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
"""Database readiness helper: wait until the DB is reachable and tables exist.
|
|
|
|
Used by the worker process to avoid crashing when the backend hasn't yet
|
|
finished running `init_db()` (which creates tables). Worker starts in parallel
|
|
with backend in docker-compose and may come up first.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from sqlalchemy import inspect, select, text
|
|
|
|
from app.db import AsyncSessionLocal, Base, engine
|
|
from app.logging_setup import get_logger
|
|
from app.models import Setting
|
|
|
|
log = get_logger("db_wait")
|
|
|
|
|
|
async def wait_for_db(
|
|
max_retries: int = 60,
|
|
delay: float = 2.0,
|
|
required_tables: tuple[str, ...] = ("settings",),
|
|
) -> None:
|
|
"""Block until the database is reachable AND all `required_tables` exist.
|
|
|
|
Retries on connection errors and on missing-table errors. Logs progress.
|
|
Raises the last error after `max_retries` attempts.
|
|
"""
|
|
last_err: Exception | None = None
|
|
for attempt in range(1, max_retries + 1):
|
|
try:
|
|
# Check raw connectivity
|
|
async with engine.connect() as conn:
|
|
await conn.execute(text("SELECT 1"))
|
|
|
|
# Check that required tables exist
|
|
async with engine.connect() as conn:
|
|
existing = await conn.run_sync(
|
|
lambda sync_conn: set(inspect(sync_conn).get_table_names())
|
|
)
|
|
missing = [t for t in required_tables if t not in existing]
|
|
if missing:
|
|
raise RuntimeError(f"required tables not yet created: {missing}")
|
|
|
|
# Smoke-test the `settings` table specifically (the worker's first query)
|
|
async with AsyncSessionLocal() as db:
|
|
await db.execute(select(Setting).limit(1))
|
|
|
|
log.info("db_ready", attempt=attempt)
|
|
return
|
|
except Exception as e:
|
|
last_err = e
|
|
log.warning(
|
|
"db_not_ready_retry",
|
|
attempt=attempt,
|
|
max_retries=max_retries,
|
|
error=f"{type(e).__name__}: {e}",
|
|
)
|
|
await asyncio.sleep(delay)
|
|
|
|
# Exhausted retries — surface the last error so the caller can decide.
|
|
assert last_err is not None
|
|
raise last_err
|
|
|
|
|
|
async def wait_for_db_or_exit(
|
|
max_retries: int = 60,
|
|
delay: float = 2.0,
|
|
required_tables: tuple[str, ...] = ("settings",),
|
|
) -> None:
|
|
"""Like `wait_for_db`, but exits the process with code 1 on failure.
|
|
|
|
Useful as the very first call in the worker entrypoint so it doesn't
|
|
spam logs forever if the DB is genuinely unreachable.
|
|
"""
|
|
try:
|
|
await wait_for_db(max_retries=max_retries, delay=delay, required_tables=required_tables)
|
|
except Exception as e:
|
|
log.error("db_wait_exhausted", error=f"{type(e).__name__}: {e}")
|
|
raise SystemExit(1)
|