pgdevkit 0.1.0__tar.gz → 0.2.0__tar.gz

{pgdevkit-0.1.0 → pgdevkit-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pgdevkit
-Version: 0.1.0
+Version: 0.2.0
 Summary: A helper for developing with Postgres
 Requires-Python: >=3.14
 Requires-Dist: psycopg[binary]>=3.2.0
@@ -93,6 +93,14 @@ is already reachable at that address and skips container management if so.
 Set `PGDEVKIT_SKIP_CONTAINER=1` to always assume it's already there and skip
 that check too.
 
+To point at a local Postgres install instead of the podman container —
+useful when podman isn't available, or you'd rather use peer authentication
+as the current OS user — set `PGDEVKIT_TESTDB_HOST` to the unix socket
+directory (e.g. `/var/run/postgresql`) and `PGDEVKIT_TESTDB_PASSWORD=""`.
+The role named by `PGDEVKIT_TESTDB_USER` must exist and match your OS user
+(`CREATE ROLE <user> SUPERUSER LOGIN;`) and `pg_hba.conf` must allow `peer`
+auth for local connections (Debian/Ubuntu Postgres ships this by default).
+
 ## `pgdevkit.db` — helpers for application code
 
 Install with the `db` extra: `pip install pgdevkit[db]`.

{pgdevkit-0.1.0 → pgdevkit-0.2.0}/README.md

@@ -76,6 +76,14 @@ is already reachable at that address and skips container management if so.
 Set `PGDEVKIT_SKIP_CONTAINER=1` to always assume it's already there and skip
 that check too.
 
+To point at a local Postgres install instead of the podman container —
+useful when podman isn't available, or you'd rather use peer authentication
+as the current OS user — set `PGDEVKIT_TESTDB_HOST` to the unix socket
+directory (e.g. `/var/run/postgresql`) and `PGDEVKIT_TESTDB_PASSWORD=""`.
+The role named by `PGDEVKIT_TESTDB_USER` must exist and match your OS user
+(`CREATE ROLE <user> SUPERUSER LOGIN;`) and `pg_hba.conf` must allow `peer`
+auth for local connections (Debian/Ubuntu Postgres ships this by default).
+
 ## `pgdevkit.db` — helpers for application code
 
 Install with the `db` extra: `pip install pgdevkit[db]`.

pgdevkit-0.2.0/pgdevkit/connection.py

@@ -0,0 +1,101 @@
+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",
+)
+_AZURE_POSTGRES_HOST_SUFFIXES = (
+    ".postgres.database.azure.com",
+    ".postgres.cosmos.azure.com",
+)
+
+_default_credential = None
+_managed_identity_credential = None
+
+
+def detect_provider(host: str) -> Literal["azure_postgres", "databricks_lakebase"]:
+    """Classify a Postgres hostname for Entra ID auth: Databricks Lakebase
+    (needs credential exchange) or the default Azure Postgres AAD token
+    flow — the latter is the fallback for any non-Lakebase host, since
+    `entra_user` is itself the caller's assertion that Entra auth applies.
+    For a strict "is this actually Azure Database for PostgreSQL" check
+    (e.g. to decide PgBouncer-aware pooling), use `is_azure_postgres_host`
+    instead."""
+    if any(host.endswith(suffix) for suffix in _LAKEBASE_HOST_SUFFIXES):
+        return "databricks_lakebase"
+    return "azure_postgres"
+
+
+def is_azure_postgres_host(host: str) -> bool:
+    """True only for actual Azure Database for PostgreSQL hostnames —
+    unlike `detect_provider`, this is not a fallback default."""
+    return any(host.endswith(suffix) for suffix in _AZURE_POSTGRES_HOST_SUFFIXES)
+
+
+def get_azure_postgres_password(
+    *,
+    managed_identity: bool = False,
+    exclude_interactive_browser_credential: bool = True,
+) -> str:
+    """Fetch an Entra ID token to use as an Azure Postgres password.
+
+    `managed_identity=True` uses `ManagedIdentityCredential` (for workloads
+    running under an Azure-assigned identity); otherwise
+    `DefaultAzureCredential`, whose credential chain already falls back to
+    managed identity when no other credential is available. Both credential
+    objects are process-cached so repeated calls (one per new pooled
+    connection) don't re-probe the credential chain every time."""
+    try:
+        from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
+    except ImportError:
+        raise ImportError("Install azure-identity extra: pip install pgdevkit[azure]")
+    global _default_credential, _managed_identity_credential
+    if managed_identity:
+        if _managed_identity_credential is None:
+            _managed_identity_credential = ManagedIdentityCredential()
+        credential = _managed_identity_credential
+    else:
+        if _default_credential is None:
+            _default_credential = DefaultAzureCredential(
+                exclude_interactive_browser_credential=exclude_interactive_browser_credential
+            )
+        credential = _default_credential
+    token = credential.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,
+    managed_identity: bool = False,
+    exclude_interactive_browser_credential: bool = True,
+) -> 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(
+            managed_identity=managed_identity,
+            exclude_interactive_browser_credential=exclude_interactive_browser_credential,
+        )
+
+    netloc = f"{quote(entra_user, safe='')}:{quote(password, safe='')}@{host}{port}"
+    return urlunparse(parsed._replace(netloc=netloc))

