db.py

"""
Turning velixir's DATABASE_URL into a SQLAlchemy engine, and waiting for the database.

velixir injects DATABASE_URL as postgres://user:pass@host:5432/db?sslmode=require. Two
things stand between that string and a working engine:

  1. The scheme. SQLAlchemy has not accepted "postgres://" since 1.4, and a bare
     "postgresql://" means whichever driver its version defaults to. "postgresql+psycopg://"
     names psycopg 3, the driver in requirements.txt.
  2. TLS, which needs nothing. sslmode stays in the URL, psycopg passes it to libpq, and
     libpq's "require" is what velixir means: encrypt the connection, but do not check the
     certificate against public CAs (it comes from velixir's own CA).
"""

import os
import threading
import time

from sqlalchemy import create_engine, text
from sqlalchemy.orm import DeclarativeBase, sessionmaker

RETRY_SECONDS = 5

# Any number works, as long as nothing else in the database takes the same advisory lock.
SCHEMA_LOCK = 7_214_202_612


def sqlalchemy_url(database_url):
    scheme, separator, rest = database_url.partition("://")
    if separator and scheme in ("postgres", "postgresql"):
        return f"postgresql+psycopg://{rest}"
    return database_url


DATABASE_URL = os.environ.get("DATABASE_URL", "").strip()

# No DATABASE_URL means setup mode: no engine, and main.py serves the setup page instead.
engine = (
    create_engine(
        sqlalchemy_url(DATABASE_URL),
        # Keep the pool small. Replicas multiply it, and a managed instance caps connections.
        pool_size=5,
        max_overflow=5,
        # Test a pooled connection before handing it out, so a database restart costs one
        # quick check rather than one failed request.
        pool_pre_ping=True,
        # Give up after five seconds when the database is unreachable, instead of hanging a
        # request (and the health check) while the kernel retries the connection.
        connect_args={"connect_timeout": 5},
    )
    if DATABASE_URL
    else None
)

Session = sessionmaker(engine, expire_on_commit=False)


class Base(DeclarativeBase):
    pass


class State:
    ready = False
    last_error = None


state = State()


def create_tables():
    # Several replicas (or uvicorn workers) starting at once would race CREATE TABLE. The
    # advisory lock makes them take turns, and ends with the transaction.
    with engine.begin() as conn:
        conn.execute(text("SELECT pg_advisory_xact_lock(:key)"), {"key": SCHEMA_LOCK})
        Base.metadata.create_all(conn)


def connect_forever():
    """Create the tables, retrying until the database answers. Never gives up, never exits.

    A database that is briefly unreachable (a failover, a rotated password waiting for a
    redeploy, a database still starting) gets an explanation on the page instead of a crash
    loop, and the app recovers by itself when the database comes back.
    """
    attempt = 0
    while True:
        attempt += 1
        try:
            create_tables()
        except Exception as exc:
            # The driver's own message, without SQLAlchemy's wrapping around it.
            state.last_error = str(getattr(exc, "orig", None) or exc).strip()
            if attempt == 1 or attempt % 12 == 0:
                print(f"database not reachable yet (attempt {attempt}): {state.last_error}", flush=True)
            time.sleep(RETRY_SECONDS)
            continue

        state.ready = True
        state.last_error = None
        print("database ready", flush=True)
        return


def connect_in_background():
    threading.Thread(target=connect_forever, name="connect-db", daemon=True).start()