SPB Git

spb/firc0001_uqo

Public
4.3 KB · 117 lines python
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
Files · e80ef398ede4f5784e4d15d7b4453d5f38c765cc
Menu

Theme