near-toll-audit / pipeline/run_bq.py · raw
#!/usr/bin/env python3"""Run pipeline/sql/*.sql against the NEAR BigQuery public dataset and snapshot results to raw/bigquery/.Needs GOOGLE_APPLICATION_CREDENTIALS (any Google account works: the sandbox gives 1 TB scanned a month, no card).Usage: run_bq.py [--only 01,02] [--cap-gb 200]Without credentials the pipeline still rebuilds from the committed snapshots."""import argparse, csv, datetime, gzip, hashlib, io, json, pathlib, sysHERE = pathlib.Path(__file__).resolve().parentROOT = HERE.parentsys.path.insert(0, str(HERE))from bq import BQ  # noqa: E402def main():    ap = argparse.ArgumentParser()    ap.add_argument("--only", default="")    ap.add_argument("--cap-gb", type=float, default=200)    ap.add_argument("--dry", action="store_true")    ap.add_argument("--total-cap-gb", type=float, default=300, help="stop before this many GB in one run")    args = ap.parse_args()    client = BQ(max_bytes=args.cap_gb * 1e9)    only = [x for x in args.only.split(",") if x]    outdir = ROOT / "raw" / "bigquery"    outdir.mkdir(parents=True, exist_ok=True)    total = 0    for path in sorted((HERE / "sql").glob("*.sql")):        if only and path.name.split("_")[0] not in only:            continue        sql = path.read_text()        if "{{INTENTS_TOKENS}}" in sql:            tokens = [t.strip() for t in (HERE / "sql" / "intents_tokens.txt").read_text().splitlines() if t.strip()]            sql = sql.replace("{{INTENTS_TOKENS}}", ",".join("'" + t.replace("'", "") + "'" for t in tokens))        if "{{TOKEN_PRICE_STRUCTS}}" in sql:            rows = list(csv.DictReader((HERE / "sql" / "token_prices.csv").open()))            structs = [f"STRUCT('{r['token'].replace(chr(39), '')}' AS token, {int(r['decimals'])} AS decimals, {float(r['price'])} AS price)" for r in rows]            sql = sql.replace("{{TOKEN_PRICE_STRUCTS}}", ", ".join(structs))        name = path.stem        if args.dry:            gb = client.dry_run(sql) / 1e9            total += gb            print(f"{name}: {gb:.2f} GB")            continue        if total + client.dry_run(sql) / 1e9 > args.total_cap_gb:            raise SystemExit(f"stopping before {name}: run budget of {args.total_cap_gb} GB would be passed")        rows, stats = client.query(sql, name)        buf = io.StringIO()        if rows:            wr = csv.DictWriter(buf, fieldnames=list(rows[0].keys()), lineterminator="\n")            wr.writeheader()            wr.writerows(rows)        raw = buf.getvalue().encode()        dest = outdir / f"{name}.csv.gz"        with gzip.GzipFile(dest, "wb", mtime=0) as fh:            fh.write(raw)        meta = {"file": dest.name, "pulled_at": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%MZ"),                "rows": len(rows), "sha256_csv": hashlib.sha256(raw).hexdigest(),                "sql_sha256": hashlib.sha256(sql.encode()).hexdigest(), "bytes_scanned": stats["bytes_processed"]}        with (outdir / "MANIFEST.jsonl").open("a", encoding="utf8") as fh:            fh.write(json.dumps(meta) + "\n")        total += stats["bytes_processed"] / 1e9        print(f"{name}: {len(rows)} rows, {stats['bytes_processed']/1e9:.2f} GB scanned, {dest.stat().st_size/1024:.0f} KB")    print(f"total {total:.2f} GB")if __name__ == "__main__":    main()