fetch.py

#! python3

"""Fetch Homebrew analytics (formula + cask install counts) and store them.

Each normal run does three things:
  1. Keeps the raw JSON payloads as zstd-compressed dumps on disk (unchanged
     from the original script's behaviour).
  2. Upserts the parsed (date, kind, name, count) rows into a SQLite
     database, with names normalized into their own table so the counts
     table only ever stores integers.
  3. Refreshes the latest package descriptions, display names and homepages.
     Metadata is stored only in SQLite; no metadata dumps are written.

Usage:
    python fetch.py              # fetch today's data (or reuse today's dump
                                  # if it's already on disk) and upsert it
    python fetch.py --backfill   # rebuild the DB from every dump already in
                                  # dumps/, instead of hitting the network
    python fetch.py --backfill-totals  # populate only totals from existing dumps
    python fetch.py --metadata-only  # refresh metadata without fetching stats
"""

import argparse
import json
import sqlite3
import sys
import urllib.request
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from datetime import date, datetime, timezone
from pathlib import Path

from compression import zstd

SCRIPT_DIR = Path(__file__).resolve().parent
DUMPS_DIR = SCRIPT_DIR / "dumps"
DB_PATH = SCRIPT_DIR / "brew_stats.db"

COMPRESSION_LEVEL = 6
MIN_INSTALLS = 500


@dataclass(frozen=True)
class KindSpec:
    kind_id: int
    name: str  # "formula" | "cask" — also the key used in each item dict
    stub: str  # URL path segment on formulae.brew.sh


KIND_SPECS: tuple[KindSpec, ...] = (
    KindSpec(0, "formula", "install-on-request"),
    KindSpec(1, "cask", "cask-install"),
)

Record = tuple[int, str, int]  # (kind_id, name, count)


# --------------------------------------------------------------------------
# Small helpers
# --------------------------------------------------------------------------


def parse_num(formatted_string: str) -> int:
    return int(formatted_string.replace(",", ""))


def dump_path(dumps_dir: Path, date_str: str, kind_name: str) -> Path:
    return dumps_dir / f"{date_str}{kind_name}.zst"


# --------------------------------------------------------------------------
# Dump (raw JSON) I/O
# --------------------------------------------------------------------------


def read_dump(path: Path) -> dict:
    with zstd.open(path, "rt", encoding="utf-8") as f:
        return json.load(f)


def write_dump(path: Path, content: str, level: int = COMPRESSION_LEVEL) -> None:
    with zstd.open(path, "wt", level=level) as f:
        f.write(content)
    ratio = len(content) / path.stat().st_size
    print(f"Dumped {path} (compression ratio {ratio:.1f}x)")


def fetch_json(url: str) -> str:
    req = urllib.request.Request(url)
    with urllib.request.urlopen(req, timeout=60) as response:
        return response.read().decode("utf-8")


def fetch_and_dump(kind: KindSpec, dumps_dir: Path) -> tuple[str, dict]:
    """Fetch a kind's latest 30d data from the API, dump it, and return it."""

    url = f"https://formulae.brew.sh/api/analytics/{kind.stub}/30d.json"
    print(f"Fetching: {kind.name} data ----------------")
    content = fetch_json(url)
    data = json.loads(content)
    end_date = data.get("end_date", "")
    date_str = end_date.replace("-", "")
    print(f"Downloaded {len(content)} bytes for {end_date}")

    path = dump_path(dumps_dir, date_str, kind.name)
    if path.exists():
        print(f"{path} already exists, not overwriting.")
    else:
        write_dump(path, content)
    return date_str, data


def get_today_data(kind: KindSpec, dumps_dir: Path) -> tuple[str, dict]:
    """Reuse today's dump if we already have it, otherwise fetch it."""
    today_str = date.today().strftime("%Y%m%d")
    path = dump_path(dumps_dir, today_str, kind.name)
    if path.exists():
        print(f"{path} already exists.")
        return today_str, read_dump(path)
    return fetch_and_dump(kind, dumps_dir)


