import sqlite3 from collections.abc import Iterator from contextlib import contextmanager from pathlib import Path from fastapi import Request SCHEMA_VERSION = 2 SCHEMA = """ CREATE TABLE IF NOT EXISTS tenants ( id INTEGER PRIMARY KEY, name TEXT NOT NULL UNIQUE COLLATE NOCASE ); -- A job is identified by (tenant, source system, the source's own job id), so two -- products can't collide and renaming a job in Veeam doesn't create a duplicate. CREATE TABLE IF NOT EXISTS jobs ( id INTEGER PRIMARY KEY, tenant_id INTEGER NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, source_system TEXT NOT NULL, external_id TEXT NOT NULL, name TEXT NOT NULL, next_run_at TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE (tenant_id, source_system, external_id) ); CREATE TABLE IF NOT EXISTS runs ( id INTEGER PRIMARY KEY, job_id INTEGER NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, external_id TEXT NOT NULL, status TEXT NOT NULL CHECK (status IN ('running', 'success', 'warning', 'failed')), started_at TEXT NOT NULL, finished_at TEXT, message TEXT NOT NULL DEFAULT '', size_bytes INTEGER, details TEXT NOT NULL DEFAULT '{}', received_at TEXT NOT NULL, CHECK ((status = 'running') = (finished_at IS NULL)), UNIQUE (job_id, external_id) ); CREATE INDEX IF NOT EXISTS runs_job_started ON runs (job_id, started_at); CREATE INDEX IF NOT EXISTS runs_job_finished ON runs (job_id, finished_at); -- Keys used by collectors to push results. Only the SHA-256 is stored. CREATE TABLE IF NOT EXISTS api_keys ( id INTEGER PRIMARY KEY, name TEXT NOT NULL UNIQUE, key_hash TEXT NOT NULL UNIQUE, created_at TEXT NOT NULL, last_used_at TEXT, revoked_at TEXT ); """ def connect(path: Path | str) -> sqlite3.Connection: # Autocommit mode; writes that must be atomic go through transaction(). # check_same_thread=False: FastAPI may run a dependency and its endpoint on different # threadpool threads, but a connection is only ever used by one request at a time. conn = sqlite3.connect(path, timeout=10, isolation_level=None, check_same_thread=False) conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys = ON") conn.execute("PRAGMA busy_timeout = 10000") return conn # Changes since SCHEMA (which is v1), applied in order to bring any database up to SCHEMA_VERSION MIGRATIONS = { 2: [ # Who created/revoked a key from the admin page, and which company it last reported for "ALTER TABLE api_keys ADD COLUMN created_by TEXT", "ALTER TABLE api_keys ADD COLUMN revoked_by TEXT", "ALTER TABLE api_keys ADD COLUMN last_used_tenant TEXT", ], } def init_db(path: Path | str) -> None: conn = connect(path) try: conn.execute("PRAGMA journal_mode = WAL") version = conn.execute("PRAGMA user_version").fetchone()[0] if version > SCHEMA_VERSION: raise RuntimeError(f"Database schema v{version} is newer than this app (v{SCHEMA_VERSION})") conn.executescript(SCHEMA) # v1 tables, for a new database with transaction(conn): for target in range(max(version, 1) + 1, SCHEMA_VERSION + 1): for statement in MIGRATIONS[target]: conn.execute(statement) conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") finally: conn.close() @contextmanager def transaction(conn: sqlite3.Connection) -> Iterator[None]: conn.execute("BEGIN IMMEDIATE") try: yield except BaseException: conn.execute("ROLLBACK") raise conn.execute("COMMIT") def get_db(request: Request) -> Iterator[sqlite3.Connection]: conn = connect(request.app.state.settings.db_path) try: yield conn finally: conn.close()