diff --git a/src/mountainash_data/backends/ibis/backend.py b/src/mountainash_data/backends/ibis/backend.py index 67a2433..43b3edf 100644 --- a/src/mountainash_data/backends/ibis/backend.py +++ b/src/mountainash_data/backends/ibis/backend.py @@ -82,10 +82,17 @@ class IbisConnection: Constructed by IbisBackend.connect() — not intended to be instantiated directly. """ - def __init__(self, ibis_conn: t.Any, dialect_spec: DialectSpec) -> None: + def __init__( + self, + ibis_conn: t.Any, + dialect_spec: DialectSpec, + *, + owns_connection: bool = True, + ) -> None: self._ibis_conn = ibis_conn self._dialect_spec = dialect_spec self._closed = False + self._owns_connection = owns_connection def list_namespaces(self, catalog: str | None = None) -> list[str]: """Return the names of all namespaces (schemas/databases) visible to this connection.""" @@ -167,10 +174,15 @@ def inspect_catalog(self, catalog: str | None = None) -> CatalogInfo: ) def close(self) -> None: - """Release the connection. Idempotent.""" + """Release the connection. Idempotent. + + When the connection was adopted (owns_connection=False), the + underlying ibis connection belongs to the caller and is left open; + only this wrapper is marked closed. + """ if not self._closed: try: - if hasattr(self._ibis_conn, "disconnect"): + if self._owns_connection and hasattr(self._ibis_conn, "disconnect"): self._ibis_conn.disconnect() except Exception: pass @@ -245,6 +257,29 @@ def __init__( "or a dialect= keyword is required" ) + @classmethod + def from_ibis_connection( + cls, + ibis_conn: t.Any, + *, + dialect: str, + owns_connection: bool = False, + ) -> IbisBackend: + """Adopt an existing live ibis connection. + + The adopted connection shares the caller's session and transaction + state — e.g. ``ibis.duckdb.from_connection(raw_duckdb_conn)`` wraps + the caller's own DuckDB connection, so ``BEGIN``/``COMMIT`` issued on + the raw connection bracket this backend's writes too. By default the + backend does NOT own the connection: ``close()`` releases the wrapper + but leaves the underlying connection open for the caller. + """ + backend = cls(dialect=dialect) + backend._conn = IbisConnection( + ibis_conn, backend._spec, owns_connection=owns_connection + ) + return backend + def _init_from_positional( self, value: str | t.Any, config: dict[str, t.Any] ) -> None: diff --git a/tests/test_unit/backends/ibis/test_backend_adopt.py b/tests/test_unit/backends/ibis/test_backend_adopt.py new file mode 100644 index 0000000..da85407 --- /dev/null +++ b/tests/test_unit/backends/ibis/test_backend_adopt.py @@ -0,0 +1,115 @@ +"""IbisBackend.from_ibis_connection — adopt an existing live ibis connection.""" + +from __future__ import annotations + +import duckdb +import ibis +import polars as pl +import pytest + +from mountainash_data.backends.ibis.backend import IbisBackend + + +@pytest.fixture +def raw_db(): + conn = duckdb.connect(":memory:") + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)") + yield conn + conn.close() + + +def _adopt(raw_db) -> IbisBackend: + return IbisBackend.from_ibis_connection( + ibis.duckdb.from_connection(raw_db), dialect="duckdb" + ) + + +def test_adopted_backend_writes_through_same_connection(raw_db): + backend = _adopt(raw_db) + backend.insert("t", pl.DataFrame({"id": [1], "name": ["a"]})) + # visible on the RAW connection — proof it is the same session + assert raw_db.execute("SELECT count(*) FROM t").fetchone()[0] == 1 + + +def test_adopted_backend_shares_transaction_state(raw_db): + backend = _adopt(raw_db) + raw_db.execute("BEGIN") + backend.insert("t", pl.DataFrame({"id": [1], "name": ["a"]})) + raw_db.execute("ROLLBACK") + assert raw_db.execute("SELECT count(*) FROM t").fetchone()[0] == 0 + + +def test_upsert_through_adopted_connection(raw_db): + backend = _adopt(raw_db) + backend.insert("t", pl.DataFrame({"id": [1], "name": ["a"]})) + backend.upsert( + "t", pl.DataFrame({"id": [1], "name": ["b"]}), + conflict_columns=["id"], update_columns=["name"], + ) + assert raw_db.execute("SELECT name FROM t WHERE id = 1").fetchone()[0] == "b" + + +def test_upsert_stages_full_rows_against_not_null_columns(raw_db): + """DuckDB validates NOT NULL on the INSERT path even when ON CONFLICT + resolves to UPDATE — consumers upserting a narrow update set must still + stage FULL rows. This pins both sides of that contract.""" + raw_db.execute( + "CREATE TABLE nn (k INTEGER PRIMARY KEY, flag BOOLEAN, req TEXT NOT NULL)" + ) + backend = _adopt(raw_db) + backend.insert("nn", pl.DataFrame({"k": [1], "flag": [True], "req": ["x"]})) + + # full-row frame + narrow update_columns: works, updates only flag + backend.upsert( + "nn", pl.DataFrame({"k": [1], "flag": [False], "req": ["IGNORED"]}), + conflict_columns=["k"], update_columns=["flag"], + ) + assert raw_db.execute("SELECT flag, req FROM nn").fetchone() == (False, "x") + + # partial-column frame: DuckDB rejects it before conflict resolution + with pytest.raises(Exception, match="NOT NULL"): + backend.upsert( + "nn", pl.DataFrame({"k": [1], "flag": [True]}), + conflict_columns=["k"], update_columns=["flag"], + ) + + +def test_close_leaves_unowned_connection_open(raw_db): + backend = _adopt(raw_db) # owns_connection defaults to False + backend.close() + # the caller's connection must survive + assert raw_db.execute("SELECT 1").fetchone()[0] == 1 + + +class _StubIbisConn: + """Counted disconnect — proves close() actually calls it (or doesn't).""" + + def __init__(self): + self.disconnects = 0 + + def disconnect(self): + self.disconnects += 1 + + +def test_close_disconnects_owned_connection(): + stub = _StubIbisConn() + backend = IbisBackend.from_ibis_connection( + stub, dialect="duckdb", owns_connection=True, + ) + backend.close() + assert stub.disconnects == 1 + + +def test_close_never_disconnects_unowned_connection(): + stub = _StubIbisConn() + backend = IbisBackend.from_ibis_connection(stub, dialect="duckdb") + backend.close() + assert stub.disconnects == 0 + + +def test_unknown_dialect_raises(): + raw = duckdb.connect(":memory:") + with pytest.raises(KeyError): + IbisBackend.from_ibis_connection( + ibis.duckdb.from_connection(raw), dialect="not-a-dialect" + )