pgdevkit 0.1.0__py3-none-any.whl

pgdevkit/introspect.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+from typing import Any
+
+import psycopg
+from psycopg.rows import dict_row
+
+from .models import (
+    ColumnDef, ConstraintDef, CompositeTypeDef, DatabaseSchema,
+    EnumDef, FunctionDef, IndexDef, TableDef, ViewDef,
+)
+
+
+def _q(conn: Any, sql: str, 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()  # type: ignore[return-value]
+
+
+def introspect_db(conninfo: str) -> DatabaseSchema:
+    with psycopg.connect(conninfo) as conn:
+        db = DatabaseSchema()
+        _load_schemas(conn, db)
+        _load_tables(conn, db)
+        _load_views(conn, db)
+        _load_functions(conn, db)
+        _load_enums(conn, db)
+        _load_composites(conn, db)
+        _load_indexes(conn, db)
+        return db
+
+
+def _load_schemas(conn: Any, db: DatabaseSchema) -> None:
+    rows = _q(conn, (
+        "SELECT schema_name FROM information_schema.schemata "
+        "WHERE schema_name NOT LIKE 'pg_%' AND schema_name != 'information_schema'"
+    ))
+    db.schemas = {r["schema_name"] for r in rows}
+
+
+def _load_tables(conn: Any, db: DatabaseSchema) -> None:
+    tables = _q(conn, """
+        SELECT n.nspname AS schema, c.relname AS name, c.relispartition AS is_partition
+        FROM pg_catalog.pg_class c
+        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+        WHERE c.relkind IN ('r', 'p')
+          AND n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+    """)
+
+    for row in tables:
+        tschema, tname = row["schema"], row["name"]
+        table = TableDef(schema=tschema, name=tname, is_partition=bool(row["is_partition"]))
+
+        for c in _q(conn, """
+            SELECT a.attname AS name,
+                   pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                   NOT a.attnotnull AS is_nullable,
+                   pg_get_expr(d.adbin, d.adrelid) AS col_default,
+                   a.attgenerated != '' AS is_generated
+            FROM pg_catalog.pg_attribute a
+            JOIN pg_catalog.pg_class c ON c.oid = a.attrelid
+            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+            LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
+            WHERE n.nspname = %(s)s AND c.relname = %(t)s
+              AND a.attnum > 0 AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """, {"s": tschema, "t": tname}):
+            table.columns.append(ColumnDef(
+                name=c["name"],
+                data_type=c["data_type"],
+                is_nullable=bool(c["is_nullable"]),
+                default=c["col_default"],
+                is_generated=bool(c["is_generated"]),
+            ))
+
+        kind_map = {"p": "PRIMARY KEY", "u": "UNIQUE", "f": "FOREIGN KEY", "c": "CHECK"}
+        for c in _q(conn, """
+            SELECT conname AS name, contype AS kind,
+                   pg_get_constraintdef(oid) AS definition
+            FROM pg_catalog.pg_constraint
+            WHERE conrelid = (
+                SELECT c.oid FROM pg_catalog.pg_class c
+                JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+                WHERE n.nspname = %(s)s AND c.relname = %(t)s
+            )
+        """, {"s": tschema, "t": tname}):
+            table.constraints.append(ConstraintDef(
+                name=c["name"],
+                kind=kind_map.get(c["kind"], c["kind"]),
+                definition=(c["definition"] or "").lower(),
+            ))
+
+        db.tables[table.qualified_name] = table
+
+
+def _load_views(conn: Any, db: DatabaseSchema) -> None:
+    for r in _q(conn, """
+        SELECT n.nspname AS schema, c.relname AS name,
+               pg_get_viewdef(c.oid, true) AS definition
+        FROM pg_catalog.pg_class c
+        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+        WHERE c.relkind = 'v'
+          AND n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+    """):
+        v = ViewDef(schema=r["schema"], name=r["name"], definition=(r["definition"] or "").lower())
+        db.views[v.qualified_name] = v
+
+
+def _load_functions(conn: Any, db: DatabaseSchema) -> None:
+    for r in _q(conn, """
+        SELECT n.nspname AS schema, p.proname AS name,
+               pg_get_function_arguments(p.oid) AS args,
+               pg_get_function_result(p.oid) AS return_type,
+               l.lanname AS language, p.prosrc AS body,
+               CASE WHEN p.prokind = 'p' THEN 'procedure' ELSE 'function' END AS kind
+        FROM pg_catalog.pg_proc p
+        JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace
+        JOIN pg_catalog.pg_language l ON l.oid = p.prolang
+        WHERE n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+    """):
+        raw_body = r["body"] or ""
+        lines = [line.strip() for line in raw_body.splitlines()]
+        body = "\n".join(line.lower() for line in lines if line)
+        func = FunctionDef(
+            schema=r["schema"], name=r["name"],
+            args=(r["args"] or "").lower(),
+            return_type=(r["return_type"] or "").lower(),
+            language=r["language"], body=body, kind=r["kind"],
+        )
+        db.functions[func.qualified_name] = func
+
+
+def _load_enums(conn: Any, db: DatabaseSchema) -> None:
+    for r in _q(conn, """
+        SELECT n.nspname AS schema, t.typname AS name,
+               array_agg(e.enumlabel ORDER BY e.enumsortorder) AS values
+        FROM pg_catalog.pg_type t
+        JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+        JOIN pg_catalog.pg_enum e ON e.enumtypid = t.oid
+        WHERE n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+        GROUP BY n.nspname, t.typname
+    """):
+        enum = EnumDef(schema=r["schema"], name=r["name"], values=list(r["values"]))
+        db.enums[enum.qualified_name] = enum
+
+
+def _load_composites(conn: Any, db: DatabaseSchema) -> None:
+    composites: dict[str, CompositeTypeDef] = {}
+    for r in _q(conn, """
+        SELECT n.nspname AS schema, t.typname AS name,
+               a.attname AS field_name,
+               pg_catalog.format_type(a.atttypid, a.atttypmod) AS field_type
+        FROM pg_catalog.pg_type t
+        JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+        JOIN pg_catalog.pg_class c ON c.oid = t.typrelid AND c.relkind = 'c'
+        JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum > 0 AND NOT a.attisdropped
+        WHERE n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+        ORDER BY n.nspname, t.typname, a.attnum
+    """):
+        key = f"{r['schema']}.{r['name']}"
+        if key not in composites:
+            composites[key] = CompositeTypeDef(schema=r["schema"], name=r["name"], fields=[])
+        composites[key].fields.append((r["field_name"], r["field_type"]))
+    db.composites = composites
+
+
+def _load_indexes(conn: Any, db: DatabaseSchema) -> None:
+    for r in _q(conn, """
+        SELECT n.nspname AS schema, t.relname AS table_name,
+               i.relname AS index_name,
+               pg_get_indexdef(ix.indexrelid) AS definition
+        FROM pg_catalog.pg_index ix
+        JOIN pg_catalog.pg_class i ON i.oid = ix.indexrelid
+        JOIN pg_catalog.pg_class t ON t.oid = ix.indrelid
+        JOIN pg_catalog.pg_namespace n ON n.oid = t.relnamespace
+        WHERE n.nspname NOT IN ('pg_catalog', 'information_schema')
+          AND n.nspname NOT LIKE 'pg_%'
+          AND NOT ix.indisprimary
+    """):
+        idx = IndexDef(
+            schema=r["schema"], table=r["table_name"],
+            name=r["index_name"],
+            definition=(r["definition"] or "").lower(),
+        )
+        db.indexes[idx.qualified_name] = idx

pgdevkit/lakebase.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import datetime
+import json
+import time
+import uuid
+from typing import TYPE_CHECKING
+from urllib import request as urllib_request
+from urllib.error import HTTPError
+
+if TYPE_CHECKING:
+    from azure.identity import DefaultAzureCredential
+
+_DATABRICKS_RESOURCE_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default"
+_REFRESH_MARGIN_SECONDS = 300
+
+_credential: DefaultAzureCredential | None = None
+_credential_cache: dict[tuple[str, str], tuple[str, float]] = {}
+_dns_cache: dict[tuple[str, str], str] = {}
+
+
+def _databricks_resource_token() -> str:
+    try:
+        from azure.identity import DefaultAzureCredential
+    except ImportError:
+        raise ImportError("Install azure-identity extra: pip install pgdevkit[azure]")
+    global _credential
+    if _credential is None:
+        _credential = DefaultAzureCredential()
+    return _credential.get_token(_DATABRICKS_RESOURCE_SCOPE).token
+
+
+def _request_json(url: str, token: str, *, body: dict | None = None) -> dict:
+    data = json.dumps(body).encode("utf-8") if body is not None else None
+    req = urllib_request.Request(
+        url,
+        data=data,
+        method="POST" if body is not None else "GET",
+        headers={
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+        },
+    )
+    try:
+        with urllib_request.urlopen(req, timeout=10) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except HTTPError as e:
+        detail = e.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"Databricks API request failed [{e.code}]: {detail}") from e
+
+
+def _parse_expiration(expiration_time: str) -> float:
+    return datetime.datetime.fromisoformat(expiration_time).timestamp()
+
+
+def get_lakebase_password(workspace_host: str, instance_name: str) -> str:
+    """Return a valid short-lived Postgres password for the given Lakebase
+    instance, fetching/refreshing the credential as needed."""
+    key = (workspace_host, instance_name)
+    cached = _credential_cache.get(key)
+    if cached is not None:
+        token, expires_at = cached
+        if time.time() < expires_at - _REFRESH_MARGIN_SECONDS:
+            return token
+
+    entra_token = _databricks_resource_token()
+    response = _request_json(
+        f"{workspace_host.rstrip('/')}/api/2.0/database/credentials",
+        entra_token,
+        body={"instance_names": [instance_name], "request_id": str(uuid.uuid4())},
+    )
+    token = response["token"]
+    expires_at = _parse_expiration(response["expiration_time"])
+    _credential_cache[key] = (token, expires_at)
+    return token
+
+
+def resolve_read_write_dns(workspace_host: str, instance_name: str) -> str:
+    """Resolve and cache the Postgres read/write DNS endpoint for a Lakebase instance."""
+    key = (workspace_host, instance_name)
+    if key in _dns_cache:
+        return _dns_cache[key]
+    entra_token = _databricks_resource_token()
+    response = _request_json(
+        f"{workspace_host.rstrip('/')}/api/2.0/database/instances/{instance_name}",
+        entra_token,
+    )
+    dns = response["read_write_dns"]
+    _dns_cache[key] = dns
+    return dns