# --------------------------------------------------------------------------
# Parsing
# --------------------------------------------------------------------------


def parse_items(data: dict, kind: KindSpec, min_installs: int) -> Iterator[Record]:
    for item in data["items"]:
        name = item[kind.name]
        count = parse_num(item["count"])
        if count >= min_installs:
            yield (kind.kind_id, name, count)


# --------------------------------------------------------------------------
# SQLite layer
# --------------------------------------------------------------------------


def init_db(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS names (
            name_id INTEGER PRIMARY KEY AUTOINCREMENT,
            kind_id INTEGER NOT NULL,
            name    TEXT    NOT NULL,
            UNIQUE (kind_id, name)
        );

        CREATE TABLE IF NOT EXISTS counts (
            date    TEXT    NOT NULL,
            kind_id INTEGER NOT NULL,
            name_id INTEGER NOT NULL REFERENCES names (name_id),
            count   INTEGER NOT NULL,
            UNIQUE (date, kind_id, name_id)
        );

        CREATE INDEX IF NOT EXISTS idx_counts_name_id ON counts (name_id);
        """
    )
    columns = {row[1] for row in conn.execute("PRAGMA table_info(names)")}
    for column in ("display_name", "description", "homepage", "metadata_refreshed_at"):
        if column not in columns:
            conn.execute(f"ALTER TABLE names ADD COLUMN {column} TEXT")
    conn.commit()
    init_totals_db(conn)


def init_totals_db(conn: sqlite3.Connection) -> None:
    with conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS totals (
                date TEXT PRIMARY KEY NOT NULL,
                formula_total INTEGER,
                cask_total INTEGER
            )
        """)


def parse_total(data: dict) -> int:
    total = data.get("total_count")
    if type(total) is not int or total < 0:
        raise ValueError("Missing or invalid total_count")
    return total


def upsert_total(
    conn: sqlite3.Connection, date_str: str, kind: KindSpec, total: int
) -> None:
    column = {0: "formula_total", 1: "cask_total"}[kind.kind_id]
    conn.execute(
        f"""INSERT INTO totals (date, {column}) VALUES (?, ?)
            ON CONFLICT (date) DO UPDATE SET {column} = excluded.{column}
            WHERE totals.{column} IS NOT excluded.{column}""",
        (date_str, total),
    )


def metadata_index(items: object, kind: KindSpec) -> dict[str, list[dict]]:
    """Validate a source and index exact identities and aliases within its kind."""
    if not isinstance(items, list) or not items:
        raise ValueError(f"Empty or invalid {kind.name} metadata response")
    index: dict[str, list[dict]] = {}
    for item in items:
        if not isinstance(item, dict):
            raise ValueError(f"Invalid {kind.name} metadata item")
        name = item.get("token" if kind.name == "cask" else "name")
        if not isinstance(name, str) or not name:
            raise ValueError(f"Missing {kind.name} package identity")
        for field in ("desc", "homepage", "tap", "full_name", "full_token"):
            if item.get(field) is not None and not isinstance(item[field], str):
                raise ValueError(f"Invalid {field} for {name}")
        display = item.get("name")
        if kind.name == "cask" and (
            not isinstance(display, list) or not all(isinstance(v, str) for v in display)
        ):
            raise ValueError(f"Invalid display names for {name}")
        keys = {name, item.get("full_name"), item.get("full_token")}
        for field in ("aliases", "oldnames", "old_tokens"):
            values = item.get(field, [])
            if not isinstance(values, list) or not all(isinstance(v, str) for v in values):
                raise ValueError(f"Invalid {field} for {name}")
            keys.update(values)
        if item.get("tap"):
            keys.add(f"{item['tap']}/{name}")
        for key in keys - {None, ""}:
            index.setdefault(key, []).append(item)
    return index


