diff --git a/pgdevkit/cli.py b/pgdevkit/cli.py index 1f2d1e2..a848677 100644 --- a/pgdevkit/cli.py +++ b/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 /_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/.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, diff --git a/pgdevkit/stats.py b/pgdevkit/stats.py new file mode 100644 index 0000000..b6d01ca --- /dev/null +++ b/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/.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 diff --git a/pyproject.toml b/pyproject.toml index 6fab0c5..54cb50c 100644 --- a/pyproject.toml +++ b/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" diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index f0866f1..f270c11 100644 --- a/skills/pgdevkit/SKILL.md +++ b/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/.json` for columns). + +--- + ## Comparing scripts to a live database ```bash diff --git a/tests/test_stats_cli.py b/tests/test_stats_cli.py new file mode 100644 index 0000000..204a6da --- /dev/null +++ b/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}"') diff --git a/uv.lock b/uv.lock index 644aad1..f07cef7 100644 --- a/uv.lock +++ b/uv.lock @@ -313,7 +313,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.9.0" +version = "0.10.0" source = { editable = "." } dependencies = [ { name = "docker" },