app/backend/fluksio/core/db.py

"""The database: one file beside the flows, or whatever ``DATABASE_URL`` says.

SQLite is the default because one engine process owns this database. Node
execution, live values and the work queue never come here — they are the state
backend's and the bus's — so what lands in a transaction is the rollups a
minute at a time, a row per cascade, and the run history. WAL lets the readers
carry on while that single writer works.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from typing import Any

from alembic import command
from alembic.config import Config
from sqlalchemy import Engine, event, inspect, make_url, text
from sqlmodel import Session, SQLModel, create_engine, select

from fluksio import crud
from fluksio.core.config import settings
from fluksio.models import Flavor, User, UserCreate

logger = logging.getLogger(__name__)


def make_engine(url: str) -> Engine:
    """One engine for the process. SQLite, and only SQLite.

    Not a limitation that snuck in: `metric_minute` and every run table are
    written with `sqlalchemy.dialects.sqlite.insert(...).on_conflict_do_update`
    and with `max(a, b)`, neither of which another dialect has. Pointing this
    at Postgres used to migrate cleanly, serve, log in — and then lose every
    observability flush into the collector's hold buffer and fail every run,
    which is a far worse answer than refusing at startup.
    """
    if make_url(url).get_backend_name() != "sqlite":
        raise RuntimeError(
            f"DATABASE_URL must be a SQLite URL, not {make_url(url).get_backend_name()!r}. "
            "The engine's own database is a file beside the flows; the "
            "analytics Postgres in the compose stack is Umami's, not this."
        )

    database = make_url(url).database
    if database and database != ":memory:":
        Path(database).parent.mkdir(parents=True, exist_ok=True)
    # The engine's threads — the run service, the metrics collector, the
    # request pool — share this engine, so connections cross threads. Each one
    # is still used by a single thread at a time; the pool sees to that.
    return create_engine(url, connect_args={"check_same_thread": False})


engine = make_engine(settings.SQLALCHEMY_DATABASE_URI)


@event.listens_for(engine, "connect")
def _sqlite_pragmas(dbapi_connection: Any, _record: Any) -> None:
    """What makes concurrent readers, cascades and waiting-out a writer work."""
    cursor = dbapi_connection.cursor()
    # Readers do not block on the writer, which is the whole reason this is
    # usable while a run is streaming metrics into it.
    cursor.execute("PRAGMA journal_mode=WAL")
    # Commit without waiting for the platter. A power cut can lose the last
    # transactions; it cannot corrupt the file.
    cursor.execute("PRAGMA synchronous=NORMAL")
    # Off by default in SQLite, and `ondelete="CASCADE"` on the OAuth tables is
    # the whole of revoking a client.
    cursor.execute("PRAGMA foreign_keys=ON")
    # A second writer waits rather than raising "database is locked".
    cursor.execute("PRAGMA busy_timeout=30000")
    cursor.close()


@contextmanager
def writing(session: Session) -> Iterator[None]:
    """Take SQLite's write lock before reading, for a read-then-write handler.

    SQLAlchemy begins deferred, so a transaction that reads and then writes
    has to upgrade — and if another writer committed in between, SQLite
    refuses with `SQLITE_BUSY_SNAPSHOT` immediately rather than waiting out
    `busy_timeout`. The engine writes constantly (the metrics flush, a run
    recording nodes), so any handler that decides *whether* to write by
    reading first can lose that race and 500.

    Only worth it where the handler really does both: this serialises against
    other writers for the length of the transaction, which is what WAL's
    concurrent readers otherwise avoid.
    """
    session.execute(text("BEGIN IMMEDIATE"))
    try:
        yield
    except Exception:
        session.rollback()
        raise


def _alembic_config(connection: Any) -> Config:
    """Alembic driven from here, so an installed wheel needs no alembic.ini."""
    config = Config()
    config.set_main_option(
        "script_location", str(Path(__file__).parents[1] / "alembic")
    )
    config.attributes["connection"] = connection
    return config


def migrate(engine: Engine) -> None:
    """Bring the schema up to date, creating it if there is nothing there."""
    with engine.begin() as connection:
        config = _alembic_config(connection)
        tables = inspect(connection).get_table_names()
        if not tables:
            # Nothing to upgrade from. Building the schema from the models is
            # both faster and exactly what the revisions would have produced.
            SQLModel.metadata.create_all(connection)
            command.stamp(config, "head")
            logger.info("created the database schema")
            return
        command.upgrade(config, "head")


def init_db(session: Session) -> None:
    """The first superuser, when the deployment names one.

    A `fluksio serve` with no configuration makes its own and prints it once;
    that path calls into `fluksio.core.bootstrap` instead of this.
    """
    if not (settings.FIRST_SUPERUSER and settings.FIRST_SUPERUSER_PASSWORD):
        return
    user = session.exec(
        select(User).where(User.email == settings.FIRST_SUPERUSER)
    ).first()
    if not user:
        user_in = UserCreate(
            email=settings.FIRST_SUPERUSER,
            password=settings.FIRST_SUPERUSER_PASSWORD,
            is_superuser=True,
        )
        user = crud.create_user(session=session, user_create=user_in)


#: The sizes an installation starts with. Written once, when there are none,
#: and editable from there — what a name means is a property of the machines
#: this installation has, and nothing here knows what those are.
SEED_FLAVORS: tuple[dict[str, Any], ...] = (
    {"name": "small", "cpus": 1, "ram": 2048, "description": "A poll, a threshold"},
    {"name": "medium", "cpus": 4, "ram": 8192, "description": "A step that computes"},
    {"name": "large", "cpus": 8, "ram": 16384, "description": "A heavy step"},
    {
        "name": "gpu-small",
        "cpus": 4,
        "gpus": 1,
        "ram": 16384,
        "description": "One card, and cores to feed it",
    },
)


def seed_flavors(session: Session) -> None:
    """Give a new installation sizes to pick from, once.

    Only when there are none at all: they are editable, and re-adding one that
    somebody deliberately removed would be an argument nobody can win.
    """
    if session.exec(select(Flavor)).first() is not None:
        return
    for row in SEED_FLAVORS:
        session.add(Flavor(**row))
    session.commit()
    logger.info("seeded %d resource flavors", len(SEED_FLAVORS))


def prepare(engine: Engine) -> None:
    """Everything that has to be true before the app serves a request."""
    migrate(engine)
    with Session(engine) as session:
        init_db(session)
        seed_flavors(session)