def refresh_metadata(conn: sqlite3.Connection) -> int:
    """Fetch both sources before writing; failures leave previous metadata intact."""
    sources = {
        kind.kind_id: json.loads(fetch_json(f"https://formulae.brew.sh/api/{kind.name}.json"))
        for kind in KIND_SPECS
    }
    refreshed_at = datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")
    return update_metadata(conn, sources, refreshed_at)


def update_metadata(conn: sqlite3.Connection, sources: dict, refreshed_at: str) -> int:
    """Validate both sources and atomically apply uniquely matched metadata."""
    indexes = {kind.kind_id: metadata_index(sources[kind.kind_id], kind) for kind in KIND_SPECS}
    rows = []
    for name_id, kind_id, name in conn.execute("SELECT name_id, kind_id, name FROM names"):
        matches = indexes[kind_id].get(name, [])
        if len(matches) != 1:
            continue  # Never guess for absent packages or ambiguous aliases.
        item = matches[0]
        display = item["name"]
        if isinstance(display, list):
            display = ", ".join(display)
        rows.append((display, item.get("desc"), item.get("homepage"), refreshed_at, name_id))
    with conn:
        conn.executemany(
            """UPDATE names SET display_name=?, description=?, homepage=?,
               metadata_refreshed_at=? WHERE name_id=?""", rows
        )
    return len(rows)


def load_name_map(conn: sqlite3.Connection) -> dict[tuple[int, str], int]:
    rows = conn.execute("SELECT name_id, kind_id, name FROM names")
    return {(kind_id, name): name_id for name_id, kind_id, name in rows}


def ensure_names(
    conn: sqlite3.Connection,
    records: Iterable[Record],
    name_map: dict[tuple[int, str], int],
) -> dict[tuple[int, str], int]:
    """Make sure every (kind_id, name) in records has a row in `names`.

    Mutates and returns name_map with any newly created ids added.
    """
    missing = {
        (kind_id, name)
        for kind_id, name, _ in records
        if (kind_id, name) not in name_map
    }
    if not missing:
        return name_map

    conn.executemany(
        "INSERT OR IGNORE INTO names (kind_id, name) VALUES (?, ?)",
        sorted(missing),
    )
    rows = conn.execute(
        "SELECT name_id, kind_id, name FROM names WHERE (kind_id, name) IN "
        f"({', '.join('(?, ?)' for _ in missing)})",
        [value for pair in missing for value in pair],
    )
    for name_id, kind_id, name in rows:
        name_map[(kind_id, name)] = name_id
    return name_map


def upsert_counts(
    conn: sqlite3.Connection,
    date_str: str,
    records: Iterable[Record],
    name_map: dict[tuple[int, str], int],
) -> int:
    rows = [
        (date_str, kind_id, name_map[(kind_id, name)], count)
        for kind_id, name, count in records
    ]
    conn.executemany(
        """
        INSERT INTO counts (date, kind_id, name_id, count)
        VALUES (?, ?, ?, ?)
        ON CONFLICT (date, kind_id, name_id) DO UPDATE SET count = excluded.count
        WHERE counts.count <> excluded.count
        """,
        rows,
    )
    return len(rows)


def store_records(
    conn: sqlite3.Connection,
    date_str: str,
    records: list[Record],
    name_map: dict[tuple[int, str], int],
) -> dict[tuple[int, str], int]:
    """Ensure names exist, upsert counts, and commit. Returns updated name_map."""
    name_map = ensure_names(conn, records, name_map)
    n = upsert_counts(conn, date_str, records, name_map)
    conn.commit()
    return name_map, n


