Add update-stats / get-stats commands (#41)

pgdevkit/cli.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import os
 import queue
 import threading
@@ -12,7 +13,7 @@
 from rich import box
 from tqdm import tqdm
 
-from . import migrate, testdb
+from . import migrate, stats, testdb
 from .backends import get_backend
 from .connection import build_conninfo
 from .diff import DiffKind, compute_diff
@@ -218,6 +219,46 @@ def fetch_missing(
     console.print(f"\nWrote {written} file(s).")
 
 
+@app.command("update-stats")
+def update_stats(
+    scripts_dir: Path = typer.Argument(..., help="The database/ folder; stats are written to <scripts_dir>/_stats/"),
+    url: str = typer.Option(..., "--url", help="PostgreSQL DSN (postgresql://user:pass@host:port/db)"),
+    entra_user: str | None = typer.Option(None, "--entra-user", help="Azure Entra user (triggers token auth)"),
+    table: list[str] = typer.Option([], "--table", help="Only update schema.name (repeatable); default is every table"),
+    exact: bool = typer.Option(False, "--exact", help="Use count(*) for row counts instead of the planner estimate"),
+    analyze: bool = typer.Option(False, "--analyze", help="Run ANALYZE first so column stats are fresh"),
+) -> None:
+    """Store table stats in _stats/_tables.json (keyed by schema.name, sorted)
+    and column stats in _stats/<schema.name>.json."""
+    if not scripts_dir.is_dir():
+        err_console.print(f"[red]Error:[/red] {scripts_dir} is not a directory")
+        raise typer.Exit(2)
+    conninfo = build_conninfo(url, entra_user)
+    try:
+        with console.status("Collecting stats..."):
+            tables, columns = stats.collect_stats(conninfo, set(table) or None, exact=exact, analyze=analyze)
+    except KeyError as e:
+        err_console.print(f"[red]Error:[/red] unknown table(s): {e.args[0]}")
+        raise typer.Exit(2)
+    path = stats.write_stats(scripts_dir, tables, columns, prune=not table)
+    console.print(f"Updated stats for {len(tables)} table(s) in {path.parent}")
+
+
+@app.command("get-stats")
+def get_stats(
+    scripts_dir: Path = typer.Argument(..., help="The database/ folder containing _stats/"),
+    tables: list[str] = typer.Argument(None, help="schema.name of the tables (default: all)"),
+    no_columns: bool = typer.Option(False, "--no-columns", help="Omit column stats"),
+) -> None:
+    """Print stored stats (read from the _stats JSON files, no DB access) as JSON."""
+    try:
+        result = stats.read_stats(scripts_dir, tables or [], columns=not no_columns)
+    except (FileNotFoundError, KeyError) as e:
+        err_console.print(f"[red]Error:[/red] {e.args[0]}")
+        raise typer.Exit(2)
+    typer.echo(json.dumps(result, indent=2, sort_keys=True))
+
+
 @testdb_app.command("up")
 def testdb_up(
     env: str = _TESTDB_ENV_OPTION,

pgdevkit/stats.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import psycopg
+import psycopg.sql
+from psycopg.rows import dict_row
+
+STATS_DIRNAME = "_stats"
+TABLES_FILE = "_tables.json"
+
+_TABLES_SQL = """
+SELECT n.nspname AS schema, c.relname AS name,
+       c.reltuples::bigint AS estimated_rows,
+       pg_table_size(c.oid) AS table_bytes,
+       pg_indexes_size(c.oid) AS index_bytes,
+       pg_total_relation_size(c.oid) AS total_bytes
+FROM pg_class c
+JOIN pg_namespace n ON n.oid = c.relnamespace
+WHERE c.relkind IN ('r', 'p')
+  AND n.nspname NOT LIKE 'pg\\_%' AND n.nspname != 'information_schema'
+ORDER BY 1, 2
+"""
+
+_COLUMNS_SQL = """
+SELECT a.attname AS name,
+       format_type(a.atttypid, a.atttypmod) AS data_type,
+       s.null_frac, s.n_distinct, s.avg_width
+FROM pg_attribute a
+LEFT JOIN pg_stats s
+  ON s.schemaname = %(schema)s AND s.tablename = %(table)s AND s.attname = a.attname
+WHERE a.attrelid = to_regclass(%(qualified)s) AND a.attnum > 0 AND NOT a.attisdropped
+ORDER BY a.attnum, s.inherited NULLS FIRST
+"""
+
+
+def stats_dir(scripts_dir: Path) -> Path:
+    return scripts_dir / STATS_DIRNAME
+
+
+def column_stats_path(scripts_dir: Path, qualified_name: str) -> Path:
+    return stats_dir(scripts_dir) / f"{qualified_name}.json"
+
+
+def _q(conn: Any, sql: Any, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+    with conn.cursor(row_factory=dict_row) as cur:
+        cur.execute(sql, params)
+        return cur.fetchall()
+
+
+def _write_json(path: Path, data: Any) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+
+def collect_stats(
+    conninfo: str, only: set[str] | None = None, exact: bool = False, analyze: bool = False
+) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]]]:
+    """Returns (table_stats, column_stats), both keyed by "schema.table".
+
+    Row counts come from pg_class.reltuples (null if the table was never
+    analyzed) unless exact=True, which runs count(*). analyze=True runs ANALYZE
+    first so column stats (null_frac/n_distinct/avg_width) are populated."""
+    tables: dict[str, dict[str, Any]] = {}
+    columns: dict[str, dict[str, Any]] = {}
+    with psycopg.connect(conninfo, autocommit=True) as conn:
+        rows = _q(conn, _TABLES_SQL)
+        if only is not None:
+            unknown = only - {f"{r['schema']}.{r['name']}" for r in rows}
+            if unknown:
+                raise KeyError(", ".join(sorted(unknown)))
+        for r in rows:
+            qn = f"{r['schema']}.{r['name']}"
+            if only is not None and qn not in only:
+                continue
+            ident = psycopg.sql.SQL("{}.{}").format(
+                psycopg.sql.Identifier(r["schema"]), psycopg.sql.Identifier(r["name"])
+            )
+            if analyze:
+                conn.execute(psycopg.sql.SQL("ANALYZE {}").format(ident))
+            if exact:
+                row_count = _q(conn, psycopg.sql.SQL("SELECT count(*) AS n FROM {}").format(ident))[0]["n"]
+            else:
+                row_count = r["estimated_rows"] if r["estimated_rows"] >= 0 else None
+            tables[qn] = {
+                "row_count": row_count,
+                "row_count_exact": exact,
+                "table_bytes": r["table_bytes"],
+                "index_bytes": r["index_bytes"],
+                "total_bytes": r["total_bytes"],
+            }
+            cols = _q(
+                conn,
+                _COLUMNS_SQL,
+                {"schema": r["schema"], "table": r["name"], "qualified": ident.as_string(conn)},
+            )
+            columns[qn] = {
+                c["name"]: {
+                    "data_type": c["data_type"],
+                    "null_fraction": c["null_frac"],
+                    "n_distinct": c["n_distinct"],
+                    "avg_width": c["avg_width"],
+                }
+                for c in cols
+            }
+    return tables, columns
+
+
+def write_stats(
+    scripts_dir: Path, tables: dict[str, dict[str, Any]], columns: dict[str, dict[str, Any]], prune: bool = False
+) -> Path:
+    """Merge into _stats/_tables.json (keys sorted; tables not in `tables` are kept,
+    unless prune=True, which drops them and their column files — use for a full run)
+    and write one _stats/<schema.table>.json per table for column stats."""
+    path = stats_dir(scripts_dir) / TABLES_FILE
+    existing: dict[str, Any] = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
+    if prune:
+        for gone in set(existing) - set(tables):
+            column_stats_path(scripts_dir, gone).unlink(missing_ok=True)
+        existing = {}
+    existing.update(tables)
+    _write_json(path, existing)
+    for qn, cols in columns.items():
+        _write_json(column_stats_path(scripts_dir, qn), cols)
+    return path
+
+
+def read_stats(scripts_dir: Path, names: list[str], columns: bool = True) -> dict[str, Any]:
+    """Stats for the given "schema.table" names (all tables if empty), read
+    from the JSON files. Unknown tables raise KeyError."""
+    path = stats_dir(scripts_dir) / TABLES_FILE
+    if not path.exists():
+        raise FileNotFoundError(f"{path} not found — run `pgdb update-stats` first")
+    all_tables: dict[str, Any] = json.loads(path.read_text(encoding="utf-8"))
+    wanted = names or sorted(all_tables)
+    missing = [n for n in wanted if n not in all_tables]
+    if missing:
+        raise KeyError(", ".join(missing))
+    result: dict[str, Any] = {}
+    for qn in wanted:
+        entry = dict(all_tables[qn])
+        cpath = column_stats_path(scripts_dir, qn)
+        if columns and cpath.exists():
+            entry["columns"] = json.loads(cpath.read_text(encoding="utf-8"))
+        result[qn] = entry
+    return result

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.9.0"
+version = "0.10.0"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"

skills/pgdevkit/SKILL.md

@@ -266,6 +266,37 @@ It diffs the live schema against `database/`, reverse-engineers DDL for anything
 
 ---
 
+## Table & column stats — `pgdb update-stats` / `pgdb get-stats`
+
+Stats live next to the schema files, in `database/_stats/`, so they can be committed and read without a DB connection:
+
+```
+database/_stats/
+├── _tables.json             # one entry per table, keyed by "schema.table", keys sorted
+└── public.users.json        # column stats for one table, keyed by column name
+```
+
+```json
+// _tables.json
+{ "public.users": { "row_count": 1200, "row_count_exact": false,
+                    "table_bytes": 98304, "index_bytes": 32768, "total_bytes": 131072 } }
+// public.users.json
+{ "email": { "data_type": "text", "null_fraction": 0.0, "n_distinct": -1, "avg_width": 24 } }
+```
+
+`row_count` is the planner estimate (`null` if the table was never analyzed) unless `--exact` was used. Column stats come from `pg_stats` (`n_distinct` < 0 means a fraction of the row count, as in Postgres) and are `null` until the table is analyzed — pass `--analyze`.
+
+```bash
+pgdb update-stats database/ --url postgresql://... --analyze            # all tables
+pgdb update-stats database/ --url postgresql://... --table public.users --exact   # partial update, keeps other entries
+pgdb get-stats database/ public.users public.orders                     # JSON to stdout, no DB needed
+pgdb get-stats database/ --no-columns                                   # all tables, table-level stats only
+```
+
+To read the stats from code or a script, just `json.load` `database/_stats/_tables.json` (and `database/_stats/<schema.table>.json` for columns).
+
+---
+
 ## Comparing scripts to a live database
 
 ```bash

tests/test_stats_cli.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import psycopg
+from typer.testing import CliRunner
+
+from pgdevkit.cli import app
+from pgdevkit.testdb import constants
+from pgdevkit.testdb.container import ensure_container
+from tests.testdb.conftest import RUN_SUFFIX, requires_podman
+
+runner = CliRunner()
+
+TEST_DB = f"pgdevkit_stats_cli_selftest_{RUN_SUFFIX}"
+
+
+def test_get_stats_reads_json_without_db(tmp_path: Path):
+    stats = tmp_path / "_stats"
+    stats.mkdir()
+    (stats / "_tables.json").write_text(json.dumps({"public.a": {"row_count": 3}, "public.b": {"row_count": 1}}))
+    (stats / "public.a.json").write_text(json.dumps({"id": {"data_type": "integer"}}))
+
+    r = runner.invoke(app, ["get-stats", str(tmp_path), "public.a"])
+    assert r.exit_code == 0, r.output
+    assert json.loads(r.stdout) == {"public.a": {"row_count": 3, "columns": {"id": {"data_type": "integer"}}}}
+
+    r = runner.invoke(app, ["get-stats", str(tmp_path), "--no-columns"])
+    assert list(json.loads(r.stdout)) == ["public.a", "public.b"]
+
+    assert runner.invoke(app, ["get-stats", str(tmp_path), "public.nope"]).exit_code == 2
+
+
+@requires_podman
+def test_update_stats_then_get_stats(tmp_path: Path):
+    ensure_container()
+    admin = constants.conninfo("postgres")
+    dsn = constants.conninfo(TEST_DB)
+    with psycopg.connect(admin, autocommit=True) as con:
+        con.execute(f'DROP DATABASE IF EXISTS "{TEST_DB}"')
+        con.execute(f'CREATE DATABASE "{TEST_DB}"')
+    try:
+        with psycopg.connect(dsn, autocommit=True) as con:
+            con.execute("CREATE TABLE public.zeta (id int PRIMARY KEY, note text)")
+            con.execute("CREATE TABLE public.alpha (id int PRIMARY KEY)")
+            con.execute("INSERT INTO public.zeta SELECT g, NULL FROM generate_series(1, 10) g")
+
+        r = runner.invoke(app, ["update-stats", str(tmp_path), "--url", dsn, "--exact", "--analyze"])
+        assert r.exit_code == 0, r.output
+        tables = json.loads((tmp_path / "_stats" / "_tables.json").read_text())
+        assert list(tables) == ["public.alpha", "public.zeta"]
+        assert tables["public.zeta"]["row_count"] == 10
+        cols = json.loads((tmp_path / "_stats" / "public.zeta.json").read_text())
+        assert cols["note"]["null_fraction"] == 1
+
+        # A partial update keeps other tables' entries.
+        r = runner.invoke(app, ["update-stats", str(tmp_path), "--url", dsn, "--table", "public.alpha"])
+        assert r.exit_code == 0, r.output
+        assert "public.zeta" in json.loads((tmp_path / "_stats" / "_tables.json").read_text())
+
+        assert runner.invoke(app, ["update-stats", str(tmp_path), "--url", dsn, "--table", "public.nope"]).exit_code == 2
+
+        # A full run prunes dropped tables.
+        with psycopg.connect(dsn, autocommit=True) as con:
+            con.execute("DROP TABLE public.alpha")
+        assert runner.invoke(app, ["update-stats", str(tmp_path), "--url", dsn]).exit_code == 0
+        assert list(json.loads((tmp_path / "_stats" / "_tables.json").read_text())) == ["public.zeta"]
+        assert not (tmp_path / "_stats" / "public.alpha.json").exists()
+
+        got = json.loads(runner.invoke(app, ["get-stats", str(tmp_path), "public.zeta"]).stdout)
+        assert got["public.zeta"]["columns"]["id"]["data_type"] == "integer"
+    finally:
+        with psycopg.connect(admin, autocommit=True) as con:
+            con.execute(f'DROP DATABASE IF EXISTS "{TEST_DB}"')