{pgdevkit-0.1.0 → pgdevkit-0.2.0}/pgdevkit/db/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from .complex_types import ComplexHelper
 from .connection import PgPool
 from .crud import (
     pg_delete,
@@ -19,6 +20,7 @@ from .loader import SqlLoader
 from .model import PostgresTableModel
 
 __all__ = [
+    "ComplexHelper",
     "PgPool",
     "PostgresTableModel",
     "SqlLoader",

pgdevkit-0.1.0/skills/pgdevkit/references/complex_helper.py → pgdevkit-0.2.0/pgdevkit/db/complex_types.py

@@ -1,51 +1,53 @@
-"""
-ComplexHelper — psycopg adapter for PostgreSQL composite types, enums, and JSONB.
+from __future__ import annotations
 
-Use this when your project has custom PostgreSQL types (composite types, enums)
-that need to be registered with psycopg before inserting test data. Plain columns
-and JSONB are handled automatically by pgdevkit.testdb.schema's _insert_test_data;
-you only need this class for USER-DEFINED types, which that function does not
-yet support natively — wrap the connection it's given before the INSERT:
-
-    from pgdevkit.testdb import schema  # for reference, not a public extension point yet
-    from your_project.pg_complex_helper import ComplexHelper
-
-    async def insert_test_data_with_complex_types(json_file, table, force_reset, con):
-        helper = ComplexHelper(con)
-        schema_name, table_name = table.split(".")
-        complex_types = await helper.load_all_complex_types((schema_name, table_name))
-        rows = json.loads(json_file.read_text(encoding="utf-8"))
-        for row in rows:
-            for col, info in complex_types.items():
-                if col in row:
-                    row[col] = await helper.recursive_convert(row[col], info, con)
-        # then insert `rows` the same way schema.py's _insert_test_data does
-"""
+from typing import Any, Callable
 
 from psycopg import AsyncConnection
 from psycopg.rows import dict_row
-from typing import Any
 from psycopg.sql import Identifier
 from psycopg.types.composite import CompositeInfo, register_composite
-from psycopg.types.json import Jsonb
 from psycopg.types.enum import EnumInfo, register_enum
+from psycopg.types.json import Jsonb
+
+ComplexTypeInfo = CompositeInfo | EnumInfo | type[Jsonb] | None
 
 
 class ComplexHelper:
-    complex_types: dict[tuple[str, str], CompositeInfo | EnumInfo] = {}
+    """psycopg adapter for PostgreSQL composite types, enums, and JSONB.
 
-    def __init__(self, con: AsyncConnection):
-        self.con = con
-        self.system_complex_type_dict = None
+    Detects a table's non-scalar columns (composite types, enums, JSONB) and
+    converts plain dict/list Python values into the psycopg-registered types
+    those columns need, recursing into nested composite fields.
 
+    `normalizers` lets a caller reshape a composite value before conversion,
+    keyed by composite type name (e.g. a project with a `locale_labels`
+    composite type that needs locale-key backfilling before it's built) —
+    this is intentionally the only project-specific extension point; nothing
+    else about a project's types is hardcoded here.
+    """
+
+    def __init__(
+        self,
+        con: AsyncConnection,
+        normalizers: dict[str, Callable[[dict], dict]] | None = None,
+    ) -> None:
+        self.con = con
+        self.system_complex_type_dict: dict[Any, tuple[str, str]] | None = None
+        # Instance-scoped: a CompositeInfo/EnumInfo carries OIDs from a
+        # specific connection/database, so caching it on the class (shared
+        # across every connection) would leak stale OIDs across databases
+        # that happen to reuse the same type name — exactly the case for
+        # pgdevkit's per-worktree isolated test databases.
+        self.complex_types: dict[tuple[str, str], CompositeInfo | EnumInfo] = {}
         self.registered: set[CompositeInfo | EnumInfo] = set()
+        self._normalizers = normalizers or {}
 
-    async def load_complex_type_dict(self):
+    async def load_complex_type_dict(self) -> None:
         async with self.con.cursor(row_factory=dict_row) as cur:
             await cur.execute("""
         SELECT t.oid,
                 pg_catalog.format_type ( t.oid, NULL ) AS obj_name,
-				t.typtype
+                t.typtype
             FROM pg_catalog.pg_type t
             JOIN pg_catalog.pg_namespace n
                 ON n.oid = t.typnamespace
@@ -57,13 +59,9 @@ class ComplexHelper:
                 AND n.nspname <> 'information_schema'
                 AND n.nspname !~ '^pg_toast'""")
             system_complex_types = await cur.fetchall()
-            self.system_complex_type_dict = {
-                r["oid"]: (r["obj_name"], r["typtype"]) for r in system_complex_types
-            }
+            self.system_complex_type_dict = {r["oid"]: (r["obj_name"], r["typtype"]) for r in system_complex_types}
 
-    async def _load_complex_type_from_colinfos(
-        self, res: dict[str, Any] | None
-    ) -> CompositeInfo | EnumInfo | type[Jsonb] | None:
+    async def _load_complex_type_from_colinfos(self, res: dict[str, Any] | None) -> ComplexTypeInfo:
         if not res:
             return None
         if res["data_type"].lower() == "jsonb":
@@ -78,15 +76,13 @@ class ComplexHelper:
             return None
         udt_schema: str = res["udt_schema"]
         udt_name: str = res["udt_name"]
-        c = await self._get_complex_type(
-            f"{udt_schema}.{udt_name}", res["is_enum"], self.con
-        )
+        c = await self._get_complex_type(f"{udt_schema}.{udt_name}", res["is_enum"], self.con)
         await self._recurse_register(c, self.con)
         return c
 
     async def load_all_complex_types(
         self, table_name: tuple[str, str], include_generated: bool = False
-    ) -> dict[str, CompositeInfo | type[Jsonb] | EnumInfo | None]:
+    ) -> dict[str, ComplexTypeInfo]:
         if self.system_complex_type_dict is None:
             await self.load_complex_type_dict()
         colquery = """
@@ -100,7 +96,8 @@ class ComplexHelper:
             udt_schema, udt_name,
             e.enum_name is not null as is_enum
               from information_schema.columns c
-                 left join enum_types e on e.enum_schema=c.udt_schema and e.enum_name=c.udt_name
+                 left join enum_types e on e.enum_schema=c.udt_schema
+                     and (e.enum_name=c.udt_name or (c.data_type='ARRAY' and c.udt_name='_'||e.enum_name))
               where table_schema=%(schema)s and table_name = %(tbl)s and (is_generated <> 'ALWAYS' or %(include_generated)s)"""
         async with self.con.cursor(row_factory=dict_row) as cur:
             await cur.execute(
@@ -112,14 +109,9 @@ class ComplexHelper:
                 },
             )
             res = await cur.fetchall()
