"""SQLite index of the SUVI archive: frames, header metadata, and detection results. This replaces the ``_f``/``_e`` filename suffixes the pipeline used to carry its verdicts in. Keeping labels out of filenames means a detector can be re-run or re-tuned without touching the archive, and several detectors can disagree about the same frame without anyone having to pick a winner on disk. Everything except ``detector_run``/``detection`` is reconstructible: ``frame`` from a directory walk, ``header`` from a cheap re-read. The database is therefore safe to delete, and the bench can point at a scratch copy via ``SUVI_DB``. Detection rows store the *continuous scores* a detector produced, not just its verdict, so sweeping a threshold is a query rather than another pass over 1.6M files. Writes are expected to come from a single process. Worker processes read, and push results down a queue to one writer (see ``writer_loop``); SQLite tolerates many readers but only one writer, and this archive lives on a shared mount where lock contention is expensive. """ import json import os import queue import sqlite3 import time from . import paths SCHEMA_VERSION = 2 #: Versions that upgrade to the current one by nothing more than running the schema #: script again. Every change so far has been a new table, and every table is #: declared IF NOT EXISTS, so an older index gains it on the next open. A change #: that alters or drops an existing column must not be added here -- it needs a real #: migration step instead. ADDITIVE_UPGRADES_FROM = (1,) #: Header keywords cached in the ``header`` table. Each entry maps a column name to #: the FITS keyword it comes from and its SQLite type. Defined once so inserts and #: reads cannot drift apart. HEADER_COLUMNS = ( ("empty", "EMPTY", "INTEGER"), ("degraded", "DEGRADED", "INTEGER"), ("eclipse", "ECLIPSE", "INTEGER"), ("num_imgs", "NUM_IMGS", "INTEGER"), ("n_long", "N_LONG", "INTEGER"), ("n_short", "N_SHORT", "INTEGER"), ("n_sh_fl", "N_SH_FL", "INTEGER"), ("exptime", "EXPTIME", "REAL"), ("img_min", "IMG_MIN", "REAL"), ("img_max", "IMG_MAX", "REAL"), ("img_mean", "IMG_MEAN", "REAL"), ("img_sdev", "IMG_SDEV", "REAL"), ("imgtii", "IMGTII", "REAL"), ("imgtir", "IMGTIR", "REAL"), ("diam_sun", "DIAM_SUN", "REAL"), ("crpix1", "CRPIX1", "REAL"), ("crpix2", "CRPIX2", "REAL"), ("crota", "CROTA", "REAL"), ("cdelt1", "CDELT1", "REAL"), ("cdelt2", "CDELT2", "REAL"), ("yaw_flip", "YAW_FLIP", "INTEGER"), ("solar_b0", "SOLAR_B0", "REAL"), ("dsun_obs", "DSUN_OBS", "REAL"), ("wavelnth", "WAVELNTH", "INTEGER"), ("date_beg", "DATE-BEG", "TEXT"), ("date_obs", "DATE-OBS", "TEXT"), ("date_end", "DATE-END", "TEXT"), ("datasum", "DATASUM", "TEXT"), ("checksum", "CHECKSUM", "TEXT"), ) HEADER_FIELDS = tuple(name for name, _, _ in HEADER_COLUMNS) _HEADER_DDL = ",\n ".join(f"{name} {sqltype}" for name, _, sqltype in HEADER_COLUMNS) SCHEMA = f""" CREATE TABLE IF NOT EXISTS meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL ); -- One row per FITS file present in the archive. `path` is archive-relative and -- always uses forward slashes, so the index is portable between Windows and Linux. CREATE TABLE IF NOT EXISTS frame ( id INTEGER PRIMARY KEY, path TEXT NOT NULL UNIQUE, satellite INTEGER NOT NULL, wavelength INTEGER NOT NULL, t_start INTEGER NOT NULL, t_end INTEGER NOT NULL, version TEXT, size_bytes INTEGER, mtime REAL ); CREATE UNIQUE INDEX IF NOT EXISTS frame_slot ON frame (satellite, wavelength, t_start); CREATE INDEX IF NOT EXISTS frame_time ON frame (t_start); -- Cached FITS header metadata. Populated by a ~18 KB read per file; no pixel -- decompression. read_ok=0 means the file could not be parsed, and read_error says -- why -- itself a strong bad-frame signal. CREATE TABLE IF NOT EXISTS header ( frame_id INTEGER PRIMARY KEY REFERENCES frame (id) ON DELETE CASCADE, {_HEADER_DDL}, read_ok INTEGER NOT NULL, read_error TEXT, scanned_at REAL NOT NULL ); -- One row per (detector, configuration) execution, so results from different -- methods and different thresholds coexist and can be compared. CREATE TABLE IF NOT EXISTS detector_run ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, config_json TEXT NOT NULL, code_version TEXT, created_at REAL NOT NULL, notes TEXT ); CREATE INDEX IF NOT EXISTS detector_run_name ON detector_run (name); CREATE TABLE IF NOT EXISTS detection ( run_id INTEGER NOT NULL REFERENCES detector_run (id) ON DELETE CASCADE, frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, verdict TEXT NOT NULL, reason TEXT, scores_json TEXT, elapsed_us INTEGER, PRIMARY KEY (run_id, frame_id) ) WITHOUT ROWID; CREATE INDEX IF NOT EXISTS detection_frame ON detection (frame_id); -- One row per day-directory, holding the mtime it had when we last read it. -- -- This is what makes re-indexing cheap. A directory's mtime changes whenever an -- entry is added or removed, so comparing it is an exact test for "did anything -- change in here" -- no guessing, no walking the files inside. The archive holds -- ~360 frames per day-directory, so checking ~7,400 directory stats replaces -- 2.65M file lookups: a 360x reduction, and the difference between an operation -- this filesystem sustains and one that exhausts it. CREATE TABLE IF NOT EXISTS dir_scan ( path TEXT PRIMARY KEY, mtime REAL NOT NULL, n_frames INTEGER NOT NULL, scanned_at REAL NOT NULL ) WITHOUT ROWID; -- Download bookkeeping, migrated out of the 800 MB file_database.json that -- puller_fits.py used to parse into memory on every run. CREATE TABLE IF NOT EXISTS remote_file ( url TEXT PRIMARY KEY, remote_mtime REAL NOT NULL, remote_size INTEGER, local_path TEXT, fetched_at REAL ) WITHOUT ROWID; -- ---------------------------------------------------------------- test bench ---- -- A stretch of archive vetted as known-good, frozen so that repeated experiments -- are measured against the same ground truth. CREATE TABLE IF NOT EXISTS bench_window ( id INTEGER PRIMARY KEY, name TEXT NOT NULL UNIQUE, t_start INTEGER NOT NULL, t_end INTEGER NOT NULL, satellites TEXT NOT NULL, wavelengths TEXT NOT NULL, vetted_at REAL, notes TEXT ); CREATE TABLE IF NOT EXISTS bench_truth ( window_id INTEGER NOT NULL REFERENCES bench_window (id) ON DELETE CASCADE, frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, is_good INTEGER NOT NULL, -- 'legacy' (old filename suffix), 'header' (cross-check), or 'manual'. vet_source TEXT NOT NULL, PRIMARY KEY (window_id, frame_id) ) WITHOUT ROWID; -- One experiment: a window plus a seeded plan of deletions and corruptions. CREATE TABLE IF NOT EXISTS bench_case ( id INTEGER PRIMARY KEY, window_id INTEGER NOT NULL REFERENCES bench_window (id) ON DELETE CASCADE, name TEXT NOT NULL UNIQUE, seed INTEGER NOT NULL, plan_json TEXT NOT NULL, overlay_dir TEXT NOT NULL, created_at REAL NOT NULL ); -- What was done to each affected slot. mode='delete' means the frame was withheld -- entirely; anything else names a corruption from suvi.corruptions.CATALOG. CREATE TABLE IF NOT EXISTS bench_injection ( case_id INTEGER NOT NULL REFERENCES bench_case (id) ON DELETE CASCADE, frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, mode TEXT NOT NULL, severity REAL NOT NULL DEFAULT 1.0, params_json TEXT, overlay_path TEXT, gap_index INTEGER NOT NULL DEFAULT 0, gap_length INTEGER NOT NULL DEFAULT 1, PRIMARY KEY (case_id, frame_id) ) WITHOUT ROWID; -- Fill quality, one row per reconstructed frame, so results survive between runs. CREATE TABLE IF NOT EXISTS bench_fill_result ( case_id INTEGER NOT NULL REFERENCES bench_case (id) ON DELETE CASCADE, frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, filler TEXT NOT NULL, source TEXT NOT NULL, -- 'oracle' or a detector name scores_json TEXT NOT NULL, PRIMARY KEY (case_id, frame_id, filler, source) ) WITHOUT ROWID; """ #: Verdicts a detector may record. 'unknown' covers frames a detector declined to #: judge (e.g. a temporal detector with no usable neighbours), which must not be #: silently conflated with 'good'. VERDICTS = ("good", "bad", "unknown") def connect(path=None, readonly=False, timeout=60.0): """Open the index, creating and initialising it if needed. WAL is preferred but not required: this archive lives on a shared mount whose locking primitives may not support it, so a failure to enter WAL falls back to the rollback journal rather than aborting. """ path = path or paths.default_db_path() if readonly: if not os.path.exists(path): raise FileNotFoundError(f"No SUVI index at {path}") conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=timeout) else: parent = os.path.dirname(os.path.abspath(path)) if parent: os.makedirs(parent, exist_ok=True) conn = sqlite3.connect(path, timeout=timeout) conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys = ON") conn.execute(f"PRAGMA busy_timeout = {int(timeout * 1000)}") if not readonly: try: conn.execute("PRAGMA journal_mode = WAL") except sqlite3.Error: pass conn.execute("PRAGMA synchronous = NORMAL") init_schema(conn) return conn def init_schema(conn): """Create the schema if absent, upgrading an older index in place. Idempotent.""" conn.executescript(SCHEMA) row = conn.execute("SELECT value FROM meta WHERE key = 'schema_version'").fetchone() if row is None: conn.execute( "INSERT INTO meta (key, value) VALUES ('schema_version', ?)", (str(SCHEMA_VERSION),), ) conn.commit() return found = int(row["value"]) if found == SCHEMA_VERSION: return if found in ADDITIVE_UPGRADES_FROM: # executescript above already created any new tables. conn.execute( "UPDATE meta SET value = ? WHERE key = 'schema_version'", (str(SCHEMA_VERSION),) ) conn.commit() return raise RuntimeError( f"Index schema version {found} cannot be upgraded to {SCHEMA_VERSION} " f"automatically; migrate or rebuild the index." ) def set_meta(conn, key, value): conn.execute( "INSERT INTO meta (key, value) VALUES (?, ?) " "ON CONFLICT (key) DO UPDATE SET value = excluded.value", (key, str(value)), ) def get_meta(conn, key, default=None): row = conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone() return row["value"] if row else default # --------------------------------------------------------------------------- frames def upsert_frame(conn, name, relpath, size_bytes=None, mtime=None): """Insert or update one frame row. Returns its id. Keyed on path, with the (satellite, wavelength, t_start) slot kept unique so a duplicate observation stored under a second version string is rejected loudly rather than silently shadowing the first. """ conn.execute( """ INSERT INTO frame (path, satellite, wavelength, t_start, t_end, version, size_bytes, mtime) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (path) DO UPDATE SET size_bytes = excluded.size_bytes, mtime = excluded.mtime """, ( relpath, name.satellite, name.wavelength, name.t_start, name.t_end, name.version, size_bytes, mtime, ), ) return conn.execute("SELECT id FROM frame WHERE path = ?", (relpath,)).fetchone()["id"] def frame_id_by_slot(conn, satellite, wavelength, t_start): row = conn.execute( "SELECT id FROM frame WHERE satellite = ? AND wavelength = ? AND t_start = ?", (satellite, wavelength, t_start), ).fetchone() return row["id"] if row else None def frames_in_dir(conn, dir_relpath): """Frames the index believes live in one directory: {filename: frame_id}.""" prefix = dir_relpath.rstrip("/") + "/" rows = conn.execute( "SELECT id, path FROM frame WHERE path >= ? AND path < ?", (prefix, prefix + "￿"), ).fetchall() return {row["path"].rsplit("/", 1)[-1]: row["id"] for row in rows} def delete_frames(conn, frame_ids): """Drop frames that have disappeared from disk, and their dependent rows.""" conn.executemany("DELETE FROM frame WHERE id = ?", [(i,) for i in frame_ids]) # ------------------------------------------------------------------ directory scans def get_dir_mtime(conn, dir_relpath): row = conn.execute("SELECT mtime FROM dir_scan WHERE path = ?", (dir_relpath,)).fetchone() return row["mtime"] if row else None def dir_mtimes(conn): """Every recorded directory mtime, as one dict. Read in a single query rather than per directory: the whole point of this table is to answer thousands of "has this changed?" questions without touching the filesystem, and it would be perverse to pay a round trip each time. """ return { row["path"]: row["mtime"] for row in conn.execute("SELECT path, mtime FROM dir_scan") } def record_dir_scans(conn, records): """`records` yields (dir_relpath, mtime, n_frames).""" now = time.time() conn.executemany( """ INSERT INTO dir_scan (path, mtime, n_frames, scanned_at) VALUES (?, ?, ?, ?) ON CONFLICT (path) DO UPDATE SET mtime = excluded.mtime, n_frames = excluded.n_frames, scanned_at = excluded.scanned_at """, [(path, mtime, count, now) for path, mtime, count in records], ) def forget_dir_scans(conn, dir_relpaths): conn.executemany("DELETE FROM dir_scan WHERE path = ?", [(p,) for p in dir_relpaths]) # ----------------------------------------------------------------------- detections def create_detector_run(conn, name, config, code_version=None, notes=None): """Register a new detector execution and return its id.""" cur = conn.execute( """ INSERT INTO detector_run (name, config_json, code_version, created_at, notes) VALUES (?, ?, ?, ?, ?) """, (name, json.dumps(config, sort_keys=True), code_version, time.time(), notes), ) conn.commit() return cur.lastrowid def latest_run_id(conn, name, config=None): """Most recent run of a named detector, or None. With `config`, only runs recorded under exactly that configuration match. This matters more than it looks: detector runs accumulate, one per detector per bench case, so "the latest run of header_v1" is whichever case was processed most recently -- not the one being asked about. Scoring one case's verdicts against another case's injections produces numbers that are wrong rather than merely imprecise, so callers working with a specific case must pass its config. """ if config is None: row = conn.execute( "SELECT id FROM detector_run WHERE name = ? " "ORDER BY created_at DESC, id DESC LIMIT 1", (name,), ).fetchone() else: row = conn.execute( "SELECT id FROM detector_run WHERE name = ? AND config_json = ? " "ORDER BY created_at DESC, id DESC LIMIT 1", (name, json.dumps(config, sort_keys=True)), ).fetchone() return row["id"] if row else None def record_detections(conn, run_id, results): """Write detection rows. `results` yields (frame_id, verdict, reason, scores, us).""" conn.executemany( """ INSERT INTO detection (run_id, frame_id, verdict, reason, scores_json, elapsed_us) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT (run_id, frame_id) DO UPDATE SET verdict = excluded.verdict, reason = excluded.reason, scores_json = excluded.scores_json, elapsed_us = excluded.elapsed_us """, [ ( run_id, frame_id, _check_verdict(verdict), reason, json.dumps(scores, sort_keys=True) if scores is not None else None, elapsed_us, ) for frame_id, verdict, reason, scores, elapsed_us in results ], ) def _check_verdict(verdict): if verdict not in VERDICTS: raise ValueError(f"Unknown verdict {verdict!r}; expected one of {VERDICTS}") return verdict def good_slots(conn, run_id, t_start, t_end, satellite, wavelengths=paths.WAVELENGTHS): """Frames a run judged good, in a half-open time window. This is the query that replaces ``merger_FITS.py``'s ``_f.fits$`` regex. """ placeholders = ",".join("?" * len(wavelengths)) rows = conn.execute( f""" SELECT f.id, f.path, f.satellite, f.wavelength, f.t_start FROM frame f JOIN detection d ON d.frame_id = f.id AND d.run_id = ? WHERE d.verdict = 'good' AND f.satellite = ? AND f.t_start >= ? AND f.t_start < ? AND f.wavelength IN ({placeholders}) ORDER BY f.t_start, f.wavelength """, (run_id, satellite, t_start, t_end, *wavelengths), ).fetchall() return rows def unscanned_frames(conn, limit=None): """Frames with no cached header yet.""" sql = """ SELECT f.id, f.path FROM frame f LEFT JOIN header h ON h.frame_id = f.id WHERE h.frame_id IS NULL ORDER BY f.t_start """ if limit is not None: sql += f" LIMIT {int(limit)}" return conn.execute(sql).fetchall() # --------------------------------------------------------------------------- headers _HEADER_INSERT = f""" INSERT INTO header (frame_id, {", ".join(HEADER_FIELDS)}, read_ok, read_error, scanned_at) VALUES (?{", ?" * len(HEADER_FIELDS)}, ?, ?, ?) ON CONFLICT (frame_id) DO UPDATE SET {", ".join(f"{f} = excluded.{f}" for f in HEADER_FIELDS)}, read_ok = excluded.read_ok, read_error = excluded.read_error, scanned_at = excluded.scanned_at """ def record_headers(conn, records): """Write cached header rows. `records` yields (frame_id, values_dict, error).""" now = time.time() conn.executemany( _HEADER_INSERT, [ ( frame_id, *(values.get(field) for field in HEADER_FIELDS), 0 if error else 1, error, now, ) for frame_id, values, error in records ], ) # ---------------------------------------------------------------------- remote files def get_remote_mtime(conn, url): row = conn.execute( "SELECT remote_mtime FROM remote_file WHERE url = ?", (url,) ).fetchone() return row["remote_mtime"] if row else None def record_remote_files(conn, records): """`records` yields (url, remote_mtime, remote_size, local_path, fetched_at).""" conn.executemany( """ INSERT INTO remote_file (url, remote_mtime, remote_size, local_path, fetched_at) VALUES (?, ?, ?, ?, ?) ON CONFLICT (url) DO UPDATE SET remote_mtime = excluded.remote_mtime, remote_size = excluded.remote_size, local_path = excluded.local_path, fetched_at = excluded.fetched_at """, list(records), ) # ---------------------------------------------------------------- single-writer loop #: Sentinel pushed onto a writer queue to shut it down. STOP = None def writer_loop(db_path, work_queue, batch_size=512, flush_seconds=5.0): """Drain (kind, payload) messages from a queue into the index. Runs in its own process so worker processes never contend for the write lock. Recognised kinds: 'detection', 'header', 'remote_file' -- each payload being the row tuple that the corresponding record_* helper expects. """ handlers = { "detection": _flush_detections, "header": record_headers, "remote_file": record_remote_files, } conn = connect(db_path) pending = {kind: [] for kind in handlers} last_flush = time.monotonic() def flush(): nonlocal last_flush for kind, rows in pending.items(): if rows: handlers[kind](conn, rows) rows.clear() conn.commit() last_flush = time.monotonic() try: while True: try: message = work_queue.get(timeout=1.0) except queue.Empty: if time.monotonic() - last_flush > flush_seconds: flush() continue if message is STOP: break kind, payload = message if kind not in pending: raise ValueError(f"Unknown writer message kind: {kind!r}") pending[kind].append(payload) if sum(len(v) for v in pending.values()) >= batch_size: flush() flush() finally: conn.commit() conn.close() def _flush_detections(conn, rows): """Adapter: writer payloads carry run_id per row, record_detections does not.""" by_run = {} for run_id, *rest in rows: by_run.setdefault(run_id, []).append(tuple(rest)) for run_id, results in by_run.items(): record_detections(conn, run_id, results)