feat: IbisBackend adoption capabilities — raw driver access, cross-dialect transactions, session-option preservation (#100)

src/mountainash_data/backends/ibis/_adoption.py

@@ -0,0 +1,70 @@
+"""Session-option snapshot / restore / apply for adoption (Gap 1)."""
+
+from __future__ import annotations
+
+import typing as t
+import warnings
+
+from mountainash_data.backends.ibis._raw import raw_execute, raw_fetch_scalar
+from mountainash_data.backends.ibis.dialects._registry import SessionOption
+from mountainash_data.core._warn import warn_once
+
+
+def snapshot_options(
+    raw_handle: t.Any, options: tuple[SessionOption, ...]
+) -> dict[str, t.Any]:
+    """Read the current value of each option that has a read_sql, via the shared
+    _raw transport (finding 3 — cursor-safe across drivers).
+
+    An option that cannot be read is NOT silently skipped — it WARNS, because a
+    value we cannot snapshot cannot be restored, and "faithful" preservation must
+    signal when it can't be faithful (Codex review). Options with read_sql=None
+    are skipped without a warning (nothing to snapshot by design)."""
+    snap: dict[str, t.Any] = {}
+    for opt in options:
+        if opt.read_sql is None:
+            continue
+        try:
+            value = raw_fetch_scalar(raw_handle, opt.read_sql)
+        except Exception as exc:  # noqa: BLE001 — warn, don't fail adoption
+            warnings.warn(
+                f"could not snapshot session option {opt.name!r}; it will not be "
+                f"restored ({exc!r})",
+                stacklevel=2,
+            )
+            continue
+        snap[opt.name] = value
+    return snap
+
+
+def restore_options(
+    raw_handle: t.Any,
+    options: tuple[SessionOption, ...],
+    snapshot: dict[str, t.Any],
+) -> None:
+    """Replay each captured value via its render_set statement (shared transport)."""
+    by_name = {o.name: o for o in options}
+    for name, value in snapshot.items():
+        opt = by_name.get(name)
+        if opt is not None:
+            raw_execute(raw_handle, opt.render_set(value))
+
+
+def apply_options(
+    raw_handle: t.Any,
+    options: tuple[SessionOption, ...],
+    values: dict[str, t.Any],
+) -> None:
+    """Apply caller-declared end-state values. Unknown names are ignored, with a
+    warning (each name warns once — see warn_once)."""
+    by_name = {o.name: o for o in options}
+    for name, value in values.items():
+        opt = by_name.get(name)
+        if opt is not None:
+            raw_execute(raw_handle, opt.render_set(value))
+        else:
+            warn_once(
+                f"apply_options:{name}",
+                f"session option {name!r} is not a declared adoption mutation "
+                f"for this backend; ignored",
+            )

src/mountainash_data/backends/ibis/_raw.py

@@ -0,0 +1,74 @@
+"""Shared native-handle statement transport (Gap 3, fable finding 3).
+
+The single seam for "run one statement on the raw driver handle", shared by
+_transaction (BEGIN/COMMIT/ROLLBACK) and, later, _adoption (session
+snapshot/restore). ``.execute()`` is NOT uniform across DBAPI drivers, so this
+falls back to ``.cursor().execute()``. A per-dialect
+``DialectSpec.raw_execute_hook`` overrides the write transport entirely.
+"""
+from __future__ import annotations
+
+import typing as t
+
+
+def raw_execute(
+    handle: t.Any,
+    sql: str,
+    *,
+    hook: t.Optional[t.Callable[[t.Any, str], None]] = None,
+) -> None:
+    """Execute ``sql`` on the native handle (no result).
+
+    hook, if given, is the whole transport. Else use ``handle.execute`` when
+    present (duckdb / sqlite / psycopg3 / pyodbc), otherwise a cursor
+    (mysqlclient / oracledb / trino), closing the cursor afterward.
+    """
+    if hook is not None:
+        hook(handle, sql)
+        return
+    execute = getattr(handle, "execute", None)
+    if callable(execute):
+        execute(sql)
+        return
+    cur = handle.cursor()
+    try:
+        cur.execute(sql)
+    finally:
+        close = getattr(cur, "close", None)
+        if callable(close):
+            close()
+
+
+def raw_fetch_scalar(
+    handle: t.Any,
+    sql: str,
+    *,
+    hook: t.Optional[t.Callable[[t.Any, str], None]] = None,
+) -> t.Any:
+    """Run ``sql`` and return the first column of the first row, or ``None``.
+
+    Same execute-or-cursor transport as :func:`raw_execute`. A void ``hook``
+    cannot return rows, so reads always go through the direct execute/cursor
+    path; ``hook`` is accepted for signature symmetry and ignored for the
+    fetch (no dialect sets ``raw_execute_hook`` today).
+    """
+    execute = getattr(handle, "execute", None)
+    if callable(execute):
+        result = execute(sql)
+        fetchone = getattr(result, "fetchone", None)
+        if callable(fetchone):
+            row = fetchone()
+            return row[0] if row else None
+        # A handle with a callable .execute has already run the SQL once;
+        # falling through to the cursor path would re-execute it. If the
+        # result has no fetchone, there is nothing more to try.
+        return None
+    cur = handle.cursor()
+    try:
+        cur.execute(sql)
+        row = cur.fetchone()
+        return row[0] if row else None
+    finally:
+        close = getattr(cur, "close", None)
+        if callable(close):
+            close()

src/mountainash_data/backends/ibis/_transaction.py

@@ -0,0 +1,134 @@
+"""Reentrant, cross-dialect unit-of-work machinery (Gap 3).
+
+Ambient registry keyed on id(raw_handle) under a module lock: the outermost
+transaction() issues the dialect's begin statement, nested calls join it, the
+outermost COMMITs, and any exception (or a poisoned-by-caught-nested-failure
+state) ROLLBACKs the whole unit. Flat semantics — no savepoints. Never toggles
+the driver's autocommit flag. BEGIN/COMMIT/ROLLBACK go through the shared
+`_raw.raw_execute` transport (honouring `raw_execute_hook`) because .execute()
+is not uniform across DBAPI drivers.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import threading
+import typing as t
+from dataclasses import dataclass
+
+from mountainash_data.backends.ibis._raw import raw_execute
+from mountainash_data.backends.ibis.dialects._registry import TransactionSupport
+from mountainash_data.core._warn import warn_once
+from mountainash_data.core.errors import (
+    TransactionUnsupportedError,
+    TransactionPoisonedError,
+    TransactionIntegrityError,
+)
+
+
+@dataclass
+class _TxState:
+    depth: int = 0
+    poisoned: bool = False
+
+
+_ACTIVE: dict[int, _TxState] = {}
+_LOCK = threading.Lock()
+
+
+@contextlib.contextmanager
+def run_transaction(
+    raw_handle: t.Any,
+    *,
+    support: TransactionSupport,
+    begin_statement: t.Optional[str],
+    dialect: str,
+    required: bool,
+    autocommit_probe: t.Optional[t.Callable[[t.Any], t.Optional[bool]]] = None,
+    in_transaction_probe: t.Optional[t.Callable[[t.Any], t.Optional[bool]]] = None,
+    raw_execute_hook: t.Optional[t.Callable[[t.Any, str], None]] = None,
+) -> t.Iterator[None]:
+    if support is TransactionSupport.NONE:
+        if required:
+            raise TransactionUnsupportedError(
+                f"{dialect!r} has no transaction concept; call transaction("
+                f"required=False) to run as a best-effort no-op."
+            )
+        warn_once(dialect, f"{dialect!r} has no transaction support; transaction() is a no-op.")
+        yield
+        return
+
+    def _exec(sql: str) -> None:
+        raw_execute(raw_handle, sql, hook=raw_execute_hook)
+
+    key = id(raw_handle)
+    with _LOCK:
+        state = _ACTIVE.get(key)
+        is_outer = state is None
+
+    # NOTE: the outer-entry check above and the _ACTIVE[key] insert below are
+    # deliberately two separate critical sections, not one — the BEGIN must run
+    # between them (register-after-BEGIN, so a failed BEGIN leaves no stale
+    # entry). This is safe because a single raw driver connection is not safe
+    # for concurrent use across threads at the DBAPI level, so two threads
+    # racing to open the outer transaction on ONE handle is already
+    # unsupported; the registry's job is reentrancy for sequential/nested
+    # reuse of one connection, not cross-thread arbitration.
+    if is_outer:
+        # Entry precondition (finding 1): ibis interleaves commits on autocommit-off
+        # connections, so a transaction() that cannot guarantee atomicity refuses.
+        if autocommit_probe is not None and autocommit_probe(raw_handle) is False:
+            raise TransactionIntegrityError(
+                f"{dialect!r} connection has autocommit disabled; ibis would interleave "
+                f"commits inside transaction(). Enable autocommit on the driver."
+            )
+        # Register AFTER a successful BEGIN so a failed BEGIN leaves no stale entry.
+        if begin_statement is not None:
+            _exec(begin_statement)
+        state = _TxState(depth=1)
+        with _LOCK:
+            _ACTIVE[key] = state
+        try:
+            yield
+        except BaseException as original:
+            try:
+                _exec("ROLLBACK")
+            except Exception as rollback_error:
+                original.__context__ = rollback_error
+            raise
+        else:
+            if state.poisoned:
+                _exec("ROLLBACK")
+                raise TransactionPoisonedError(
+                    "unit of work was poisoned by a caught nested failure; rolled back"
+                )
+            # Commit-time integrity (finding 1): if ibis rolled the server tx back
+            # underneath us, refuse rather than commit nothing.
+            if in_transaction_probe is not None and in_transaction_probe(raw_handle) is False:
+                raise TransactionIntegrityError(
+                    "server transaction vanished before COMMIT (ibis interleaved a "
+                    "commit/rollback inside the unit of work)"
+                )
+            _exec("COMMIT")
+        finally:
+            with _LOCK:
+                _ACTIVE.pop(key, None)
+        return
+
+    # Nested: join the in-flight unit of work (all state mutations under the lock).
+    assert state is not None  # is_outer is False here, so _ACTIVE.get(key) was not None
+    with _LOCK:
+        if state.poisoned:
+            raise TransactionPoisonedError(
+                "transaction is poisoned by a prior failure in this unit of work"
+            )
+        state.depth += 1
+    try:
+        yield
+    except BaseException:
+        with _LOCK:
+            state.poisoned = True
+        raise
+    finally:
+        with _LOCK:
+            state.depth -= 1