From a212a4f5eed91611e095245d9705307c7ec8e4f5 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Wed, 30 Sep 2026 13:49:43 +0200 Subject: [PATCH 1/5] WIP: table/column stats (#41) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01XCPhFUUdP3NVRQunZ8hbyb From 585cfbf463affcf3764de3f4d639f3f6104b9b92 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Wed, 30 Sep 2026 13:51:41 +0200 Subject: [PATCH 2/5] Add update-stats and get-stats commands, bump to 0.10.0 (#41) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01XCPhFUUdP3NVRQunZ8hbyb --- pgdevkit/cli.py | 43 ++++++++++++- pgdevkit/stats.py | 132 +++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- skills/pgdevkit/SKILL.md | 31 +++++++++ tests/test_stats_cli.py | 66 ++++++++++++++++++++ uv.lock | 2 +- 6 files changed, 273 insertions(+), 3 deletions(-) create mode 100644 pgdevkit/stats.py create mode 100644 tests/test_stats_cli.py diff --git a/pgdevkit/cli.py b/pgdevkit/cli.py index 1f2d1e2..5617d2a 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) + with console.status("Collecting stats..."): + tables, columns = stats.collect_stats(conninfo, set(table) or None, exact=exact, analyze=analyze) + unknown = set(table) - set(tables) + if unknown: + err_console.print(f"[red]Error:[/red] unknown table(s): {', '.join(sorted(unknown))}") + raise typer.Exit(2) + path = stats.write_stats(scripts_dir, tables, columns) + 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..bbbf3e1 --- /dev/null +++ b/pgdevkit/stats.py @@ -0,0 +1,132 @@ +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 +""" + + +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 _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, row_factory=dict_row) as conn: + rows = conn.execute(_TABLES_SQL).fetchall() + 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 = conn.execute(psycopg.sql.SQL("SELECT count(*) AS n FROM {}").format(ident)).fetchone()["n"] # type: ignore[index] + 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 = conn.execute( + _COLUMNS_SQL, + {"schema": r["schema"], "table": r["name"], "qualified": str(ident.as_string(conn))}, + ).fetchall() + 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]] +) -> Path: + """Merge into _stats/tables.json (keys sorted; tables not in `tables` are kept) + 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 {} + 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..8d01d89 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..75ba4f9 --- /dev/null +++ b/tests/test_stats_cli.py @@ -0,0 +1,66 @@ +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()) + + 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" }, From 818194858cfc3fac2e723a6b98aad15c055a80e7 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Wed, 30 Sep 2026 13:52:17 +0200 Subject: [PATCH 3/5] Rename stats table file to _tables.json so it sorts first (#41) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01XCPhFUUdP3NVRQunZ8hbyb --- pgdevkit/cli.py | 2 +- pgdevkit/stats.py | 4 ++-- skills/pgdevkit/SKILL.md | 6 +++--- tests/test_stats_cli.py | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pgdevkit/cli.py b/pgdevkit/cli.py index 5617d2a..89d9ebd 100644 --- a/pgdevkit/cli.py +++ b/pgdevkit/cli.py @@ -228,7 +228,7 @@ def update_stats( 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) + """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") diff --git a/pgdevkit/stats.py b/pgdevkit/stats.py index bbbf3e1..66df881 100644 --- a/pgdevkit/stats.py +++ b/pgdevkit/stats.py @@ -9,7 +9,7 @@ from psycopg.rows import dict_row STATS_DIRNAME = "_stats" -TABLES_FILE = "tables.json" +TABLES_FILE = "_tables.json" _TABLES_SQL = """ SELECT n.nspname AS schema, c.relname AS name, @@ -100,7 +100,7 @@ def collect_stats( def write_stats( scripts_dir: Path, tables: dict[str, dict[str, Any]], columns: dict[str, dict[str, Any]] ) -> Path: - """Merge into _stats/tables.json (keys sorted; tables not in `tables` are kept) + """Merge into _stats/_tables.json (keys sorted; tables not in `tables` are kept) 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 {} diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index 8d01d89..f270c11 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -272,12 +272,12 @@ Stats live next to the schema files, in `database/_stats/`, so they can be commi ``` database/_stats/ -├── tables.json # one entry per table, keyed by "schema.table", keys sorted +├── _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 +// _tables.json { "public.users": { "row_count": 1200, "row_count_exact": false, "table_bytes": 98304, "index_bytes": 32768, "total_bytes": 131072 } } // public.users.json @@ -293,7 +293,7 @@ pgdb get-stats database/ public.users public.orders # JSON t 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). +To read the stats from code or a script, just `json.load` `database/_stats/_tables.json` (and `database/_stats/.json` for columns). --- diff --git a/tests/test_stats_cli.py b/tests/test_stats_cli.py index 75ba4f9..1dc5eed 100644 --- a/tests/test_stats_cli.py +++ b/tests/test_stats_cli.py @@ -19,7 +19,7 @@ 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 / "_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"]) @@ -48,7 +48,7 @@ def test_update_stats_then_get_stats(tmp_path: Path): 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()) + 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()) @@ -57,7 +57,7 @@ def test_update_stats_then_get_stats(tmp_path: Path): # 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 "public.zeta" in json.loads((tmp_path / "_stats" / "_tables.json").read_text()) got = json.loads(runner.invoke(app, ["get-stats", str(tmp_path), "public.zeta"]).stdout) assert got["public.zeta"]["columns"]["id"]["data_type"] == "integer" From b41ee93cff56caf9ea0a128484d03ca9e7bcffce Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Wed, 30 Sep 2026 13:52:52 +0200 Subject: [PATCH 4/5] Stats: prune dropped tables on full run, validate --table up front, prefer inherited column stats (#41) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01XCPhFUUdP3NVRQunZ8hbyb --- pgdevkit/cli.py | 12 ++++++------ pgdevkit/stats.py | 15 ++++++++++++--- tests/test_stats_cli.py | 9 +++++++++ 3 files changed, 27 insertions(+), 9 deletions(-) diff --git a/pgdevkit/cli.py b/pgdevkit/cli.py index 89d9ebd..a848677 100644 --- a/pgdevkit/cli.py +++ b/pgdevkit/cli.py @@ -234,13 +234,13 @@ def update_stats( err_console.print(f"[red]Error:[/red] {scripts_dir} is not a directory") raise typer.Exit(2) conninfo = build_conninfo(url, entra_user) - with console.status("Collecting stats..."): - tables, columns = stats.collect_stats(conninfo, set(table) or None, exact=exact, analyze=analyze) - unknown = set(table) - set(tables) - if unknown: - err_console.print(f"[red]Error:[/red] unknown table(s): {', '.join(sorted(unknown))}") + 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) + path = stats.write_stats(scripts_dir, tables, columns, prune=not table) console.print(f"Updated stats for {len(tables)} table(s) in {path.parent}") diff --git a/pgdevkit/stats.py b/pgdevkit/stats.py index 66df881..5438470 100644 --- a/pgdevkit/stats.py +++ b/pgdevkit/stats.py @@ -32,7 +32,7 @@ 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 +ORDER BY a.attnum, s.inherited NULLS FIRST """ @@ -61,6 +61,10 @@ def collect_stats( columns: dict[str, dict[str, Any]] = {} with psycopg.connect(conninfo, autocommit=True, row_factory=dict_row) as conn: rows = conn.execute(_TABLES_SQL).fetchall() + 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: @@ -98,12 +102,17 @@ def collect_stats( def write_stats( - scripts_dir: Path, tables: dict[str, dict[str, Any]], columns: dict[str, dict[str, Any]] + 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) + """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(): diff --git a/tests/test_stats_cli.py b/tests/test_stats_cli.py index 1dc5eed..204a6da 100644 --- a/tests/test_stats_cli.py +++ b/tests/test_stats_cli.py @@ -59,6 +59,15 @@ def test_update_stats_then_get_stats(tmp_path: Path): 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: From 6ec2bcbc33701f1a7361b57e40182805025ccef6 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Wed, 30 Sep 2026 13:55:05 +0200 Subject: [PATCH 5/5] Fix ty errors in stats row handling (#41) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01XCPhFUUdP3NVRQunZ8hbyb --- pgdevkit/stats.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/pgdevkit/stats.py b/pgdevkit/stats.py index 5438470..b6d01ca 100644 --- a/pgdevkit/stats.py +++ b/pgdevkit/stats.py @@ -44,6 +44,12 @@ 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") @@ -59,8 +65,8 @@ def collect_stats( 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, row_factory=dict_row) as conn: - rows = conn.execute(_TABLES_SQL).fetchall() + 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: @@ -75,7 +81,7 @@ def collect_stats( if analyze: conn.execute(psycopg.sql.SQL("ANALYZE {}").format(ident)) if exact: - row_count = conn.execute(psycopg.sql.SQL("SELECT count(*) AS n FROM {}").format(ident)).fetchone()["n"] # type: ignore[index] + 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] = { @@ -85,10 +91,11 @@ def collect_stats( "index_bytes": r["index_bytes"], "total_bytes": r["total_bytes"], } - cols = conn.execute( + cols = _q( + conn, _COLUMNS_SQL, - {"schema": r["schema"], "table": r["name"], "qualified": str(ident.as_string(conn))}, - ).fetchall() + {"schema": r["schema"], "table": r["name"], "qualified": ident.as_string(conn)}, + ) columns[qn] = { c["name"]: { "data_type": c["data_type"],