Mdmapp migration

pgdevkit/connection.py

@@ -7,23 +7,63 @@
     ".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: Databricks Lakebase (needs credential
-    exchange) or the default Azure Postgres AAD token flow."""
+    """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 get_azure_postgres_password() -> str:
+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
+        from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
     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")
+    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
 
 
@@ -33,6 +73,8 @@ def build_conninfo(
     *,
     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
@@ -50,7 +92,10 @@ def build_conninfo(
 
         password = get_lakebase_password(databricks_workspace_host, databricks_instance)
     else:
-        password = get_azure_postgres_password()
+        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/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 .model import PostgresTableModel
 
 __all__ = [
+    "ComplexHelper",
     "PgPool",
     "PostgresTableModel",
     "SqlLoader",

skills/pgdevkit/references/complex_helper.py → 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 @@ async def load_complex_type_dict(self):
                 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 @@ async def _load_complex_type_from_colinfos(
             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 @@ async def load_all_complex_types(
             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 @@ async def load_all_complex_types(
                 },
             )
             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 @@ async def load_complex_type(
             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 @@ async def load_complex_type(
                 {"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 @@ async def _get_complex_type(
             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 @@ async def _recurse_register(
     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 @@ async def recursive_convert(
             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 @@ async def recursive_convert(
                 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