-            return {
-                r["column_name"]: await self._load_complex_type_from_colinfos(r)
-                for r in res
-            }
+            return {r["column_name"]: await self._load_complex_type_from_colinfos(r) for r in res}
 
-    async def load_complex_type(
-        self, table_name: tuple[str, str], col_name: str
-    ) -> CompositeInfo | type[Jsonb] | EnumInfo | None:
+    async def load_complex_type(self, table_name: tuple[str, str], col_name: str) -> ComplexTypeInfo:
         if self.system_complex_type_dict is None:
             await self.load_complex_type_dict()
         colquery = """
@@ -133,7 +125,8 @@ class ComplexHelper:
             udt_schema, udt_name,
             e.enum_name is not null as is_enum
               from information_schema.columns c
-            left join enum_types e on e.enum_schema=c.udt_schema and e.enum_name=c.udt_name
+            left join enum_types e on e.enum_schema=c.udt_schema
+                and (e.enum_name=c.udt_name or (c.data_type='ARRAY' and c.udt_name='_'||e.enum_name))
               where table_schema=%(schema)s and table_name = %(tbl)s
                 and column_name = %(col)s"""
         async with self.con.cursor(row_factory=dict_row) as cur:
@@ -142,18 +135,13 @@ class ComplexHelper:
                 {"schema": table_name[0], "tbl": table_name[1], "col": col_name},
             )
             res = await cur.fetchone()
