interlaced 2.0.0a2__py3-none-any.whl

interlace/engines/duckdb.py

@@ -0,0 +1,249 @@
+"""DuckDB engine adapter — the default local engine and federation hub.
+
+Everything crosses the boundary as Arrow: :meth:`fetch` streams results as a
+``pyarrow.RecordBatchReader`` (zero-copy, single pass) and :meth:`load` registers
+an Arrow reader and writes it with one ``CREATE TABLE AS`` / ``INSERT``. Blocking
+DuckDB calls run in a worker thread; each call uses its own ``cursor()`` so reads
+proceed concurrently (DuckDB MVCC), while the DAG guarantees no two tasks write
+the same table at once.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import Sequence
+from uuid import uuid4
+
+import duckdb
+import pyarrow as pa
+import tenacity
+from sqlglot import exp
+
+from interlace.engines.base import EngineAdapter, EngineCaps, LoadMode
+from interlace.ir.relation import TableRef
+
+_DUCKDB_CAPS = EngineCaps(
+    supports_create_or_replace=True,
+    supports_star_exclude=True,
+)
+
+
+# DuckLake uses optimistic concurrency: a concurrent writer's commit surfaces as
+# a TransactionException. Our write batches are whole-transaction idempotent
+# (CREATE OR REPLACE / DELETE+INSERT run as one unit), so a short retry is safe.
+_commit_retry = tenacity.retry(
+    retry=tenacity.retry_if_exception_type(duckdb.TransactionException),
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_exponential_jitter(initial=0.1, max=1.0),
+    reraise=True,
+)
+
+
+def _affected(cur: duckdb.DuckDBPyConnection) -> int:
+    """DML/CTAS/COPY return their affected-row count as a one-cell result; DDL returns
+    nothing. Never raises — row stats are best-effort decoration, not correctness."""
+    try:
+        rows = cur.fetchall()
+        return int(rows[0][0]) if rows and rows[0] and isinstance(rows[0][0], int) else 0
+    except Exception:
+        return 0
+
+
+class DuckDBAdapter(EngineAdapter):
+    """Executes canonical ASTs and moves Arrow data in and out of a DuckDB database."""
+
+    dialect = "duckdb"
+    caps = _DUCKDB_CAPS
+
+    def __init__(self, connection: duckdb.DuckDBPyConnection, session_init: Sequence[str] = ()) -> None:
+        self._conn = connection
+        # Statements re-applied on every cursor — SESSION-LOCAL state only (USE).
+        # Anything instance-wide (LOAD, secrets, ATTACH) belongs at connect time:
+        # re-running catalog writes here races across concurrent cursors.
+        self._session_init = list(session_init)
+        self._attached: list[str] = []  # aliases to DETACH on close (see close())
+
+    def _cursor(self) -> duckdb.DuckDBPyConnection:
+        cur = self._conn.cursor()
+        for statement in self._session_init:
+            cur.execute(statement)
+        return cur
+
+    @classmethod
+    def in_memory(cls) -> DuckDBAdapter:
+        return cls(duckdb.connect(":memory:"))
+
+    @classmethod
+    def connect(cls, path: str) -> DuckDBAdapter:
+        return cls(duckdb.connect(path))
+
+    @classmethod
+    def connect_ducklake(
+        cls,
+        catalog: str,
+        *,
+        alias: str = "warehouse",
+        data_path: str | None = None,
+        metadata_schema: str | None = None,
+        secrets: Sequence[str] = (),
+        extensions: Sequence[str] = (),
+    ) -> DuckDBAdapter:
+        """Open a DuckLake warehouse that needs attach options and/or credentials —
+        remote catalogs (``postgres:…``) and object-store ``data_path``s can't ride the
+        plain ``duckdb.connect("ducklake:…")`` shortcut. Opens ``:memory:``, installs
+        the extensions, issues the ``CREATE SECRET`` statements, ATTACHes the DuckLake
+        with the options, and makes it the default catalog."""
+        conn = duckdb.connect(":memory:")
+        for extension in extensions:
+            conn.execute(f"INSTALL {extension}; LOAD {extension};")
+        for statement in secrets:
+            conn.execute(statement)
+        options: list[str] = []
+        if data_path:
+            options.append(f"DATA_PATH '{data_path.replace(chr(39), chr(39) * 2)}'")
+        if metadata_schema:
+            options.append(f"METADATA_SCHEMA '{metadata_schema.replace(chr(39), chr(39) * 2)}'")
+        options_sql = f" ({', '.join(options)})" if options else ""
+        escaped = catalog.replace("'", "''")
+        alias_sql = exp.to_identifier(alias).sql("duckdb")
+        conn.execute(f"ATTACH IF NOT EXISTS '{escaped}' AS {alias_sql}{options_sql}")
+        conn.execute(f"USE {alias_sql}")
+        # LOAD, secrets, and ATTACH are all instance-wide — they carry into every
+        # cursor and must run ONCE (re-running CREATE OR REPLACE SECRET per cursor
+        # races: concurrent cursors hit "catalog write-write conflict on alter").
+        # Only the default catalog is session state, so that is all a cursor re-applies.
+        return cls(conn, session_init=[f"USE {alias_sql}"])
+
+    def close(self) -> None:
+        # DETACH long-lived attaches first: DuckLake leaks its DatabaseInstance when
+        # concurrent cursors were used (duckdb 1.5.4), which would otherwise keep the
+        # attached databases' file handles locked for the rest of the process.
+        for alias in self._attached:
+            with contextlib.suppress(Exception):
+                self._conn.execute(f"DETACH {exp.to_identifier(alias).sql('duckdb')}")
+        self._attached.clear()
+        self._conn.close()
+
+    def attach(self, alias: str, uri: str) -> None:
+        """ATTACH another database (duckdb/sqlite/postgres/... URI) under ``alias``."""
+        escaped = uri.replace("'", "''")
+        self._conn.execute(f"ATTACH IF NOT EXISTS '{escaped}' AS {exp.to_identifier(alias).sql('duckdb')}")
+        self._attached.append(alias)
+
+    # --- identifier helpers -------------------------------------------------
+
+    def _table_sql(self, table: TableRef) -> str:
+        return table.to_expr().sql(dialect=self.dialect)
+
+    # --- EngineAdapter ------------------------------------------------------
+
+    async def execute(self, ast: exp.Expression) -> None:
+        await self.execute_sql(self.transpile(ast))
+
+    async def execute_all(self, statements: Sequence[exp.Expression]) -> list[int]:
+        return await asyncio.to_thread(self._execute_all_sync, [self.transpile(s) for s in statements])
+
+    async def fetch(self, ast: exp.Expression) -> pa.RecordBatchReader:
+        return await self.fetch_sql(self.transpile(ast))
+
+    async def load(self, table: TableRef, reader: pa.RecordBatchReader, mode: LoadMode) -> int:
+        return await asyncio.to_thread(self._load_sync, table, reader, mode)
+
+    async def create_view(self, name: TableRef, target: TableRef) -> None:
+        await self.execute_sql(
+            f"CREATE OR REPLACE VIEW {self._table_sql(name)} AS SELECT * FROM {self._table_sql(target)}"
+        )
+
+    # --- raw / convenience (used by the state store and tests) --------------
+
+    async def execute_sql(self, sql: str) -> None:
+        await asyncio.to_thread(self._execute_sync, sql)
+
+    async def fetch_sql(self, sql: str) -> pa.RecordBatchReader:
+        return await asyncio.to_thread(self._fetch_sync, sql)
+
+    async def create_schema(self, name: str) -> None:
+        await self.execute_sql(f"CREATE SCHEMA IF NOT EXISTS {exp.to_identifier(name).sql(dialect=self.dialect)}")
+
+    async def table_exists(self, table: TableRef) -> bool:
+        return await asyncio.to_thread(self._table_exists_sync, table)
+
+    async def describe(self, table: TableRef) -> dict[str, str]:
+        return await asyncio.to_thread(self._describe_sync, table)
+
+    # --- sync workers (run in a thread) -------------------------------------
+
+    @_commit_retry
+    def _execute_sync(self, sql: str) -> None:
+        cur = self._cursor()
+        try:
+            cur.execute(sql)
+        finally:
+            cur.close()
+
+    @_commit_retry
+    def _execute_all_sync(self, sqls: list[str]) -> list[int]:
+        cur = self._cursor()
+        counts: list[int] = []
+        try:
+            cur.execute("BEGIN")
+            for sql in sqls:
+                cur.execute(sql)
+                counts.append(_affected(cur))
+            cur.execute("COMMIT")
+        except Exception:
+            cur.execute("ROLLBACK")
+            raise
+        finally:
+            cur.close()
+        return counts
+
+    def _fetch_sync(self, sql: str) -> pa.RecordBatchReader:
+        # The reader keeps the underlying result alive after the cursor is dropped.
+        cur = self._cursor()
+        cur.execute(sql)
+        return cur.to_arrow_reader()
+
+    @_commit_retry
+    def _load_sync(self, table: TableRef, reader: pa.RecordBatchReader, mode: LoadMode) -> int:
+        cur = self._cursor()
+        src = f"__interlace_src_{uuid4().hex}"
+        cur.register(src, reader)
+        try:
+            target = self._table_sql(table)
+            if mode == "create":
+                cur.execute(f"CREATE OR REPLACE TABLE {target} AS SELECT * FROM {src}")
+            else:
+                cur.execute(f"INSERT INTO {target} SELECT * FROM {src}")
+            return _affected(cur)
+        finally:
+            cur.unregister(src)
+            cur.close()
+
+    def _table_exists_sync(self, table: TableRef) -> bool:
+        # information_schema spans every attached catalog: pin to the ref's catalog
+        # (or the session default) so same-named tables elsewhere don't collide.
+        cur = self._cursor()
+        try:
+            row = cur.execute(
+                "SELECT count(*) FROM information_schema.tables WHERE table_schema = ? AND table_name = ? "
+                "AND table_catalog = coalesce(?, current_database())",
+                [table.schema, table.name, table.catalog],
+            ).fetchone()
+        finally:
+            cur.close()
+        return bool(row and row[0])
+
+    def _describe_sync(self, table: TableRef) -> dict[str, str]:
+        cur = self._cursor()
+        try:
+            rows = cur.execute(
+                "SELECT column_name, data_type FROM information_schema.columns "
+                "WHERE table_schema = ? AND table_name = ? "
+                "AND table_catalog = coalesce(?, current_database()) ORDER BY ordinal_position",
+                [table.schema, table.name, table.catalog],
+            ).fetchall()
+        finally:
+            cur.close()
+        return dict(rows)

interlace/engines/postgres.py

@@ -0,0 +1,157 @@
+"""Postgres engine adapter — the first native remote engine (ADBC transport).
+
+Strategies execute *inside* Postgres: canonical ASTs transpile to the postgres
+dialect and run over one ADBC connection; results come back as Arrow and bulk
+loads go in via ``adbc_ingest`` — columnar end to end, no row-format hop.
+
+Capability honesty drives the strategy fallbacks: Postgres has no
+``CREATE OR REPLACE TABLE`` (FullRefresh falls back to DROP+CREATE) and no
+star-EXCLUDE projection (scd_type_2 refuses with a clear error). ``merge_by_key``
+and ``full_merge`` are portable by construction (DELETE+INSERT / set difference).
+
+The ADBC connection is synchronous: calls run in a worker thread behind a lock
+(one statement at a time per engine — remote engines parallelise internally).
+Requires the ``adbc`` extra (``pip install 'interlaced[adbc]'``).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Sequence
+from typing import Any
+
+import pyarrow as pa
+from sqlglot import exp
+
+from interlace.engines.base import EngineAdapter, EngineCaps, LoadMode
+from interlace.exceptions import ConfigurationError
+from interlace.ir.relation import TableRef
+
+_POSTGRES_CAPS = EngineCaps(
+    supports_create_or_replace=False,  # no CREATE OR REPLACE TABLE -> DROP+CREATE fallback
+    supports_star_exclude=False,  # no SELECT * EXCLUDE -> scd_type_2 unsupported for now
+)
+
+
+class PostgresAdapter(EngineAdapter):
+    """Executes canonical ASTs inside Postgres; Arrow in and out via ADBC."""
+
+    dialect = "postgres"
+    caps = _POSTGRES_CAPS
+
+    def __init__(self, connection: Any) -> None:  # adbc_driver_postgresql.dbapi.Connection
+        self._conn = connection
+        self._lock = threading.Lock()
+
+    @classmethod
+    def connect(cls, dsn: str) -> PostgresAdapter:
+        try:
+            import adbc_driver_postgresql.dbapi as dbapi  # type: ignore[import-untyped]
+        except ImportError as exc:  # pragma: no cover - import guard
+            raise ConfigurationError(
+                "the postgres engine needs the 'adbc' extra: pip install 'interlaced[adbc]'"
+            ) from exc
+        return cls(dbapi.connect(dsn))
+
+    def close(self) -> None:
+        self._conn.close()
+
+    # --- EngineAdapter ------------------------------------------------------
+
+    async def execute(self, ast: exp.Expression) -> None:
+        await self.execute_sql(self.transpile(ast))
+
+    async def execute_all(self, statements: Sequence[exp.Expression]) -> list[int]:
+        return await asyncio.to_thread(self._execute_all_sync, [self.transpile(s) for s in statements])
+
+    async def fetch(self, ast: exp.Expression) -> pa.RecordBatchReader:
+        return await self.fetch_sql(self.transpile(ast))
+
+    async def load(self, table: TableRef, reader: pa.RecordBatchReader, mode: LoadMode) -> int:
+        return await asyncio.to_thread(self._load_sync, table, reader, mode)
+
+    async def create_view(self, name: TableRef, target: TableRef) -> None:
+        await self.execute_sql(
+            f"CREATE OR REPLACE VIEW {self._table_sql(name)} AS SELECT * FROM {self._table_sql(target)}"
+        )
+
+    async def execute_sql(self, sql: str) -> None:
+        await asyncio.to_thread(self._execute_sync, sql)
+
+    async def fetch_sql(self, sql: str) -> pa.RecordBatchReader:
+        return await asyncio.to_thread(self._fetch_sync, sql)
+
+    async def create_schema(self, name: str) -> None:
+        await self.execute_sql(f"CREATE SCHEMA IF NOT EXISTS {exp.to_identifier(name).sql(dialect=self.dialect)}")
+
+    async def table_exists(self, table: TableRef) -> bool:
+        return await asyncio.to_thread(self._table_exists_sync, table)
+
+    async def describe(self, table: TableRef) -> dict[str, str]:
+        return await asyncio.to_thread(self._describe_sync, table)
+
+    # --- sync workers (one connection; ADBC is synchronous) ------------------
+
+    def _table_sql(self, table: TableRef) -> str:
+        return table.to_expr().sql(dialect=self.dialect)
+
+    def _execute_sync(self, sql: str) -> None:
+        with self._lock, self._conn.cursor() as cur:
+            cur.execute(sql)
+            self._conn.commit()
+
+    def _execute_all_sync(self, sqls: list[str]) -> list[int]:
+        counts: list[int] = []
+        with self._lock:
+            try:
+                with self._conn.cursor() as cur:
+                    for sql in sqls:
+                        cur.execute(sql)
+                        rowcount = getattr(cur, "rowcount", -1)
+                        counts.append(rowcount if isinstance(rowcount, int) and rowcount > 0 else 0)
+                self._conn.commit()  # ADBC autocommit is off: this is one transaction
+            except Exception:
+                self._conn.rollback()
+                raise
+        return counts
+
+    def _fetch_sync(self, sql: str) -> pa.RecordBatchReader:
+        with self._lock, self._conn.cursor() as cur:
+            cur.execute(sql)
+            table = cur.fetch_arrow_table()  # materialised: keeps the reader cursor-independent
+            self._conn.commit()
+        return table.to_reader()
+
+    def _load_sync(self, table: TableRef, reader: pa.RecordBatchReader, mode: LoadMode) -> int:
+        ingest_mode = "replace" if mode == "create" else "append"  # "create" == CREATE OR REPLACE semantics
+        with self._lock:
+            try:
+                with self._conn.cursor() as cur:
+                    loaded = cur.adbc_ingest(table.name, reader, mode=ingest_mode, db_schema_name=table.schema)
+                self._conn.commit()
+            except Exception:
+                self._conn.rollback()
+                raise
+        return int(loaded) if isinstance(loaded, int) and loaded > 0 else 0
+
+    def _table_exists_sync(self, table: TableRef) -> bool:
+        with self._lock, self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT count(*) FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2",
+                (table.schema, table.name),
+            )
+            row = cur.fetchone()
+            self._conn.commit()
+        return bool(row and row[0])
+
+    def _describe_sync(self, table: TableRef) -> dict[str, str]:
+        with self._lock, self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT column_name, data_type FROM information_schema.columns "
+                "WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position",
+                (table.schema, table.name),
+            )
+            rows = cur.fetchall()
+            self._conn.commit()
+        return dict(rows)

