pgdevkit 0.1.0__py3-none-any.whl

pgdevkit/cli.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import psycopg
+import typer
+from rich.console import Console
+from rich.table import Table
+from rich import box
+
+from . import testdb
+from .connection import build_conninfo
+from .diff import DiffKind, compute_diff
+from .fetch_missing import SUBFOLDER, find_missing_objects, layer_folder_for, reconstruct_ddl
+from .introspect import introspect_db
+from .parser import parse_directory
+
+app = typer.Typer(name="pgdb", help="PostgreSQL database schema tools")
+console = Console()
+err_console = Console(stderr=True)
+
+testdb_app = typer.Typer(name="testdb", help="Manage the shared local Postgres test container")
+app.add_typer(testdb_app, name="testdb")
+
+
+@app.command()
+def compare(
+    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)"),
+    databricks_workspace_host: str | None = typer.Option(
+        None,
+        "--databricks-workspace-host",
+        help="Databricks workspace URL, e.g. https://adb-....azuredatabricks.net (required for Lakebase hosts)",
+    ),
+    databricks_instance: str | None = typer.Option(
+        None, "--databricks-instance", help="Lakebase instance name (required for Lakebase hosts)"
+    ),
+    report_extra_db: bool = typer.Option(False, "--report-extra-db", help="Report objects in DB but not in scripts"),
+    scripts_dir: Path = typer.Argument(..., help="Directory containing SQL scripts"),
+) -> None:
+    """Compare SQL scripts to a live PostgreSQL database and report differences."""
+    if not scripts_dir.is_dir():
+        err_console.print(f"[red]Error:[/red] {scripts_dir} is not a directory")
+        raise typer.Exit(2)
+
+    try:
+        conninfo = build_conninfo(
+            url,
+            entra_user,
+            databricks_workspace_host=databricks_workspace_host,
+            databricks_instance=databricks_instance,
+        )
+    except ValueError as e:
+        err_console.print(f"[red]Error:[/red] {e}")
+        raise typer.Exit(2)
+
+    with console.status("Parsing SQL scripts..."):
+        scripts_schema = parse_directory(scripts_dir)
+
+    with console.status("Introspecting database..."):
+        db_schema = introspect_db(conninfo)
+
+    diffs = compute_diff(scripts_schema, db_schema, report_extra_db=report_extra_db)
+
+    if not diffs:
+        console.print("[green]No differences found.[/green]")
+        return
+
+    table = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+    table.add_column("Kind", style="cyan", min_width=20)
+    table.add_column("Type", style="magenta", min_width=15)
+    table.add_column("Object")
+    table.add_column("Detail", style="dim")
+
+    kind_style = {
+        DiffKind.MISSING_IN_DB: "[yellow]missing in DB[/yellow]",
+        DiffKind.MISSING_IN_SCRIPTS: "[blue]missing in scripts[/blue]",
+        DiffKind.MISMATCH: "[red]mismatch[/red]",
+    }
+    for d in diffs:
+        table.add_row(kind_style[d.kind], d.object_type, d.object_name, d.detail)
+
+    console.print(table)
+    console.print(f"\n[bold red]{len(diffs)} difference(s) found.[/bold red]")
+    raise typer.Exit(1)
+
+
+@app.command("fetch-missing")
+def fetch_missing(
+    scripts_dir: Path = typer.Argument(..., help="The database/ folder to compare against and write into"),
+    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)"),
+    write: bool = typer.Option(False, "--write", help="Write the reconstructed .sql files (default: dry run)"),
+    only: list[str] = typer.Option([], "--only", help="Only fetch schema.name (repeatable); default is everything"),
+) -> None:
+    """Find tables/views/functions that exist in the database but aren't
+    tracked under scripts_dir, and reverse-engineer their DDL into new files."""
+    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("Comparing database/ against the live schema..."):
+        missing = find_missing_objects(scripts_dir, conninfo)
+
+    if only:
+        wanted = set(only)
+        missing = [m for m in missing if m.qualified_name in wanted]
+
+    if not missing:
+        console.print("[green]No missing objects.[/green]")
+        return
+
+    table = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+    table.add_column("Type", style="magenta")
+    table.add_column("Object")
+    table.add_column("Destination", style="dim")
+    for m in missing:
+        dest = layer_folder_for(scripts_dir, m.schema) / SUBFOLDER[m.object_type] / f"{m.name}.sql"
+        table.add_row(m.object_type, m.qualified_name, str(dest))
+    console.print(table)
+
+    if not write:
+        console.print("\n[yellow]Dry run[/yellow] — pass --write to create these files.")
+        return
+
+    written = 0
+    with psycopg.connect(conninfo) as conn:
+        for m in missing:
+            dest_dir = layer_folder_for(scripts_dir, m.schema) / SUBFOLDER[m.object_type]
+            dest = dest_dir / f"{m.name}.sql"
+            if dest.exists():
+                console.print(f"  [yellow]SKIP[/yellow] {dest} (already exists)")
+                continue
+            try:
+                ddl = reconstruct_ddl(conn, m)
+            except Exception as e:  # noqa: BLE001
+                err_console.print(f"[red]Error[/red] reconstructing {m.qualified_name}: {e}")
+                continue
+            dest_dir.mkdir(parents=True, exist_ok=True)
+            dest.write_text(ddl, encoding="utf-8")
+            console.print(f"  [green]WROTE[/green] {dest}")
+            written += 1
+
+    console.print(f"\nWrote {written} file(s).")
+
+
+@testdb_app.command("up")
+def testdb_up() -> None:
+    """Ensure the container is running, the workspace DB exists, and schema is applied."""
+    testdb.ensure_testdb()
+    info = testdb.status()
+    console.print(f"[green]Test DB ready:[/green] {info['database']} ({info['dsn']})")
+
+
+@testdb_app.command("reset")
+def testdb_reset() -> None:
+    """Drop and recreate only this workspace's database, then reapply schema + seed data."""
+    testdb.reset_testdb()
+    info = testdb.status()
+    console.print(f"[green]Test DB reset:[/green] {info['database']}")
+
+
+@testdb_app.command("run-sql")
+def testdb_run_sql(
+    file: Path | None = typer.Argument(None, help="Path to a .sql file"),
+    sql: str | None = typer.Option(None, "--sql", help="Inline SQL string"),
+    results: bool = typer.Option(False, "--results", help="Print query results as a table"),
+) -> None:
+    """Run SQL against this workspace's database."""
+    if (file is None) == (sql is None):
+        err_console.print("[red]Error:[/red] pass exactly one of FILE or --sql")
+        raise typer.Exit(2)
+    statement = file.read_text(encoding="utf-8") if file else sql
+    assert statement is not None
+    rows = testdb.run_sql(statement)
+
+    if rows is None:
+        console.print("OK")
+        return
+    if not results:
+        console.print(f"OK — {len(rows)} row(s)")
+        return
+    if not rows:
+        console.print("(0 row(s))")
+        return
+    table = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+    for col in rows[0]:
+        table.add_column(col)
+    for row in rows:
+        table.add_row(*(str(v) for v in row.values()))
+    console.print(table)
+    console.print(f"({len(rows)} row(s))")
+
+
+@testdb_app.command("status")
+def testdb_status() -> None:
+    """Show container state, this workspace's database name, and DSN."""
+    for key, value in testdb.status().items():
+        console.print(f"{key}: {value}")
+
+
+@testdb_app.command("shell")
+def testdb_shell() -> None:
+    """Drop into psql against this workspace's database."""
+    os.execvp("psql", ["psql", testdb.dsn_for()])
+
+
+@testdb_app.command("clean")
+def testdb_clean(
+    all: bool = typer.Option(False, "--all", help="Drop every database belonging to this project"),
+) -> None:
+    """Drop this workspace's database (or every database of this project with --all)."""
+    testdb.clean_testdb(all=all)
+    console.print("[green]Cleaned.[/green]")

