Blame for intel/export.py

"""
Exports the store as gzipped CSVs plus a manifest, for the Colab notebooks in
ml/notebooks/. The manifest carries Layer 1's freshness report, so every
notebook run states what data it was trained on and how old it was.
"""
import csv
import gzip
import json
import sqlite3
from pathlib import Path

from intel.db import now_iso
from intel.freshness import report

TABLES = {
    "locations": "select id, level, name, parent_id, lat, lon from locations order by id",
    "streams": "select * from streams order by id",
    "observations": """select stream_id, location_id, period_start, value from observations
                       where is_forecast = 0 order by stream_id, location_id, period_start""",
    "forecasts": """select stream_id, location_id, period_start, issued_at, value from observations
                    where is_forecast = 1 order by issued_at, stream_id, location_id, period_start""",
    "events": """select id, family, type, location_id, starts_on, ends_on, title, is_estimate, source_id
                 from events order by starts_on""",
    "places": """select id, name, category, subsector, confidence, operating_status, lat, lon, location_id, locality
                 from places order by location_id""",
}


def export(con: sqlite3.Connection, out_dir: Path) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    counts = {}
    for name, sql in TABLES.items():
        cur = con.execute(sql)
        cols = [d[0] for d in cur.description]
        n = 0
        with gzip.open(out_dir / f"{name}.csv.gz", "wt", newline="", encoding="utf8") as f:
            w = csv.writer(f)
            w.writerow(cols)
            for row in cur:
                w.writerow(row)
                n += 1
        counts[name] = n
    manifest = {"exported_at": now_iso(), "rows": counts, "freshness": report(con)}
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf8")
    return manifest