db.py

"""NGUYEN-RH-1 research catalog — SQLite helper.

Usage:
    python db.py init
    python db.py log-experiment E-001 --name "first 10^4 zeros" \
        --script research/numerics/scripts/zeros_basic.py \
        --params '{"n_zeros": 10000}' --dps 50 \
        --csv research/numerics/data/E-001-zeros.csv \
        --json research/numerics/data/E-001-zeros.json --session 1
    python db.py list experiments
    python db.py add-result R-001 --title "..." --tag PROVED --track A \
        --file research/results/R-001-....md
"""
import argparse
import hashlib
import json
import sqlite3
import subprocess
from datetime import date
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[3]
DB_PATH = REPO_ROOT / "research" / "research.db"
SCHEMA = Path(__file__).with_name("schema.sql")

TABLES = ("sessions", "results", "graveyard", "derivations", "experiments")


def connect() -> sqlite3.Connection:
    con = sqlite3.connect(DB_PATH)
    con.row_factory = sqlite3.Row
    return con


def cmd_init(_args) -> None:
    con = connect()
    con.executescript(SCHEMA.read_text(encoding="utf-8"))
    con.commit()
    print(f"initialized {DB_PATH}")


def git_hash() -> str:
    try:
        return subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True, text=True, cwd=REPO_ROOT,
        ).stdout.strip()
    except OSError:
        return ""


def sha256_of(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def cmd_log_experiment(a) -> None:
    csv_path = REPO_ROOT / a.csv
    con = connect()
    con.execute(
        """INSERT OR REPLACE INTO experiments
           (id, name, script, git_hash, parameters, precision_dps,
            csv_file, json_file, row_count, checksum, session, date)
           VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
        (a.id, a.name, a.script, git_hash(), a.params, a.dps,
         a.csv, a.json, a.rows, sha256_of(csv_path), a.session,
         date.today().isoformat()),
    )
    con.commit()
    print(f"logged experiment {a.id}")


def cmd_add_result(a) -> None:
    con = connect()
    con.execute(
        """INSERT OR REPLACE INTO results
           (id, title, tag, track, file, date, dependencies)
           VALUES (?,?,?,?,?,?,?)""",
        (a.id, a.title, a.tag, a.track, a.file,
         date.today().isoformat(), a.deps),
    )
    con.commit()
    print(f"logged result {a.id}")


def cmd_add_derivation(a) -> None:
    con = connect()
    con.execute(
        """INSERT OR REPLACE INTO derivations
           (id, title, task, status, file, date) VALUES (?,?,?,?,?,?)""",
        (a.id, a.title, a.task, a.status, a.file, date.today().isoformat()),
    )
    con.commit()
    print(f"logged derivation {a.id}")


def cmd_add_graveyard(a) -> None:
    con = connect()
    con.execute(
        """INSERT OR REPLACE INTO graveyard
           (id, title, track, obstruction, file, date) VALUES (?,?,?,?,?,?)""",
        (a.id, a.title, a.track, a.obstruction, a.file, date.today().isoformat()),
    )
    con.commit()
    print(f"logged graveyard entry {a.id}")


def cmd_list(a) -> None:
    con = connect()
    for row in con.execute(f"SELECT * FROM {a.table}"):
        print(json.dumps(dict(row), ensure_ascii=False))


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("init").set_defaults(func=cmd_init)

    e = sub.add_parser("log-experiment")
    e.add_argument("id")
    e.add_argument("--name", required=True)
    e.add_argument("--script", required=True)
    e.add_argument("--params", default="{}")
    e.add_argument("--dps", type=int, default=50)
    e.add_argument("--csv", required=True)
    e.add_argument("--json", required=True)
    e.add_argument("--rows", type=int, default=0)
    e.add_argument("--session", type=int, default=0)
    e.set_defaults(func=cmd_log_experiment)

    r = sub.add_parser("add-result")
    r.add_argument("id")
    r.add_argument("--title", required=True)
    r.add_argument("--tag", required=True)
    r.add_argument("--track", default="")
    r.add_argument("--file", required=True)
    r.add_argument("--deps", default="")
    r.set_defaults(func=cmd_add_result)

    d = sub.add_parser("add-derivation")
    d.add_argument("id")
    d.add_argument("--title", required=True)
    d.add_argument("--task", default="")
    d.add_argument("--status", default="in-progress")
    d.add_argument("--file", required=True)
    d.set_defaults(func=cmd_add_derivation)

    g = sub.add_parser("add-graveyard")
    g.add_argument("id")
    g.add_argument("--title", required=True)
    g.add_argument("--track", default="")
    g.add_argument("--obstruction", required=True)
    g.add_argument("--file", required=True)
    g.set_defaults(func=cmd_add_graveyard)

    l = sub.add_parser("list")
    l.add_argument("table", choices=TABLES)
    l.set_defaults(func=cmd_list)

    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()