noaa-goes-visualization/suvi/db.py

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