-
             return await self._load_complex_type_from_colinfos(res)
 
-    async def _get_complex_type(
-        self, name: str, is_enum: bool, con: AsyncConnection
-    ) -> CompositeInfo | EnumInfo:
+    async def _get_complex_type(self, name: str, is_enum: bool, con: AsyncConnection) -> CompositeInfo | EnumInfo:
         if name.endswith("[]"):
             name = name[:-2]
-        schema, type_name = name.split(".")
-        if type_name.startswith(
-            "_"
-        ):  # the array type in PostgreSQL starts with an underscore
+        schema, type_name = name.split(".") if "." in name else ("public", name)
+        if type_name.startswith("_"):  # the array type in PostgreSQL starts with an underscore
             type_name = type_name[1:]
         if is_enum:
             ci = await EnumInfo.fetch(con, Identifier(schema, type_name))
@@ -165,12 +153,8 @@ class ComplexHelper:
             self.complex_types[(schema, type_name)] = ci
         return self.complex_types[(schema, type_name)]
 
-    async def _recurse_register(
-        self, info: CompositeInfo | EnumInfo, con: AsyncConnection
-    ):
-        assert self.system_complex_type_dict is not None, (
-            "System complex type dictionary not loaded"
-        )
+    async def _recurse_register(self, info: CompositeInfo | EnumInfo, con: AsyncConnection) -> None:
+        assert self.system_complex_type_dict is not None, "System complex type dictionary not loaded"
         if info not in self.registered:
             if isinstance(info, EnumInfo):
                 register_enum(info, con)
@@ -188,7 +172,7 @@ class ComplexHelper:
     async def recursive_convert(
         self,
         value: Any,
-        info: CompositeInfo | EnumInfo | type[Jsonb] | None,
+        info: ComplexTypeInfo,
         con: AsyncConnection,
     ) -> Any:
         if info is None:
@@ -197,20 +181,21 @@ class ComplexHelper:
             return None
         if self.system_complex_type_dict is None:
             await self.load_complex_type_dict()
+        if info == Jsonb:
+            # A JSONB column's value is wrapped whole, even if it's a list —
+            # only an array-of-composite/enum column recurses per element.
+            return Jsonb(value)
         if isinstance(value, list):
             return [await self.recursive_convert(item, info, con) for item in value]
         prms = {}
-        if info == Jsonb:
-            return Jsonb(value)
         if isinstance(value, str):
             assert isinstance(info, EnumInfo), f"Expected EnumInfo, got {type(info)}"
             return getattr(info.enum, value)  # Enum
-        assert isinstance(info, CompositeInfo), (
-            f"Expected CompositeInfo, got {type(info)}"
-        )
-        assert self.system_complex_type_dict is not None, (
-            "System complex type dictionary not loaded"
-        )
+        assert isinstance(info, CompositeInfo), f"Expected CompositeInfo, got {type(info)}"
+        assert self.system_complex_type_dict is not None, "System complex type dictionary not loaded"
+        normalizer = self._normalizers.get(info.name)
+        if normalizer is not None:
+            value = normalizer(value)
         for k, v in value.items():
             if v is None:
                 prms[k] = None
@@ -221,14 +206,10 @@ class ComplexHelper:
                 name, typtype = self.system_complex_type_dict[type_oid]
                 ci = await self._get_complex_type(name, typtype == "e", con)
                 if name.endswith("[]"):
-                    prms[k] = [
-                        await self.recursive_convert(item, ci, con) for item in v
-                    ]
+                    prms[k] = [await self.recursive_convert(item, ci, con) for item in v]
                 else:
                     prms[k] = await self.recursive_convert(v, ci, con)
             else:
                 prms[k] = v
