import aiosqlite from pathlib import Path from app.config import settings SCHEMA = """ CREATE TABLE IF NOT EXISTS drives ( id INTEGER PRIMARY KEY AUTOINCREMENT, truenas_disk_id TEXT UNIQUE NOT NULL, devname TEXT NOT NULL, serial TEXT, model TEXT, size_bytes INTEGER, temperature_c INTEGER, smart_health TEXT DEFAULT 'UNKNOWN', last_seen_at TEXT NOT NULL, last_polled_at TEXT NOT NULL, notes TEXT, location TEXT ); CREATE TABLE IF NOT EXISTS smart_tests ( id INTEGER PRIMARY KEY AUTOINCREMENT, drive_id INTEGER NOT NULL REFERENCES drives(id) ON DELETE CASCADE, test_type TEXT NOT NULL CHECK(test_type IN ('short', 'long')), state TEXT NOT NULL DEFAULT 'idle', percent INTEGER DEFAULT 0, truenas_job_id INTEGER, started_at TEXT, eta_at TEXT, finished_at TEXT, error_text TEXT, UNIQUE(drive_id, test_type) ); CREATE TABLE IF NOT EXISTS burnin_jobs ( id INTEGER PRIMARY KEY AUTOINCREMENT, drive_id INTEGER NOT NULL REFERENCES drives(id), profile TEXT NOT NULL, state TEXT NOT NULL DEFAULT 'queued', percent INTEGER DEFAULT 0, stage_name TEXT, operator TEXT NOT NULL, created_at TEXT NOT NULL, started_at TEXT, finished_at TEXT, error_text TEXT ); CREATE TABLE IF NOT EXISTS burnin_stages ( id INTEGER PRIMARY KEY AUTOINCREMENT, burnin_job_id INTEGER NOT NULL REFERENCES burnin_jobs(id) ON DELETE CASCADE, stage_name TEXT NOT NULL, state TEXT NOT NULL DEFAULT 'pending', percent INTEGER DEFAULT 0, started_at TEXT, finished_at TEXT, duration_seconds REAL, error_text TEXT ); CREATE TABLE IF NOT EXISTS audit_events ( id INTEGER PRIMARY KEY AUTOINCREMENT, event_type TEXT NOT NULL, drive_id INTEGER REFERENCES drives(id), burnin_job_id INTEGER REFERENCES burnin_jobs(id), operator TEXT, message TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) ); CREATE INDEX IF NOT EXISTS idx_smart_drive_type ON smart_tests(drive_id, test_type); CREATE INDEX IF NOT EXISTS idx_burnin_jobs_drive ON burnin_jobs(drive_id, state); CREATE INDEX IF NOT EXISTS idx_burnin_stages_job ON burnin_stages(burnin_job_id); CREATE INDEX IF NOT EXISTS idx_audit_events_job ON audit_events(burnin_job_id); """ # Migrations for existing databases that predate schema additions. # Each entry is tried with try/except — SQLite raises OperationalError # ("duplicate column name") if the column already exists, which is safe to ignore. _MIGRATIONS = [ "ALTER TABLE drives ADD COLUMN notes TEXT", "ALTER TABLE drives ADD COLUMN location TEXT", # Stage 7: SSH command output + SMART attribute storage "ALTER TABLE burnin_stages ADD COLUMN log_text TEXT", "ALTER TABLE burnin_stages ADD COLUMN bad_blocks INTEGER DEFAULT 0", "ALTER TABLE drives ADD COLUMN smart_attrs TEXT", "ALTER TABLE smart_tests ADD COLUMN raw_output TEXT", # Stage 8: track last reset time so dashboard burn-in col clears after reset "ALTER TABLE drives ADD COLUMN last_reset_at TEXT", # 1.0.0-15: pool-membership lock "ALTER TABLE drives ADD COLUMN pool_name TEXT", "ALTER TABLE drives ADD COLUMN pool_role TEXT", "ALTER TABLE drives ADD COLUMN pool_seen_at TEXT", # 1.0.0-44: per-pattern badblocks progress for the drive drawer's # 4-meter UI. bb_phase is 1-8 (1=write 0xaa, 2=verify 0xaa, 3=write # 0x55, 4=verify 0x55, 5=write 0xff, 6=verify 0xff, 7=write 0x00, # 8=verify 0x00). bb_phase_pct is 0-100 within the current phase. "ALTER TABLE burnin_stages ADD COLUMN bb_phase INTEGER", "ALTER TABLE burnin_stages ADD COLUMN bb_phase_pct REAL", # 1.0.0-46: live write/read throughput for the per-pattern meters. # Computed from successive `XX% done` lines in badblocks output: # delta_bytes = (overall_pct_delta / 800) * drive_size_bytes. # Updated on every progress line; NULL until the second progress # line arrives (need two samples to compute a rate). "ALTER TABLE burnin_stages ADD COLUMN bb_mbps REAL", # 1.0.0-47: per-pattern duration history. JSON map of # {"1": "2026-05-09T05:39:44+00:00", "2": ..., ...} where each key # is the phase number (1-8) and the value is when the parser first # observed that phase. Drawer derives "0xaa: 14h 22m" by diffing # consecutive phase-1 keys. "ALTER TABLE burnin_stages ADD COLUMN bb_phase_history TEXT", # 1.0.0-19: enforce one active burn-in per drive at the storage layer. # Closes the read-then-insert race in burnin.start_job — without this, # two concurrent /api/v1/burnin/start requests for the same drive could # both observe zero active jobs and both insert queued rows. """CREATE UNIQUE INDEX IF NOT EXISTS uniq_active_burnin_per_drive ON burnin_jobs (drive_id) WHERE state IN ('queued', 'running')""", # 1.0.0-22: app-level login (username + bcrypt password) """CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, full_name TEXT, is_admin INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL, last_login_at TEXT )""", # 1.0.0-28: case-insensitive uniqueness. The base UNIQUE on username # is case-sensitive but login does NOCASE — without this index two # users `Admin` and `admin` could coexist and shadow each other. """CREATE UNIQUE INDEX IF NOT EXISTS uniq_users_username_nocase ON users (username COLLATE NOCASE)""", ] async def _run_migrations(db: aiosqlite.Connection) -> None: for sql in _MIGRATIONS: try: await db.execute(sql) except Exception: pass # Column already exists — harmless # Remove the old CHECK(profile IN ('quick','full')) constraint if present. # SQLite can't ALTER a CHECK — requires a full table rebuild. cur = await db.execute( "SELECT sql FROM sqlite_master WHERE type='table' AND name='burnin_jobs'" ) row = await cur.fetchone() if row and "CHECK" in (row[0] or ""): await db.executescript(""" PRAGMA foreign_keys=OFF; CREATE TABLE burnin_jobs_new ( id INTEGER PRIMARY KEY AUTOINCREMENT, drive_id INTEGER NOT NULL REFERENCES drives(id), profile TEXT NOT NULL, state TEXT NOT NULL DEFAULT 'queued', percent INTEGER DEFAULT 0, stage_name TEXT, operator TEXT NOT NULL, created_at TEXT NOT NULL, started_at TEXT, finished_at TEXT, error_text TEXT ); INSERT INTO burnin_jobs_new SELECT * FROM burnin_jobs; DROP TABLE burnin_jobs; ALTER TABLE burnin_jobs_new RENAME TO burnin_jobs; CREATE INDEX IF NOT EXISTS idx_burnin_jobs_drive ON burnin_jobs(drive_id, state); PRAGMA foreign_keys=ON; """) async def init_db() -> None: Path(settings.db_path).parent.mkdir(parents=True, exist_ok=True) async with aiosqlite.connect(settings.db_path) as db: await db.execute("PRAGMA journal_mode=WAL") await db.execute("PRAGMA busy_timeout=60000") await db.execute("PRAGMA foreign_keys=ON") await db.executescript(SCHEMA) await _run_migrations(db) await db.commit() async def get_db(): db = await aiosqlite.connect(settings.db_path) db.row_factory = aiosqlite.Row try: await db.execute("PRAGMA journal_mode=WAL") await db.execute("PRAGMA busy_timeout=60000") await db.execute("PRAGMA foreign_keys=ON") yield db finally: await db.close()