from __future__ import annotations import asyncio import logging import time from typing import TYPE_CHECKING from sqlalchemy import text from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import ( create_async_engine, async_sessionmaker, AsyncEngine, AsyncSession, ) from sqlalchemy.pool import NullPool if TYPE_CHECKING: from quart import Quart logger = logging.getLogger(__name__) # Recycle pooled connections after 30 minutes. Nothing in Postgres closes an # idle connection by default, but the path to it is not always durable: NAT and # conntrack tables drop idle flows, and a Docker network rebuild silently # strands existing sockets. Capping connection age means a stranded connection # is retired on a timer instead of surfacing as a failed query later. POOL_RECYCLE_SECONDS = 1800 # Total budget for the database to become reachable at startup, and the backoff # ceiling between attempts. 60s comfortably covers Postgres WAL recovery after # an unclean shutdown (observed ~7s) plus Docker DNS coming up on a host reboot. DB_CONNECT_TIMEOUT_SECONDS = 60.0 _BACKOFF_INITIAL_SECONDS = 0.5 _BACKOFF_MAX_SECONDS = 5.0 # Per-REQUEST retry schedule (seconds between attempts), total ~1.75s across 4 # tries. Far shorter than the startup budget on purpose: nobody is watching a # container boot, but somebody is watching this page load. Long enough to ride # out a pool blip or a fast reconnect, short enough that a genuinely-down # database gets an honest answer instead of a spinner. _REQUEST_RETRY_DELAYS = (0.25, 0.5, 1.0) # Postgres SQLSTATEs that will never resolve by waiting — retrying these just # delays a clear error behind a full timeout budget. _FATAL_SQLSTATES = { "28P01", # invalid_password "28000", # invalid_authorization_specification "3D000", # invalid_catalog_name — database does not exist } def init_db(app: "Quart") -> None: """Create async engine and attach db_sessionmaker to app. Called during create_app() after config is loaded. Does not create tables — Alembic handles migrations. """ db_url: str = app.config["DATABASE_URL"] engine = create_async_engine( db_url, echo=False, # Check a pooled connection is still alive before handing it out. When # the database restarts, every connection already in the pool is dead # at the socket level; without this, SQLAlchemy only discovers that by # failing a real query, so the first operation after a DB restart # errors out on whatever triggered it (a request, a scheduled poll). # The check is a cheap round-trip and it makes a DB restart invisible. pool_pre_ping=True, pool_recycle=POOL_RECYCLE_SECONDS, ) app.db_sessionmaker: async_sessionmaker[AsyncSession] = async_sessionmaker( engine, expire_on_commit=False ) app._db_engine = engine def _fatal_sqlstate(exc: BaseException) -> str | None: """Return the SQLSTATE if this error chain carries a non-retryable one.""" seen: set[int] = set() cur: BaseException | None = exc while cur is not None and id(cur) not in seen: seen.add(id(cur)) sqlstate = getattr(cur, "sqlstate", None) if sqlstate in _FATAL_SQLSTATES: return sqlstate cur = cur.__cause__ or cur.__context__ return None def _describe(exc: BaseException) -> str: """Innermost cause, which is the part that says what actually went wrong. SQLAlchemy wraps DBAPI errors several layers deep; the outer message is boilerplate, so surface the root for the waiting-for-database log line. """ cur: BaseException = exc seen: set[int] = {id(cur)} while True: nxt = cur.__cause__ or cur.__context__ if nxt is None or id(nxt) in seen: return f"{type(cur).__name__}: {cur}" seen.add(id(nxt)) cur = nxt def wait_for_database( db_url: str, timeout_seconds: float = DB_CONNECT_TIMEOUT_SECONDS, ) -> None: """Block until the database accepts a connection, or raise after the budget. Steward otherwise assumes the database is both resolvable and *ready* the first time it asks, which is false in two ordinary situations: • Host reboot / full stack restart — the app container can try to resolve the `db` service name before Docker's embedded DNS has the record, giving `gaierror -2 Name or service not known`. • Unclean shutdown — Postgres is listening but still replaying WAL, and refuses connections with "the database system is not yet accepting connections" until recovery reaches a consistent state. compose's `depends_on: service_healthy` covers ordering on a clean `up`, but not either of the above. Both are transient and self-healing, so retry rather than dumping a traceback and dying. Credential and missing-database errors are NOT transient and fail immediately. """ deadline = time.monotonic() + timeout_seconds delay = _BACKOFF_INITIAL_SECONDS attempt = 0 last_exc: BaseException | None = None while True: attempt += 1 try: asyncio.run(_probe(db_url)) except Exception as exc: # broad by design — classified just below sqlstate = _fatal_sqlstate(exc) if sqlstate is not None: # Wrong password / missing database: waiting cannot fix it. raise last_exc = exc remaining = deadline - time.monotonic() if remaining <= 0: break logger.warning( "Database not ready (attempt %d, %.0fs budget left): %s — " "retrying in %.1fs", attempt, remaining, _describe(exc), min(delay, remaining), ) time.sleep(min(delay, remaining)) delay = min(delay * 2, _BACKOFF_MAX_SECONDS) else: if attempt > 1: logger.info("Database ready after %d attempt(s)", attempt) return detail = f" Last error: {_describe(last_exc)}" if last_exc is not None else "" raise RuntimeError( f"Database did not become available within {timeout_seconds:.0f}s " f"({attempt} attempts).{detail}" ) from last_exc class DatabaseUnavailable(Exception): """The database could not be reached while serving a request.""" async def ensure_database_reachable(engine: AsyncEngine) -> None: """Check out one pooled connection, retrying briefly, or raise. Startup readiness (wait_for_database) does not help once the app is already serving: if the database goes away underneath a running Steward, the next request needs a connection, the pooled ones are dead, and establishing a new one fails — which is how a login attempt turned into an opaque 500. pool_pre_ping already makes recovery automatic *once the database is back*. What it cannot do is wait: while the container is genuinely down, its replacement connect fails too. So retry briefly here to ride out a restart, then give up and let the caller render an honest "database unavailable" page rather than a generic error. The budget is deliberately short — a person is waiting on this request, and a page that hangs for half a minute is worse than one that says plainly what is wrong and offers a retry. """ last_exc: BaseException | None = None for attempt, delay in enumerate(_REQUEST_RETRY_DELAYS + (None,)): try: async with engine.connect(): if attempt: logger.info( "Database reachable again after %d retry attempt(s)", attempt) return except (SQLAlchemyError, OSError) as exc: # Bad credentials / missing database will never resolve by waiting, # and they are not what this guard is for — let them surface. if _fatal_sqlstate(exc) is not None: raise last_exc = exc if delay is not None: await asyncio.sleep(delay) raise DatabaseUnavailable(_describe(last_exc)) from last_exc async def _probe(db_url: str) -> None: """Open one throwaway connection and round-trip a trivial query. Uses its own engine with NullPool: this runs before the app engine exists, and a probe connection must never be left in a pool for real work to reuse. """ engine = create_async_engine(db_url, echo=False, poolclass=NullPool) try: async with engine.connect() as conn: await conn.execute(text("SELECT 1")) finally: await engine.dispose()