pgdevkit/connection.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Literal
+from urllib.parse import quote, urlparse, urlunparse
+
+_LAKEBASE_HOST_SUFFIXES = (
+    ".database.azuredatabricks.net",
+    ".database.cloud.databricks.com",
+)
+
+
+def detect_provider(host: str) -> Literal["azure_postgres", "databricks_lakebase"]:
+    """Classify a Postgres hostname: Databricks Lakebase (needs credential
+    exchange) or the default Azure Postgres AAD token flow."""
+    if any(host.endswith(suffix) for suffix in _LAKEBASE_HOST_SUFFIXES):
+        return "databricks_lakebase"
+    return "azure_postgres"
+
+
+def get_azure_postgres_password() -> str:
+    try:
+        from azure.identity import DefaultAzureCredential
+    except ImportError:
+        raise ImportError("Install azure-identity extra: pip install pgdevkit[azure]")
+    cred = DefaultAzureCredential()
+    token = cred.get_token("https://ossrdbms-aad.database.windows.net/.default")
+    return token.token
+
+
+def build_conninfo(
+    url: str,
+    entra_user: str | None = None,
+    *,
+    databricks_workspace_host: str | None = None,
+    databricks_instance: str | None = None,
+) -> str:
+    if entra_user is None:
+        return url
+
+    parsed = urlparse(url)
+    host = parsed.hostname or ""
+    port = f":{parsed.port}" if parsed.port else ""
+
+    if detect_provider(host) == "databricks_lakebase":
+        if not databricks_workspace_host or not databricks_instance:
+            raise ValueError(
+                "Lakebase host detected — pass --databricks-workspace-host and --databricks-instance"
+            )
+        from .lakebase import get_lakebase_password
+
+        password = get_lakebase_password(databricks_workspace_host, databricks_instance)
+    else:
+        password = get_azure_postgres_password()
+
+    netloc = f"{quote(entra_user, safe='')}:{quote(password, safe='')}@{host}{port}"
+    return urlunparse(parsed._replace(netloc=netloc))