-        assert info.python_type is not None, (
-            f"Python type for {info.name} is null, maybe an array?"
-        )
+        assert info.python_type is not None, f"Python type for {info.name} is null, maybe an array?"
         return info.python_type(**prms) if prms else None

pgdevkit-0.2.0/pgdevkit/db/connection.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Literal
+
+from psycopg_pool import AsyncConnectionPool, AsyncNullConnectionPool
+
+from ..connection import detect_provider, get_azure_postgres_password, is_azure_postgres_host
+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`.
+
+    `use_null_pool="auto"` (the default) switches to a null pool — no local
+    pooling — when the host is Azure Postgres, since Azure's own PgBouncer
+    (transaction-mode) does the pooling; `connection_kwargs` then
+    auto-includes `prepare_threshold=None` too, since PgBouncer transaction
+    mode doesn't support server-side prepared statements. Pass an explicit
+    `use_null_pool`/`connection_kwargs` to override either.
+    """
+
+    def __init__(
+        self,
+        env_prefix: str = "POSTGRES_",
+        max_size: int = 40,
+        *,
+        entra_user: str | None = None,
+        credential_kind: Literal["default_azure", "managed_identity"] = "default_azure",
+        exclude_interactive_browser_credential: bool = True,
+        max_lifetime: float | None = None,
+        use_null_pool: bool | Literal["auto"] = "auto",
+        connection_kwargs: dict | None = None,
+        dsn_params: dict[str, str] | None = None,
+    ) -> None:
+        self._env_prefix = env_prefix
+        self._max_size = max_size
+        self._entra_user = entra_user
+        self._credential_kind = credential_kind
+        self._exclude_interactive_browser_credential = exclude_interactive_browser_credential
+        self._max_lifetime = max_lifetime
+        self._use_null_pool = use_null_pool
+        self._connection_kwargs = connection_kwargs
+        self._dsn_params = dsn_params or {}
+        self._pool: AsyncConnectionPool | AsyncNullConnectionPool | None = None
+
+    def _is_azure_postgres(self, host: str) -> bool:
+        return is_azure_postgres_host(host)
+
+    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,
+                    managed_identity=self._credential_kind == "managed_identity",
+                    exclude_interactive_browser_credential=self._exclude_interactive_browser_credential,
+                )
+
+        dsn = f"host={host} port={port} dbname={dbname} user={user} password={password}"
+        for key, value in self._dsn_params.items():
+            dsn += f" {key}={value}"
+        return dsn
+
+    async def open(self) -> None:
+        if self._pool is None:
+            host = os.environ[self._env_prefix + "HOST"]
+            is_azure_postgres = self._is_azure_postgres(host)
+            use_null_pool = is_azure_postgres if self._use_null_pool == "auto" else self._use_null_pool
+            connection_kwargs = self._connection_kwargs
+            if connection_kwargs is None:
+                # PgBouncer (transaction mode) doesn't support prepared statements
+                connection_kwargs = {"prepare_threshold": None} if is_azure_postgres else {}
+            max_lifetime = 3600.0 if self._max_lifetime is None else self._max_lifetime
+            if use_null_pool:
+                self._pool = AsyncNullConnectionPool(
+                    conninfo=self._dsn,
+                    open=False,
+                    max_size=self._max_size,
+                    max_lifetime=max_lifetime,
+                    check=AsyncNullConnectionPool.check_connection,
+                    kwargs=connection_kwargs,
+                )
+            else:
+                self._pool = AsyncConnectionPool(
+                    conninfo=self._dsn,
+                    open=False,
+                    max_size=self._max_size,
+                    max_lifetime=max_lifetime,
+                    check=AsyncConnectionPool.check_connection,
+                    kwargs=connection_kwargs,
+                )
+        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()
+
+    @property
+    def raw_pool(self) -> AsyncConnectionPool | AsyncNullConnectionPool:
+        """The underlying psycopg_pool object — for integrating with other
+        pool-consuming libraries (e.g. Procrastinate's `open_async(pool=...)`)
+        that need it directly rather than going through `connection()`."""
+        if self._pool is None:
+            raise RuntimeError("Call open() first (e.g. in app startup).")
+        return self._pool