interlace/engines/quack.py

@@ -0,0 +1,113 @@
+"""Quack engine adapter — a remote DuckDB warehouse over the quack protocol.
+
+Connects to a warehouse served by ``interlace serve --quack`` (or any
+``CALL quack_serve(...)``). Every statement is shipped to the server via the
+``quack_query`` table function — full SQL pass-through with results streamed
+back as Arrow — because quack's catalog ATTACH only resolves the server's main
+schema while the protocol is in beta. Arrow loads go through the attached
+remote catalog (DDL resolves schema-qualified names correctly over the wire).
+Multi-statement plans are sent as one BEGIN/COMMIT payload so they stay atomic
+server-side. Requires DuckDB >= 1.5.3 (quack is a core extension there).
+"""
+
+from __future__ import annotations
+
+from uuid import uuid4
+
+import duckdb
+import pyarrow as pa
+from sqlglot import exp
+
+from interlace.engines.base import LoadMode
+from interlace.engines.duckdb import DuckDBAdapter
+from interlace.ir.relation import TableRef
+
+_REMOTE = "__interlace_remote"
+
+
+def sql_literal(value: str) -> str:
+    return exp.Literal.string(value).sql(dialect="duckdb")
+
+
+class QuackAdapter(DuckDBAdapter):
+    """Executes canonical ASTs against a quack-served warehouse."""
+
+    def __init__(self, connection: duckdb.DuckDBPyConnection, uri: str) -> None:
+        super().__init__(connection)
+        self._uri = uri
+
+    @classmethod
+    def connect(cls, path: str, token: str | None = None) -> QuackAdapter:
+        """Connect to a ``quack:<host>:<port>`` warehouse; ``token`` is the serve token."""
+        conn = duckdb.connect(":memory:")
+        if token:
+            conn.execute(f"CREATE SECRET {_REMOTE}_secret (TYPE quack, TOKEN {sql_literal(token)})")
+        conn.execute(f"ATTACH {sql_literal(path)} AS {_REMOTE}")  # used for Arrow loads; also validates reachability
+        return cls(conn, path)
+
+    # --- sync workers: route SQL through quack_query -------------------------
+
+    def _remote_sync(self, cur: duckdb.DuckDBPyConnection, sql: str) -> duckdb.DuckDBPyConnection:
+        return cur.execute("FROM quack_query(?, ?)", [self._uri, sql])
+
+    def _execute_sync(self, sql: str) -> None:
+        cur = self._conn.cursor()
+        try:
+            self._remote_sync(cur, sql)
+        finally:
+            cur.close()
+
+    def _execute_all_sync(self, sqls: list[str]) -> list[int]:
+        # One payload, transactional server-side; quack executes it as a unit —
+        # per-statement affected counts are not observable over the wire.
+        self._execute_sync(";\n".join(["BEGIN", *sqls, "COMMIT"]))
+        return [0] * len(sqls)
+
+    def _fetch_sync(self, sql: str) -> pa.RecordBatchReader:
+        cur = self._conn.cursor()
+        return self._remote_sync(cur, sql).to_arrow_reader()
+
+    def _load_sync(self, table: TableRef, reader: pa.RecordBatchReader, mode: LoadMode) -> int:
+        # Register the reader locally; schema-qualified DDL against the attached
+        # remote catalog streams the batches over the wire. The attach caches the
+        # remote catalog, so refresh it to see schemas created since we connected.
+        cur = self._conn.cursor()
+        cur.execute("CALL quack_clear_cache()")
+        src = f"__interlace_src_{uuid4().hex}"
+        cur.register(src, reader)
+        try:
+            target = exp.table_(table.name, db=table.schema, catalog=_REMOTE).sql(dialect=self.dialect)
+            if mode == "create":
+                cur.execute(f"CREATE OR REPLACE TABLE {target} AS SELECT * FROM {src}")
+            else:
+                cur.execute(f"INSERT INTO {target} SELECT * FROM {src}")
+            rows = cur.fetchall()
+            return int(rows[0][0]) if rows and rows[0] and isinstance(rows[0][0], int) else 0
+        finally:
+            cur.unregister(src)
+            cur.close()
+
+    def _table_exists_sync(self, table: TableRef) -> bool:
+        sql = (
+            "SELECT count(*) FROM information_schema.tables "
+            f"WHERE table_schema = {sql_literal(table.schema)} AND table_name = {sql_literal(table.name)}"
+        )
+        cur = self._conn.cursor()
+        try:
+            row = self._remote_sync(cur, sql).fetchone()
+        finally:
+            cur.close()
+        return bool(row and row[0])
+
+    def _describe_sync(self, table: TableRef) -> dict[str, str]:
+        sql = (
+            "SELECT column_name, data_type FROM information_schema.columns "
+            f"WHERE table_schema = {sql_literal(table.schema)} AND table_name = {sql_literal(table.name)} "
+            "ORDER BY ordinal_position"
+        )
+        cur = self._conn.cursor()
+        try:
+            rows = self._remote_sync(cur, sql).fetchall()
+        finally:
+            cur.close()
+        return dict(rows)

