pgdevkit 0.10.0__tar.gz → 0.11.0__tar.gz

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

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

{pgdevkit-0.10.0 → pgdevkit-0.11.0}/pgdevkit/cli.py

@@ -227,16 +227,22 @@ def update_stats(
     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"),
+    dialect: str = typer.Option("postgres", "--dialect", help="postgres (default) or mssql"),
 ) -> 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)
+    try:
+        dialect = get_backend(dialect).dialect.name
+    except ValueError as e:
+        err_console.print(f"[red]Error:[/red] {e}")
+        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)
+            tables, columns = stats.collect_stats(conninfo, set(table) or None, exact=exact, analyze=analyze, dialect=dialect)
     except KeyError as e:
         err_console.print(f"[red]Error:[/red] unknown table(s): {e.args[0]}")
         raise typer.Exit(2)

{pgdevkit-0.10.0 → pgdevkit-0.11.0}/pgdevkit/stats.py

@@ -8,6 +8,8 @@ import psycopg
 import psycopg.sql
 from psycopg.rows import dict_row
 
+from .dialect import resolve_dialect
+
 STATS_DIRNAME = "_stats"
 TABLES_FILE = "_tables.json"
 
@@ -55,14 +57,132 @@ def _write_json(path: Path, data: Any) -> None:
     path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
 
 
