spb/firc0001_uqo
Public
1"""Shared helpers for the FIRC-0001 Stage D read-only audits (run on M3U96b).23Every script opens the two DuckDB files with read_only=True, SET threads=4,4SET memory_limit='24GB', writes CSV/JSON under ~/firc0001_review/results/<name>/5and a manifest.json (script sha256, DuckDB mtimes, start/end, row counts).6Nothing is ever written under ~/UQO_RESEARCH_DB.7"""8from __future__ import annotations910import csv11import datetime as dt12import hashlib13import json14import os15import sys16import time17from pathlib import Path1819import duckdb2021HOME = Path(os.environ.get("HOME", "~")).expanduser()22RDB = HOME / "UQO_RESEARCH_DB"23SPID_DB = RDB / "project-db" / "data" / "db" / "spid.duckdb"24EDGAR_DB = RDB / "data" / "duckdb" / "edgar_research.duckdb"25WORK = HOME / "firc0001_review"26RESULTS = WORK / "results"27THREADS = int(os.environ.get("FIRC_THREADS", "4"))28MEMLIMIT = os.environ.get("FIRC_MEMLIMIT", "24GB")29SEED = 20261105303132def connect(path: Path = SPID_DB) -> duckdb.DuckDBPyConnection:33 con = duckdb.connect(str(path), read_only=True)34 con.execute(f"SET threads={THREADS}")35 con.execute(f"SET memory_limit='{MEMLIMIT}'")36 con.execute("SET preserve_insertion_order=false")37 # spill files must never land next to the databases (default = <db>.tmp)38 (WORK / "tmp").mkdir(parents=True, exist_ok=True)39 con.execute(f"SET temp_directory='{WORK / 'tmp'}'")40 return con414243def attach_edgar(con: duckdb.DuckDBPyConnection, alias: str = "edgar") -> None:44 con.execute(f"ATTACH '{EDGAR_DB}' AS {alias} (READ_ONLY)")454647def sha256_file(p: Path) -> str:48 h = hashlib.sha256()49 with open(p, "rb") as f:50 for chunk in iter(lambda: f.read(1 << 20), b""):51 h.update(chunk)52 return h.hexdigest()535455class Run:56 """Context for one audit script: output dir, timing, manifest."""5758 def __init__(self, name: str, script: str | None = None, out: str | None = None):59 self.name = name60 self.script = Path(script or sys.argv[0]).resolve()61 self.out = Path(out) if out else RESULTS / name62 self.out.mkdir(parents=True, exist_ok=True)63 self.t0 = time.time()64 self.started = dt.datetime.now().isoformat(timespec="seconds")65 self.rows: dict[str, int] = {}66 self.notes: dict = {}6768 def write_csv(self, fname: str, cols: list[str], rows) -> int:69 rows = list(rows)70 with open(self.out / fname, "w", newline="") as f:71 w = csv.writer(f)72 w.writerow(cols)73 w.writerows(rows)74 self.rows[fname] = len(rows)75 print(f"[{self.name}] {fname}: {len(rows)} rows", flush=True)76 return len(rows)7778 def dump(self, fname: str, con, sql: str, params=None):79 cur = con.execute(sql, params) if params else con.execute(sql)80 cols = [d[0] for d in cur.description]81 return self.write_csv(fname, cols, cur.fetchall())8283 def write_json(self, fname: str, obj) -> None:84 with open(self.out / fname, "w") as f:85 json.dump(obj, f, indent=2, default=str)86 print(f"[{self.name}] {fname} written", flush=True)8788 def finish(self, extra: dict | None = None) -> None:89 man = {90 "package": self.name,91 "script": str(self.script),92 "script_sha256": sha256_file(self.script) if self.script.exists() else None,93 "started": self.started,94 "ended": dt.datetime.now().isoformat(timespec="seconds"),95 "wall_seconds": round(time.time() - self.t0, 1),96 "node": os.uname().nodename,97 "python": sys.version.split()[0],98 "duckdb": duckdb.__version__,99 "threads": THREADS,100 "memory_limit": MEMLIMIT,101 "seed": SEED,102 "databases": {103 str(SPID_DB): {"bytes": SPID_DB.stat().st_size,104 "mtime": dt.datetime.fromtimestamp(SPID_DB.stat().st_mtime).isoformat()},105 str(EDGAR_DB): {"bytes": EDGAR_DB.stat().st_size,106 "mtime": dt.datetime.fromtimestamp(EDGAR_DB.stat().st_mtime).isoformat()},107 },108 "read_only": True,109 "row_counts": self.rows,110 "notes": self.notes,111 }112 if extra:113 man.update(extra)114 with open(self.out / "manifest.json", "w") as f:115 json.dump(man, f, indent=2, default=str)116 print(f"[{self.name}] done in {man['wall_seconds']} s", flush=True)117