interlace/engines/registry.py

@@ -0,0 +1,125 @@
+"""Named engine registry — multi-engine routing without a global singleton.
+
+A project may declare several engines in ``interlace.yaml``. The registry opens
+them lazily (so unused remote engines stay cold), resolves models by name, and
+closes everything on teardown. Single-engine projects use one entry named
+``default`` synthesised from the top-level warehouse fields.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator, Mapping
+from typing import TypeVar
+
+from interlace.engines.base import EngineAdapter
+from interlace.exceptions import ConfigurationError, PlanError
+
+_T = TypeVar("_T", bound=EngineAdapter)
+
+Opener = Callable[[str], EngineAdapter]
+
+
+class EngineRegistry(Mapping[str, EngineAdapter]):
+    """Lazy map of engine name → open :class:`EngineAdapter`.
+
+    ``get(name)`` opens on first use. Unknown names raise
+    :class:`ConfigurationError`. ``default`` is the project's default engine name.
+    """
+
+    def __init__(
+        self,
+        names: Iterable[str],
+        opener: Opener,
+        *,
+        default: str = "default",
+        attach_uris: Mapping[str, str] | None = None,
+    ) -> None:
+        self._names = frozenset(names)
+        self._opener = opener
+        self.default = default
+        self._cache: dict[str, EngineAdapter] = {}
+        # engine name -> URI another DuckDB-family engine could ATTACH (transfer
+        # fast lane); absent/None = only reachable through its own adapter.
+        self.attach_uris: dict[str, str] = dict(attach_uris or {})
+        if default not in self._names:
+            raise ConfigurationError(
+                f"default_engine {default!r} is not a configured engine",
+                details={"engines": sorted(self._names)},
+            )
+
+    def __getitem__(self, name: str) -> EngineAdapter:
+        return self.get(name)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(sorted(self._names))
+
+    def __len__(self) -> int:
+        return len(self._names)
+
+    def __contains__(self, name: object) -> bool:
+        return isinstance(name, str) and name in self._names
+
+    def get(self, name: str | None = None) -> EngineAdapter:  # type: ignore[override]
+        """Return the adapter for ``name`` (or the default engine), opening if needed."""
+        key = self.default if name is None else name
+        if key not in self._names:
+            raise ConfigurationError(
+                f"unknown engine {key!r}",
+                details={"engines": sorted(self._names)},
+            )
+        if key not in self._cache:
+            self._cache[key] = self._opener(key)
+        return self._cache[key]
+
+    def require(self, name: str, *, model: str | None = None) -> EngineAdapter:
+        """Like :meth:`get`, but raise :class:`PlanError` when the engine is missing."""
+        if name not in self._names:
+            where = f" for model {model!r}" if model else ""
+            raise PlanError(
+                f"engine {name!r}{where} is not configured",
+                details={"engines": sorted(self._names)},
+            )
+        return self.get(name)
+
+    def close(self) -> None:
+        """Close every opened adapter (idempotent for already-closed adapters)."""
+        for adapter in self._cache.values():
+            close = getattr(adapter, "close", None)
+            if callable(close):
+                close()
+        self._cache.clear()
+
+
+def as_registry(
+    engine: EngineAdapter | None = None,
+    engines: Mapping[str, EngineAdapter] | EngineRegistry | None = None,
+    *,
+    default: str = "default",
+) -> EngineRegistry:
+    """Normalise the ``engine=`` / ``engines=`` apply kwargs into a registry.
+
+    Existing call sites that pass a single ``engine=`` keep working: that adapter
+    is registered under ``default``.
+    """
+    if isinstance(engines, EngineRegistry):
+        if engine is not None:
+            # Prefer the registry; a bare engine is only a fallback for default.
+            pass
+        return engines
+
+    cache: dict[str, EngineAdapter] = dict(engines or {})
+    if engine is not None:
+        cache.setdefault(default, engine)
+    if not cache:
+        raise PlanError("apply requires engine= or engines=")
+
+    def opener(name: str) -> EngineAdapter:
+        try:
+            return cache[name]
+        except KeyError as exc:
+            raise ConfigurationError(
+                f"unknown engine {name!r}",
+                details={"engines": sorted(cache)},
+            ) from exc
+
+    return EngineRegistry(cache.keys(), opener, default=default if default in cache else next(iter(cache)))