Backup-Dashboard/backend/app/db.py

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