pgdevkit/db/__init__.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from .connection import PgPool
+from .crud import (
+    pg_delete,
+    pg_delete_dict,
+    pg_insert,
+    pg_insert_many,
+    pg_retrieve,
+    pg_retrieve_many,
+    pg_update,
+    pg_update_dict,
+    pg_upsert,
+    pg_upsert_dict,
+    pg_upsert_many,
+    pg_upsert_many_dict,
+)
+from .loader import SqlLoader
+from .model import PostgresTableModel
+
+__all__ = [
+    "PgPool",
+    "PostgresTableModel",
+    "SqlLoader",
+    "pg_delete",
+    "pg_delete_dict",
+    "pg_insert",
+    "pg_insert_many",
+    "pg_retrieve",
+    "pg_retrieve_many",
+    "pg_update",
+    "pg_update_dict",
+    "pg_upsert",
+    "pg_upsert_dict",
+    "pg_upsert_many",
+    "pg_upsert_many_dict",
+]

pgdevkit/db/connection.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import asyncio
+import os
+
+from psycopg_pool import AsyncConnectionPool
+
+from ..connection import detect_provider, get_azure_postgres_password
+from ..lakebase import get_lakebase_password
+
+
+class PgPool:
+    """A connection pool keyed off `{env_prefix}HOST/PORT/DB/USER/PASSWORD`
+    environment variables (e.g. env_prefix="MDM_POSTGRES_").
+
+    Pass `entra_user` to authenticate via Entra ID instead of a static
+    password. The Postgres host is inspected to pick the Azure Postgres AAD
+    token flow or Databricks Lakebase credential exchange; for the latter
+    also set `{env_prefix}DATABRICKS_WORKSPACE_HOST` and
+    `{env_prefix}DATABRICKS_INSTANCE`.
+    """
+
+    def __init__(
+        self,
+        env_prefix: str = "POSTGRES_",
+        max_size: int = 40,
+        *,
+        entra_user: str | None = None,
+    ) -> None:
+        self._env_prefix = env_prefix
+        self._max_size = max_size
+        self._entra_user = entra_user
+        self._pool: AsyncConnectionPool | None = None
+
+    async def _dsn(self) -> str:
+        p = self._env_prefix
+        host = os.environ[p + "HOST"]
+        port = os.environ[p + "PORT"]
+        dbname = os.environ[p + "DB"]
+
+        if self._entra_user is None:
+            user = os.environ[p + "USER"]
+            password = os.environ[p + "PASSWORD"]
+        else:
+            user = self._entra_user
+            if detect_provider(host) == "databricks_lakebase":
+                workspace_host = os.environ[p + "DATABRICKS_WORKSPACE_HOST"]
+                instance_name = os.environ[p + "DATABRICKS_INSTANCE"]
+                password = await asyncio.to_thread(get_lakebase_password, workspace_host, instance_name)
+            else:
+                password = await asyncio.to_thread(get_azure_postgres_password)
+
+        return f"host={host} port={port} dbname={dbname} user={user} password={password}"
+
+    async def open(self) -> None:
+        if self._pool is None:
+            self._pool = AsyncConnectionPool(
+                conninfo=self._dsn,
+                open=False,
+                max_size=self._max_size,
+                check=AsyncConnectionPool.check_connection,
+            )
+        if not self._pool._opened:
+            await self._pool.open()
+
+    async def close(self) -> None:
+        if self._pool is not None:
+            await self._pool.close()
+
+    def connection(self):
+        """Return an async connection context manager from the pool."""
+        if self._pool is None:
+            raise RuntimeError("Call open() first (e.g. in app startup).")
+        return self._pool.connection()