pgdevkit/models.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ColumnDef:
+    name: str
+    data_type: str
+    is_nullable: bool
+    default: str | None
+    is_generated: bool = False
+
+
+@dataclass
+class ConstraintDef:
+    name: str | None
+    kind: str  # PRIMARY KEY, UNIQUE, FOREIGN KEY, CHECK
+    definition: str
+
+
+@dataclass
+class TableDef:
+    schema: str
+    name: str
+    columns: list[ColumnDef] = field(default_factory=list)
+    constraints: list[ConstraintDef] = field(default_factory=list)
+    is_partition: bool = False
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+
+@dataclass
+class ViewDef:
+    schema: str
+    name: str
+    definition: str
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+
+@dataclass
+class FunctionDef:
+    schema: str
+    name: str
+    args: str
+    return_type: str
+    language: str
+    body: str
+    kind: str = "function"
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+
+@dataclass
+class EnumDef:
+    schema: str
+    name: str
+    values: list[str]
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+
+@dataclass
+class CompositeTypeDef:
+    schema: str
+    name: str
+    fields: list[tuple[str, str]]
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+
+@dataclass
+class IndexDef:
+    schema: str
+    table: str
+    name: str
+    definition: str
+
+    @property
+    def qualified_name(self) -> str:
+        return self.name
+
+
+@dataclass
+class DatabaseSchema:
+    schemas: set[str] = field(default_factory=set)
+    tables: dict[str, TableDef] = field(default_factory=dict)
+    views: dict[str, ViewDef] = field(default_factory=dict)
+    functions: dict[str, FunctionDef] = field(default_factory=dict)
+    enums: dict[str, EnumDef] = field(default_factory=dict)
+    composites: dict[str, CompositeTypeDef] = field(default_factory=dict)
+    indexes: dict[str, IndexDef] = field(default_factory=dict)