pgdevkit 0.8.0__tar.gz → 0.10.0__tar.gz

{pgdevkit-0.8.0 → pgdevkit-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pgdevkit
-Version: 0.8.0
+Version: 0.10.0
 Summary: A helper for developing with Postgres
 Requires-Python: >=3.14
 Requires-Dist: docker>=7.1.0

{pgdevkit-0.8.0 → pgdevkit-0.10.0}/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.table import Table
 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-0.8.0 → pgdevkit-0.10.0}/pgdevkit/migrate.py

@@ -303,6 +303,14 @@ def record_applied(conninfo: str, tracking_table: str, filename: str) -> bool:
             return False
 
 
+def created_table_names(sql: str) -> list[str]:
+    """Table names any CREATE TABLE statement in this raw SQL script targets. Public
+    wrapper around the same detection `apply_migration` uses internally, for callers that
+    run a script directly (e.g. via `execute_sql_script`) instead of through a tracked
+    migration file, and still want to know what tables -- if any -- it created."""
+    return _created_table_names(_split_sql(sql))
+
+
 def verify_created_tables(conninfo: str, stmts: list[str]) -> list[str]:
     """Table names from this migration's CREATE TABLE statements that do NOT exist in the
     database. Empty means everything landed."""
@@ -318,6 +326,39 @@ def verify_created_tables(conninfo: str, stmts: list[str]) -> list[str]:
     return missing
 
 
+def missing_privileges(conninfo: str, role: str, tables: list[str], privilege: str = "select") -> list[str]:
+    """Table names from `tables` that `role` cannot currently exercise `privilege` on
+    (checked via Postgres's own `has_table_privilege`). Empty means the role can access
+    all of them. For guarding against the classic "migration creates a table, nobody
+    grants it to the app's runtime role" gap: check the tables a migration just created
+    against the role that will actually query them at runtime."""
+    if not tables:
+        return []
+    missing = []
+    with psycopg.connect(conninfo) as con:
+        for tbl in tables:
+            row = con.execute("select has_table_privilege(%s, %s, %s)", (role, tbl, privilege)).fetchone()
+            if not (row and row[0]):
+                missing.append(tbl)
+    return missing
+
+
+def _execute_stmts(conninfo: str, stmts: list[str]) -> None:
+    with psycopg.connect(conninfo) as con:
+        for stmt in stmts:
+            con.execute(cast(LiteralString, stmt))
+        con.commit()
+
+
+def execute_sql_script(conninfo: str, sql: str) -> None:
+    """Run a raw SQL script as one committed transaction, split into statements the same
+    statement-boundary-safe way apply_migration is (dollar-quoted blocks, string literals,
+    and line comments never get split mid-statement). Unlike apply_migration, this does
+    no tracking-table bookkeeping and isn't forward-only -- for scripts meant to re-run
+    every time, like an idempotent `GRANT ... ON ALL TABLES IN SCHEMA` privilege sync."""
+    _execute_stmts(conninfo, _split_sql(sql))
+
+
 @dataclass
 class ApplyResult:
     filename: str
@@ -337,11 +378,7 @@ def apply_migration(
     stmts = _split_sql(sql)
 
     if not already_done:
-        # DDL in its own committed transaction.
-        with psycopg.connect(conninfo) as con:
-            for stmt in stmts:
-                con.execute(cast(LiteralString, stmt))
-            con.commit()
+        _execute_stmts(conninfo, stmts)
 
     # Tracking insert is a separate connection/transaction so a missing tracking table
     # never rolls back the DDL that was just applied.

pgdevkit-0.10.0/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

{pgdevkit-0.8.0 → pgdevkit-0.10.0}/pyproject.toml

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

{pgdevkit-0.8.0 → pgdevkit-0.10.0}/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

pgdevkit-0.10.0/tests/test_migrate_grants.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import psycopg
+import pytest
+
+from pgdevkit.migrate import created_table_names, execute_sql_script, missing_privileges
+
+ROLE = "pgdevkit_test_grants_role"
+
+
+@pytest.fixture
+def app_role(clean_db: str):
+    """A throwaway, no-login role to check privileges against -- dropped afterwards since
+    roles are cluster-wide, not scoped to the (per-test) database like clean_db's schemas."""
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute(f"DROP ROLE IF EXISTS {ROLE}")
+        con.execute(f"CREATE ROLE {ROLE} NOLOGIN")
+    yield ROLE
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        # DROP OWNED BY revokes any grants left on tables the role can still see --
+        # otherwise DROP ROLE fails with "cannot be dropped because some objects depend
+        # on it" for as long as the granted-on table (e.g. widgets) still exists.
+        con.execute(f"DROP OWNED BY {ROLE}")
+        con.execute(f"DROP ROLE IF EXISTS {ROLE}")
+
+
+def test_created_table_names_finds_tables_in_a_raw_script():
+    sql = "GRANT SELECT ON foo TO some_role; CREATE TABLE public.widgets (id int primary key);"
+    assert created_table_names(sql) == ["public.widgets"]
+
+
+def test_created_table_names_empty_for_a_script_with_no_create_table():
+    assert created_table_names("GRANT SELECT ON ALL TABLES IN SCHEMA public TO some_role;") == []
+
+
+def test_execute_sql_script_runs_multiple_statements_in_one_transaction(clean_db: str):
+    execute_sql_script(
+        clean_db,
+        """
+        -- a leading comment shouldn't confuse statement splitting
+        CREATE TABLE public.widgets (id int primary key);
+        INSERT INTO public.widgets VALUES (1);
+        """,
+    )
+    with psycopg.connect(clean_db) as con:
+        row = con.execute("select count(*) from public.widgets").fetchone()
+        assert row == (1,)
+
+
+def test_execute_sql_script_rolls_back_all_statements_on_any_failure(clean_db: str):
+    with pytest.raises(psycopg.Error):
+        execute_sql_script(
+            clean_db,
+            "CREATE TABLE public.widgets (id int primary key); NOT VALID SQL HERE;",
+        )
+    with psycopg.connect(clean_db) as con:
+        row = con.execute("select to_regclass('public.widgets')").fetchone()
+        assert row is not None and row[0] is None
+
+
+def test_missing_privileges_reports_table_without_grant(clean_db: str, app_role: str):
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == ["public.widgets"]
+
+
+def test_missing_privileges_empty_once_granted(clean_db: str, app_role: str):
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+        con.execute(f"GRANT SELECT ON public.widgets TO {app_role}")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == []
+
+
+def test_missing_privileges_empty_for_no_tables(clean_db: str, app_role: str):
+    assert missing_privileges(clean_db, app_role, []) == []
+
+
+def test_execute_sql_script_grant_on_all_tables_covers_tables_created_after_the_fact(
+    clean_db: str, app_role: str
+):
+    """The actual use case: a re-run of `GRANT ... ON ALL TABLES IN SCHEMA` after a new
+    table appears picks it up automatically -- unlike ALTER DEFAULT PRIVILEGES, which only
+    covers tables created later by the exact role that ran the ALTER DEFAULT PRIVILEGES
+    statement itself."""
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == ["public.widgets"]
+
+    execute_sql_script(clean_db, f"GRANT SELECT ON ALL TABLES IN SCHEMA public TO {app_role}")
+
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == []

pgdevkit-0.10.0/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}"')

{pgdevkit-0.8.0 → pgdevkit-0.10.0}/uv.lock

@@ -313,7 +313,7 @@ wheels = [
 
 [[package]]
 name = "pgdevkit"
-version = "0.8.0"
+version = "0.10.0"
 source = { editable = "." }
 dependencies = [
     { name = "docker" },