"""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 instance starts with. Written once, when there are none, #: and editable from there — what a name means is a property of the machines #: this instance 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 instance 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)