FabledSteward/steward/database.py

from __future__ import annotations
import asyncio
import logging
import time
from typing import TYPE_CHECKING
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, 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

# 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


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()