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