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 pathlib import Path
from typing import Any

from alembic import command
from alembic.config import Config
from sqlalchemy import Engine, event, inspect, make_url
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, configured for whichever dialect it names."""
    if make_url(url).get_backend_name() != "sqlite":
        # A connection idle across a server restart is dead but still pooled;
        # the pre-ping spends a round trip to find out instead of failing the
        # request.
        return create_engine(url, pool_pre_ping=True)

    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."""
    if engine.dialect.name != "sqlite":
        return
    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()


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)