pgdevkit/db/crud.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Mapping, Optional, Sequence, Type, TypeVar
+
+from psycopg.connection_async import AsyncConnection
+from psycopg.rows import dict_row
+from psycopg.sql import SQL, Identifier, Placeholder
+
+from .model import PostgresTableModel
+
+T = TypeVar("T", bound=PostgresTableModel)
+
+
+async def pg_retrieve(con: AsyncConnection, data_type: Type[T], pks: dict) -> T | None:
+    """Fetch a single row by primary key(s)."""
+    async with con.cursor(row_factory=dict_row) as cur:
+        schema, table = data_type.get_table_name()
+        query = SQL("SELECT * FROM {tbl} WHERE {where}").format(
+            tbl=Identifier(schema, table),
+            where=SQL(" AND ").join(SQL("{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in pks),
+        )
+        await cur.execute(query, pks)
+        row = await cur.fetchone()
+    return data_type(**row) if row else None
+
+
+async def pg_retrieve_many(
+    con: AsyncConnection,
+    data_type: Type[T],
+    filters: dict,
+    *,
+    from_dict: Optional[Callable[[Mapping], T]] = None,
+) -> Sequence[T]:
+    """Fetch multiple rows matching all filter key=value pairs."""
+    async with con.cursor(row_factory=dict_row) as cur:
+        schema, table = data_type.get_table_name()
+        if filters:
+            query = SQL("SELECT * FROM {tbl} WHERE {where}").format(
+                tbl=Identifier(schema, table),
+                where=SQL(" AND ").join(
+                    SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in filters
+                ),
+            )
+        else:
+            query = SQL("SELECT * FROM {tbl}").format(tbl=Identifier(schema, table))
+        await cur.execute(query, filters)
+        rows = await cur.fetchall()
+    fn = from_dict or (lambda d: data_type(**d))
+    return [fn(r) for r in rows]
+
+
+async def pg_insert(con: AsyncConnection, table_name: tuple[str, str], data: dict) -> dict[str, Any]:
+    """Insert one row and return the full row (RETURNING *)."""
+    query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING *").format(
+        tbl=Identifier(*table_name),
+        cols=SQL(", ").join(Identifier(k) for k in data),
+        vals=SQL(", ").join(Placeholder(k) for k in data),
+    )
+    async with con.cursor(row_factory=dict_row) as cur:
+        await cur.execute(query, data)
+        row = await cur.fetchone()
+    assert row is not None
+    return row
+
+
+async def pg_update_dict(
+    con: AsyncConnection,
+    table_name: tuple[str, str],
+    data: dict,
+    primary_keys: Sequence[str],
+) -> Any | None:
+    """Update a row identified by primary_keys. Returns the raw row tuple."""
+    set_parts = [
+        SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in data if k not in primary_keys
+    ]
+    where_parts = [SQL("{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys]
+    query = SQL("UPDATE {tbl} SET {sets} WHERE {where} RETURNING *").format(
+        tbl=Identifier(*table_name),
+        sets=SQL(", ").join(set_parts),
+        where=SQL(" AND ").join(where_parts),
+    )
+    async with con.cursor() as cur:
+        await cur.execute(query, data)
+        return await cur.fetchone()
+
+
+async def pg_update(con: AsyncConnection, data: T, data_type: type[T]) -> Any | None:
+    """Update a typed model instance."""
+    return await pg_update_dict(con, data_type.get_table_name(), data.model_dump(), data_type.get_primary_key())
+
+
+async def pg_upsert_dict(
+    con: AsyncConnection,
+    table_name: tuple[str, str],
+    data: dict,
+    primary_keys: Sequence[str],
+) -> dict:
+    """INSERT ... ON CONFLICT ... DO UPDATE, returns the row as a dict."""
+    fields = list(data)
+    updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields]
+    query = SQL(
+        "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates} RETURNING *"
+    ).format(
+        tbl=Identifier(*table_name),
+        cols=SQL(", ").join(Identifier(k) for k in fields),
+        vals=SQL(", ").join(Placeholder(k) for k in fields),
+        pks=SQL(", ").join(Identifier(pk) for pk in primary_keys),
+        updates=SQL(", ").join(updates),
+    )
+    async with con.cursor(row_factory=dict_row) as cur:
+        await cur.execute(query, data)
+        row = await cur.fetchone()
+    assert row is not None
+    return row
+
+
+async def pg_upsert(con: AsyncConnection, data: T, data_type: type[T]) -> dict:
+    """Upsert a typed model instance."""
+    return await pg_upsert_dict(con, data_type.get_table_name(), data.model_dump(), data_type.get_primary_key())
+
+
+async def pg_upsert_many_dict(
+    con: AsyncConnection,
+    table_name: tuple[str, str],
+    data: Sequence[dict],
+    primary_keys: Sequence[str],
+) -> None:
+    """Batch upsert — one round-trip via executemany."""
+    if not data:
+        return
+    fields = list(data[0])
+    updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields if k not in primary_keys]
+    query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates}").format(
+        tbl=Identifier(*table_name),
+        cols=SQL(", ").join(Identifier(k) for k in fields),
+        vals=SQL(", ").join(Placeholder(k) for k in fields),
+        pks=SQL(", ").join(Identifier(pk) for pk in primary_keys),
+        updates=SQL(", ").join(updates),
+    )
+    async with con.cursor() as cur:
+        await cur.executemany(query, data)
+
+
+async def pg_upsert_many(con: AsyncConnection, data: Sequence[T], data_type: type[T]) -> None:
+    await pg_upsert_many_dict(
+        con, data_type.get_table_name(), [d.model_dump() for d in data], data_type.get_primary_key()
+    )
+
+
+async def pg_insert_many(
+    con: AsyncConnection,
+    table_name: tuple[str, str],
+    data: Sequence[dict],
+) -> None:
+    """Batch insert — no RETURNING, one round-trip via executemany."""
+    if not data:
+        return
+    fields = list(data[0])
+    query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals})").format(
+        tbl=Identifier(*table_name),
+        cols=SQL(", ").join(Identifier(k) for k in fields),
+        vals=SQL(", ").join(Placeholder(k) for k in fields),
+    )
+    async with con.cursor() as cur:
+        await cur.executemany(query, data)
+
+
+async def pg_delete_dict(con: AsyncConnection, table_name: tuple[str, str], data: dict) -> dict | None:
+    """Delete by arbitrary key dict, returns the deleted row."""
+    where_parts = [SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in data]
+    query = SQL("DELETE FROM {tbl} WHERE {where} RETURNING *").format(
+        tbl=Identifier(*table_name),
+        where=SQL(" AND ").join(where_parts),
+    )
+    async with con.cursor(row_factory=dict_row) as cur:
+        await cur.execute(query, data)
+        return await cur.fetchone()
+
+
+async def pg_delete(con: AsyncConnection, data: T, data_type: type[T]) -> T | None:
+    """Delete a typed model instance by its primary key(s)."""
+    pk_dict = {pk: getattr(data, pk) for pk in data_type.get_primary_key()}
+    row = await pg_delete_dict(con, data_type.get_table_name(), pk_dict)
+    return data_type.model_validate(row) if row else None

pgdevkit/db/loader.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from pathlib import Path
+from typing import LiteralString, cast
+
+
+class SqlLoader:
+    """Loads and caches SQL text from `{root}/<topic>/<name>.sql`."""
+
+    def __init__(self, root: Path) -> None:
+        self._root = root
+        self._load = lru_cache(maxsize=None)(self._read)
+
+    def _read(self, topic: str, name: str) -> LiteralString:
+        return cast(LiteralString, (self._root / topic / f"{name}.sql").read_text(encoding="utf-8"))
+
+    def load_sql(self, topic: str, name: str) -> LiteralString:
+        return self._load(topic, name)

pgdevkit/db/model.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Sequence
+
+from pydantic import BaseModel
+
+
+class PostgresTableModel(BaseModel, ABC):
+    """Base class for models that map 1:1 to a database table/row.
+
+    Models representing partial results (joins, aggregations, projections)
+    should extend `pydantic.BaseModel` directly instead."""
+
+    @staticmethod
+    @abstractmethod
+    def get_table_name() -> tuple[str, str]:
+        """Return (schema, table), e.g. ('public', 'users')."""
+
+    @staticmethod
+    @abstractmethod
+    def get_primary_key() -> Sequence[str]:
+        """Return the primary key column name(s)."""