def store_snapshot(
    conn: sqlite3.Connection,
    date_str: str,
    data: dict,
    kind: KindSpec,
    min_installs: int,
    name_map: dict[tuple[int, str], int],
) -> tuple[dict[tuple[int, str], int], int]:
    """Commit a source's unfiltered total and filtered package counts together."""
    total = parse_total(data)
    records = list(parse_items(data, kind, min_installs))
    # Keep the caller's map unchanged if the transaction rolls back.
    updated_map = name_map.copy()
    with conn:
        ensure_names(conn, records, updated_map)
        n = upsert_counts(conn, date_str, records, updated_map)
        upsert_total(conn, date_str, kind, total)
    return updated_map, n


# --------------------------------------------------------------------------
# Run modes
# --------------------------------------------------------------------------


def run_fetch(conn: sqlite3.Connection, dumps_dir: Path, min_installs: int) -> None:
    """Normal run: fetch (or reuse) today's dumps and upsert them."""
    name_map = load_name_map(conn)
    for kind in KIND_SPECS:
        date_str, data = get_today_data(kind, dumps_dir)
        name_map, n = store_snapshot(conn, date_str, data, kind, min_installs, name_map)
        print(f"{kind.name}: upserted {n} rows for {date_str}")


def run_backfill(conn: sqlite3.Connection, dumps_dir: Path, min_installs: int) -> None:
    """Rebuild counts from every dump already on disk, no network calls."""
    name_map = load_name_map(conn)
    for kind in KIND_SPECS:
        paths = sorted(dumps_dir.glob(f"*{kind.name}.zst"))
        if not paths:
            print(f"No dumps found for {kind.name}.")
            continue
        for path in paths:
            date_str = path.name.removesuffix(f"{kind.name}.zst")
            data = read_dump(path)
            name_map, n = store_snapshot(conn, date_str, data, kind, min_installs, name_map)
            print(f"[backfill] {path.name}: upserted {n} rows for {date_str}")


def run_totals_backfill(conn: sqlite3.Connection, dumps_dir: Path) -> None:
    """Populate only totals from archived snapshots, without network calls."""
    for kind in KIND_SPECS:
        for path in sorted(dumps_dir.glob(f"*{kind.name}.zst")):
            date_str = path.name.removesuffix(f"{kind.name}.zst")
            total = parse_total(read_dump(path))
            with conn:
                upsert_total(conn, date_str, kind, total)
            print(f"[totals backfill] {path.name}: {total}")


# --------------------------------------------------------------------------
# Entry point
# --------------------------------------------------------------------------


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fetch Homebrew analytics into SQLite."
    )
    modes = parser.add_mutually_exclusive_group()
    modes.add_argument("--metadata-only", action="store_true", help="Refresh metadata without fetching analytics.")
    modes.add_argument(
        "--backfill",
        action="store_true",
        help="Rebuild counts from every existing dump instead of fetching new data.",
    )
    modes.add_argument(
        "--backfill-totals",
        action="store_true",
        help="Populate only totals from existing dumps, leaving other tables untouched.",
    )
    parser.add_argument(
        "--db",
        type=Path,
        default=DB_PATH,
        help=f"SQLite database path (default: {DB_PATH.name}).",
    )
    parser.add_argument(
        "--min-installs",
        type=int,
        default=MIN_INSTALLS,
        help=f"Minimum install count to keep a row (default: {MIN_INSTALLS}).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    DUMPS_DIR.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(args.db)
    try:
        if args.backfill_totals:
            init_totals_db(conn)
            run_totals_backfill(conn, DUMPS_DIR)
            return
        init_db(conn)
        if args.backfill:
            run_backfill(conn, DUMPS_DIR, args.min_installs)
        elif args.metadata_only:
            print(f"Refreshed metadata for {refresh_metadata(conn)} packages")
        else:
            run_fetch(conn, DUMPS_DIR, args.min_installs)
            try:
                print(f"Refreshed metadata for {refresh_metadata(conn)} packages")
            except (OSError, ValueError) as error:
                print(f"Metadata refresh failed; retained previous metadata: {error}", file=sys.stderr)
    finally:
        conn.close()


if __name__ == "__main__":
    main()