+_MSSQL_TABLES_SQL = """
+SELECT s.name AS [schema], t.name AS name,
+       SUM(CASE WHEN ps.index_id IN (0, 1) THEN ps.row_count ELSE 0 END) AS estimated_rows,
+       SUM(CASE WHEN ps.index_id IN (0, 1) THEN ps.used_page_count ELSE 0 END) * 8192 AS table_bytes,
+       SUM(CASE WHEN ps.index_id > 1 THEN ps.used_page_count ELSE 0 END) * 8192 AS index_bytes,
+       SUM(ps.used_page_count) * 8192 AS total_bytes
+FROM sys.tables t
+JOIN sys.schemas s ON s.schema_id = t.schema_id
+LEFT JOIN sys.dm_db_partition_stats ps ON ps.object_id = t.object_id
+GROUP BY s.name, t.name
+ORDER BY 1, 2
+"""
+
+_MSSQL_COLUMNS_SQL = """
+SELECT c.name AS name, ty.name AS base_type, c.max_length AS max_length,
+       c.precision AS precision, c.scale AS scale
+FROM sys.columns c
+JOIN sys.types ty ON ty.user_type_id = c.user_type_id
+WHERE c.object_id = OBJECT_ID(?)
+ORDER BY c.column_id
+"""
+
+# Types that can't be COUNT(DISTINCT)ed / measured with DATALENGTH.
+_MSSQL_UNMEASURABLE = {"text", "ntext", "image", "xml", "geography", "geometry", "hierarchyid", "sql_variant", "json", "vector"}
+
+
+def _mssql_q(conn: Any, sql: str, params: tuple = ()) -> list[dict[str, Any]]:
+    cur = conn.cursor()
+    try:
+        cur.execute(sql, params)
+        names = [c[0] for c in cur.description]
+        return [dict(zip(names, row)) for row in cur.fetchall()]
+    finally:
+        cur.close()
+
+
+def _mssql_ident(name: str) -> str:
+    return "[" + name.replace("]", "]]") + "]"
+
+
+def _collect_mssql_stats(
+    conninfo: str, only: set[str] | None, exact: bool, analyze: bool
+) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]]]:
+    """MSSQL flavour of collect_stats. Row counts/sizes come from
+    sys.dm_db_partition_stats. SQL Server has no pg_stats equivalent, so column
+    null_fraction/n_distinct/avg_width are only filled in with exact=True
+    (computed by scanning each table); otherwise they are null.
+    analyze=True runs UPDATE STATISTICS."""
+    # Lazy: the mssql extra isn't required for postgres-only use.
+    import mssql_python
+
+    from .mssql_introspect import _format_type, _is_system_schema
+
+    tables: dict[str, dict[str, Any]] = {}
+    columns: dict[str, dict[str, Any]] = {}
+    conn = mssql_python.connect(conninfo, autocommit=True)
+    try:
+        rows = [r for r in _mssql_q(conn, _MSSQL_TABLES_SQL) if not _is_system_schema(r["schema"])]
+        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 = f"{_mssql_ident(r['schema'])}.{_mssql_ident(r['name'])}"
+            if analyze:
+                conn.execute(f"UPDATE STATISTICS {ident}")
+            col_defs = _mssql_q(conn, _MSSQL_COLUMNS_SQL, (ident,))
+            measured: dict[str, Any] = {}
+            if exact:
+                # One scan per table, so row and column counts share a snapshot.
+                parts = ["COUNT_BIG(*) AS n"]
+                for i, c in enumerate(col_defs):
+                    if c["base_type"].lower() not in _MSSQL_UNMEASURABLE:
+                        col = _mssql_ident(c["name"])
+                        parts.append(
+                            f"COUNT_BIG({col}) AS nn{i}, COUNT_BIG(DISTINCT {col}) AS nd{i}, "
+                            f"AVG(CAST(DATALENGTH({col}) AS float)) AS w{i}"
+                        )
+                measured = _mssql_q(conn, f"SELECT {', '.join(parts)} FROM {ident}")[0]
+                row_count = measured["n"]
+            else:
+                row_count = r["estimated_rows"]
+            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"],
+            }
+            col_stats: dict[str, Any] = {}
+            for i, c in enumerate(col_defs):
+                entry: dict[str, Any] = {
+                    "data_type": _format_type(c["base_type"], c["max_length"], c["precision"], c["scale"]),
+                    "null_fraction": None,
+                    "n_distinct": None,
+                    "avg_width": None,
+                }
+                if row_count and f"nn{i}" in measured:
+                    entry["null_fraction"] = 1 - measured[f"nn{i}"] / row_count
+                    entry["n_distinct"] = measured[f"nd{i}"]
+                    w = measured[f"w{i}"]
+                    entry["avg_width"] = None if w is None else round(w)
+                col_stats[c["name"]] = entry
+            columns[qn] = col_stats
+    finally:
+        conn.close()
+    return tables, columns
+
+
 def collect_stats(
-    conninfo: str, only: set[str] | None = None, exact: bool = False, analyze: bool = False
+    conninfo: str,
+    only: set[str] | None = None,
+    exact: bool = False,
+    analyze: bool = False,
+    dialect: str = "postgres",
 ) -> 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."""
+    if resolve_dialect(dialect).name == "mssql":
+        return _collect_mssql_stats(conninfo, only, exact, analyze)
     tables: dict[str, dict[str, Any]] = {}
     columns: dict[str, dict[str, Any]] = {}
     with psycopg.connect(conninfo, autocommit=True) as conn:

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

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

{pgdevkit-0.10.0 → pgdevkit-0.11.0}/skills/pgdevkit/SKILL.md

@@ -293,6 +293,8 @@ pgdb get-stats database/ public.users public.orders                     # JSON t
 pgdb get-stats database/ --no-columns                                   # all tables, table-level stats only
 ```
 
+MSSQL: add `--dialect mssql` (needs the `mssql` extra). Row counts and sizes come from `sys.dm_db_partition_stats`; SQL Server has no `pg_stats`, so column `null_fraction`/`n_distinct`/`avg_width` are `null` unless `--exact` is given, which scans each table to compute them. `--analyze` runs `UPDATE STATISTICS`.
+
 To read the stats from code or a script, just `json.load` `database/_stats/_tables.json` (and `database/_stats/<schema.table>.json` for columns).
 
 ---

{pgdevkit-0.10.0 → pgdevkit-0.11.0}/tests/test_stats_cli.py

@@ -73,3 +73,4 @@ def test_update_stats_then_get_stats(tmp_path: Path):
     finally:
         with psycopg.connect(admin, autocommit=True) as con:
             con.execute(f'DROP DATABASE IF EXISTS "{TEST_DB}"')
+

pgdevkit-0.11.0/tests/test_stats_mssql_live.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import mssql_python
+import pytest
+from typer.testing import CliRunner
+
+from pgdevkit.cli import app
+from pgdevkit.testdb.api import clean_testdb, ensure_testdb, status
+from tests.testdb.conftest import _make_project, requires_mssql
+
+# Selected only by the dedicated `mssql-test` CI job; see test_compare_mssql_live.py.
+pytestmark = pytest.mark.mssql
+
+runner = CliRunner()
+
+
+def _exec(dsn: str, sql: str) -> None:
+    conn = mssql_python.connect(dsn, autocommit=True)
+    try:
+        conn.cursor().execute(sql)
+    finally:
+        conn.close()
+
+
+@requires_mssql
+def test_update_stats_mssql_live(tmp_path: Path):
+    project = _make_project(tmp_path, "mssqlstats", "main", engine="mssql")
+    out = tmp_path / "out"
+    out.mkdir()
+    try:
+        ensure_testdb(project)
+        dsn = status(project)["dsn"]
+        _exec(dsn, "INSERT INTO app.widget (id, name, tags) VALUES (2, N'gear', NULL), (3, N'cog', NULL)")
+
+        r = runner.invoke(app, ["update-stats", str(out), "--url", dsn, "--dialect", "mssql", "--exact", "--analyze"])
+        assert r.exit_code == 0, r.output
+        tables = json.loads((out / "_stats" / "_tables.json").read_text())
+        assert "app.widget" in tables and not any(k.startswith("sys.") for k in tables)
+        assert tables["app.widget"]["row_count"] == 3
+        assert tables["app.widget"]["row_count_exact"] is True
+        assert tables["app.widget"]["total_bytes"] > 0
+        cols = json.loads((out / "_stats" / "app.widget.json").read_text())
+        assert cols["name"]["data_type"] == "nvarchar(100)"
+        assert cols["id"]["n_distinct"] == 3 and cols["id"]["null_fraction"] == 0
+        assert cols["tags"]["null_fraction"] == 1
+
+        # Without --exact: engine-maintained row count, no column measurements.
+        r = runner.invoke(app, ["update-stats", str(out), "--url", dsn, "--dialect", "mssql", "--table", "app.widget"])
+        assert r.exit_code == 0, r.output
+        assert json.loads((out / "_stats" / "_tables.json").read_text())["app.widget"]["row_count"] == 3
+        assert json.loads((out / "_stats" / "app.widget.json").read_text())["id"]["n_distinct"] is None
+
+        bad = ["update-stats", str(out), "--url", dsn, "--dialect", "mssql", "--table", "app.nope"]
+        assert runner.invoke(app, bad).exit_code == 2
+    finally:
+        clean_testdb(project)

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

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