diff --git a/src/mountainash_data/backends/ibis/_adoption.py b/src/mountainash_data/backends/ibis/_adoption.py new file mode 100644 index 0000000..a6d01e4 --- /dev/null +++ b/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", + ) diff --git a/src/mountainash_data/backends/ibis/_raw.py b/src/mountainash_data/backends/ibis/_raw.py new file mode 100644 index 0000000..488ccba --- /dev/null +++ b/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() diff --git a/src/mountainash_data/backends/ibis/_transaction.py b/src/mountainash_data/backends/ibis/_transaction.py new file mode 100644 index 0000000..fa48544 --- /dev/null +++ b/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 diff --git a/src/mountainash_data/backends/ibis/backend.py b/src/mountainash_data/backends/ibis/backend.py index 43b3edf..4ec51e7 100644 --- a/src/mountainash_data/backends/ibis/backend.py +++ b/src/mountainash_data/backends/ibis/backend.py @@ -10,7 +10,11 @@ import typing as t -from mountainash_data.backends.ibis.dialects._registry import DIALECTS, DialectSpec +from mountainash_data.backends.ibis.dialects._registry import DIALECTS, DialectSpec, TransactionSupport +from mountainash_data.backends.ibis._transaction import run_transaction +from mountainash_data.backends.ibis._adoption import ( + apply_options, snapshot_options, restore_options, +) from mountainash_data.backends.ibis.operations import _generic_add_columns, _generic_rename_table, _generic_upsert from mountainash_data.backends.ibis._index import ( _generic_create_index, @@ -264,6 +268,7 @@ def from_ibis_connection( *, dialect: str, owns_connection: bool = False, + apply_session_options: t.Optional[dict[str, t.Any]] = None, ) -> IbisBackend: """Adopt an existing live ibis connection. @@ -273,8 +278,67 @@ def from_ibis_connection( 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. + + apply_session_options re-applies a caller-declared end-state for the + session options ibis mutates on adoption (the ibis backend is already + built, so the pre-adoption value cannot be snapshotted here — use + from_raw_connection(preserve_session=True) for faithful restore). + """ + backend = cls(dialect=dialect) + backend._conn = IbisConnection( + ibis_conn, backend._spec, owns_connection=owns_connection + ) + if apply_session_options: + apply_options( + backend.raw_driver_connection(), + backend._spec.adoption_mutations, + apply_session_options, + ) + return backend + + @classmethod + def from_raw_connection( + cls, + raw_conn: t.Any, + *, + dialect: str, + owns_connection: bool = False, + preserve_session: bool = False, + ) -> IbisBackend: + """Adopt a *raw driver* connection (not an ibis backend). + + This constructor owns the raw->ibis adoption step, so when + preserve_session=True it snapshots the session options ibis mutates on + adoption BEFORE calling ibis's from_connection, then restores them — + leaving the caller's session uncorrupted. preserve_session=False (the + default) reproduces plain ibis adoption behaviour. """ + import importlib + backend = cls(dialect=dialect) + # Gate (fable finding 4): only verified dialects have a known-good raw + # adoption path; others must use from_ibis_connection. + if not backend._spec.raw_adoption_verified: + raise NotImplementedError( + f"raw adoption not yet verified for {dialect!r}; construct the ibis " + f"connection yourself and use IbisBackend.from_ibis_connection(...)." + ) + options = backend._spec.adoption_mutations + snapshot = snapshot_options(raw_conn, options) if preserve_session else {} + + # ibis's from_connection runs _post_connect, which mutates the session + # BEFORE returning. If adoption raises after that, the caller's session is + # already stomped — restore in the finally so a failed adoption does not + # leave the session corrupted (Codex review). + ibis_backend_module = importlib.import_module( + f"ibis.backends.{backend._spec.ibis_backend_name}" + ) + try: + ibis_conn = ibis_backend_module.Backend.from_connection(raw_conn) + finally: + if preserve_session and snapshot: + restore_options(raw_conn, options, snapshot) + backend._conn = IbisConnection( ibis_conn, backend._spec, owns_connection=owns_connection ) @@ -468,6 +532,41 @@ def get_connection(self) -> IbisConnection: """Return the internal IbisConnection wrapper.""" return self._require_connected() + def raw_driver_connection(self) -> t.Any: + """Return the underlying native driver handle (see Backend protocol). + + Reads the per-dialect ``raw_handle_attr`` off the ibis backend. Works + for connections this backend opened AND adopted ones. Raises if not + connected or the handle is absent. + """ + conn = self._require_connected() + attr = self._spec.raw_handle_attr + handle = getattr(conn._ibis_conn, attr, None) + if handle is None: + raise RuntimeError( + f"No native driver handle on the {self.dialect!r} ibis backend " + f"(expected attribute {attr!r}); the connection may be closed." + ) + return handle + + @property + def supports_transactions(self) -> bool: + return self._spec.transaction_support is not TransactionSupport.NONE + + def transaction(self, *, required: bool = True) -> t.ContextManager[None]: + support = self._spec.transaction_support + raw = self.raw_driver_connection() if support is not TransactionSupport.NONE else None + return run_transaction( + raw, + support=support, + begin_statement=self._spec.begin_statement, + dialect=self.dialect, + required=required, + autocommit_probe=self._spec.autocommit_probe, + in_transaction_probe=self._spec.in_transaction_probe, + raw_execute_hook=self._spec.raw_execute_hook, + ) + # --- Inspection (terminal — delegates to IbisConnection) --- def list_tables(self, namespace: NamespaceLike = None) -> list[str]: diff --git a/src/mountainash_data/backends/ibis/dialects/_registry.py b/src/mountainash_data/backends/ibis/dialects/_registry.py index eda0d7a..fc0c8ff 100644 --- a/src/mountainash_data/backends/ibis/dialects/_registry.py +++ b/src/mountainash_data/backends/ibis/dialects/_registry.py @@ -28,6 +28,12 @@ class UpsertStyle(str, enum.Enum): ON_DUPLICATE_KEY = "on_duplicate_key" +class TransactionSupport(str, enum.Enum): + FULL = "full" + LIMITED = "limited" + NONE = "none" + + class DropScope(str, enum.Enum): SCHEMA_GLOBAL = "schema_global" # DROP INDEX name TABLE_SCOPED = "table_scoped" # DROP INDEX name ON tbl @@ -59,6 +65,19 @@ class IndexCapability: AddColumnsHook = t.Callable[..., None] +@dataclass(frozen=True) +class SessionOption: + """A session option ibis mutates on adoption (Gap 1). + + read_sql returns the current scalar value (None if unreadable); render_set + maps a value to the SQL statement that sets it. + """ + + name: str + read_sql: t.Optional[str] + render_set: t.Callable[[t.Any], str] + + @dataclass(frozen=True) class DialectSpec: """Per-dialect configuration and capability hooks.""" @@ -78,6 +97,19 @@ class DialectSpec: drop_index_hook: t.Optional[DropIndexHook] = None rename_table_hook: t.Optional[RenameTableHook] = None add_columns_hook: t.Optional[AddColumnsHook] = None + raw_handle_attr: str = "con" + # attribute on the ibis backend holding the native driver handle (Gap 2). + raw_adoption_verified: bool = False + # True once Gap 1's from_ibis_connection() adoption path has been live-verified + # for this dialect (assigned by the Gap 1 plan; declared here to avoid a + # second addition to this dataclass if Gap 1 lands after Gap 3). + transaction_support: "TransactionSupport" = TransactionSupport.NONE + begin_statement: t.Optional[str] = "BEGIN" + 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 + adoption_mutations: tuple["SessionOption", ...] = () + # session options ibis stomps on adoption; () = none (Gap 1). extras: t.Mapping[str, t.Any] = field(default_factory=dict) @@ -679,6 +711,45 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: ) +def _postgres_autocommit_probe(con: t.Any) -> t.Optional[bool]: + """psycopg Connection.autocommit — True when ibis's connect default is in force.""" + return bool(con.autocommit) + + +def _postgres_in_transaction_probe(con: t.Any) -> t.Optional[bool]: + """False when no server-side transaction is open (psycopg transaction_status IDLE == 0).""" + return con.info.transaction_status != 0 + + +def _sql_str_literal(v: t.Any) -> str: + """Escape a value as a single-quoted SQL string literal (injection-safe). + + Uses sqlglot so embedded quotes/backslashes are escaped, not raw-interpolated + (Codex review — apply_session_options values are caller-supplied). + """ + import sqlglot.expressions as exp + return exp.Literal.string(str(v)).sql() + + +def _duckdb_render_replacements(v: t.Any) -> str: + # boolean -> fixed token, never interpolated + return f"SET python_enable_replacements={'true' if v else 'false'}" + + +def _duckdb_render_timezone(v: t.Any) -> str: + return f"SET TimeZone={_sql_str_literal(v)}" + + +_DUCKDB_ADOPTION = ( + SessionOption("python_enable_replacements", + "SELECT current_setting('python_enable_replacements')", + _duckdb_render_replacements), + SessionOption("timezone", + "SELECT current_setting('TimeZone')", + _duckdb_render_timezone), +) + + DIALECTS: dict[str, DialectSpec] = { "sqlite": DialectSpec( ibis_backend_name="sqlite", @@ -693,6 +764,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=True, native_if_exists=True, index_types=frozenset(), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", ), "duckdb": DialectSpec( ibis_backend_name="duckdb", @@ -707,6 +780,10 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=True, native_if_exists=True, index_types=frozenset(), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", + adoption_mutations=_DUCKDB_ADOPTION, + raw_adoption_verified=True, ), "motherduck": DialectSpec( ibis_backend_name="duckdb", @@ -721,6 +798,10 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=True, native_if_exists=True, index_types=frozenset(), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", + adoption_mutations=_DUCKDB_ADOPTION, + raw_adoption_verified=True, ), "postgres": DialectSpec( ibis_backend_name="postgres", @@ -734,6 +815,10 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=True, native_if_exists=True, index_types=frozenset({"btree", "hash", "gist", "gin", "brin", "spgist"}), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", + autocommit_probe=_postgres_autocommit_probe, + in_transaction_probe=_postgres_in_transaction_probe, ), "mysql": DialectSpec( ibis_backend_name="mysql", @@ -747,6 +832,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=False, native_if_exists=False, index_types=frozenset({"btree"}), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", ), "mssql": DialectSpec( ibis_backend_name="mssql", @@ -760,6 +847,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=False, native_if_exists=True, index_types=frozenset(), ), + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN TRANSACTION", ), "oracle": DialectSpec( ibis_backend_name="oracle", @@ -773,6 +862,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=False, native_if_exists=False, index_types=frozenset(), ), + transaction_support=TransactionSupport.FULL, + begin_statement=None, ), "snowflake": DialectSpec( ibis_backend_name="snowflake", @@ -780,6 +871,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="snowflake://", connection_builder=_build_snowflake_connection, upsert_style=UpsertStyle.MERGE, + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", ), "bigquery": DialectSpec( ibis_backend_name="bigquery", @@ -787,6 +880,9 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="bigquery://", connection_builder=_build_bigquery_connection, upsert_style=UpsertStyle.MERGE, + raw_handle_attr="client", + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), "redshift": DialectSpec( ibis_backend_name="postgres", # Redshift uses postgres protocol @@ -794,6 +890,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="postgres://", # confirmed: redshift uses postgres:// connection_builder=_build_redshift_connection, upsert_style=UpsertStyle.MERGE, + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", ), "trino": DialectSpec( ibis_backend_name="trino", @@ -801,12 +899,16 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="trino://", connection_builder=_build_trino_connection, upsert_style=UpsertStyle.MERGE, + transaction_support=TransactionSupport.LIMITED, + begin_statement="START TRANSACTION", ), "clickhouse": DialectSpec( ibis_backend_name="clickhouse", connection_mode=_KWARGS, connection_string_scheme="clickhouse://", connection_builder=_build_clickhouse_connection, + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), "databricks": DialectSpec( ibis_backend_name="databricks", @@ -814,6 +916,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="", connection_builder=_build_databricks_connection, upsert_style=UpsertStyle.MERGE, + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), "singlestoredb": DialectSpec( ibis_backend_name="singlestoredb", @@ -827,6 +931,8 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: native_if_not_exists=False, native_if_exists=False, index_types=frozenset({"btree", "hash"}), ), + transaction_support=TransactionSupport.LIMITED, + begin_statement="BEGIN", ), "exasol": DialectSpec( ibis_backend_name="exasol", @@ -834,18 +940,24 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="exasol://", connection_builder=_build_exasol_connection, upsert_style=UpsertStyle.MERGE, + transaction_support=TransactionSupport.FULL, + begin_statement=None, ), "impala": DialectSpec( ibis_backend_name="impala", connection_mode=_KWARGS, connection_string_scheme="impala://", connection_builder=_build_impala_connection, + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), "materialize": DialectSpec( ibis_backend_name="materialize", connection_mode=_KWARGS, connection_string_scheme="materialize://", connection_builder=_build_materialize_connection, + transaction_support=TransactionSupport.FULL, + begin_statement="BEGIN", ), "risingwave": DialectSpec( ibis_backend_name="risingwave", @@ -853,17 +965,24 @@ def _build_pyspark_connection(**config: t.Any) -> t.Any: connection_string_scheme="risingwave://", connection_builder=_build_risingwave_connection, upsert_style=UpsertStyle.ON_CONFLICT, + transaction_support=TransactionSupport.LIMITED, + begin_statement="BEGIN", ), "druid": DialectSpec( ibis_backend_name="druid", connection_mode=_KWARGS, connection_string_scheme="druid://", connection_builder=_build_druid_connection, + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), "pyspark": DialectSpec( ibis_backend_name="pyspark", connection_mode=_CONNECTION_STRING, connection_string_scheme="pyspark://", connection_builder=_build_pyspark_connection, + raw_handle_attr="_session", + transaction_support=TransactionSupport.NONE, + begin_statement=None, ), } diff --git a/src/mountainash_data/backends/iceberg/backend.py b/src/mountainash_data/backends/iceberg/backend.py index ea19e44..2504363 100644 --- a/src/mountainash_data/backends/iceberg/backend.py +++ b/src/mountainash_data/backends/iceberg/backend.py @@ -2,10 +2,13 @@ from __future__ import annotations +import contextlib import typing as t from mountainash_data.backends.iceberg.catalogs.rest import IcebergRestConnection from mountainash_data.backends.iceberg.connection import IcebergConnectionBase +from mountainash_data.core.errors import TransactionUnsupportedError +from mountainash_data.core._warn import warn_once _CATALOG_REGISTRY: dict[str, type[IcebergConnectionBase]] = { @@ -77,3 +80,26 @@ def inspect_namespace(self, name: str) -> t.Any: def inspect_catalog(self, catalog: str | None = None) -> t.Any: return self._require_connected().inspect_catalog(catalog=catalog) + + def raw_driver_connection(self) -> t.Any: + """Return the underlying pyiceberg Catalog (native handle).""" + return self._require_connected().catalog_backend + + @property + def supports_transactions(self) -> bool: + return False + + @contextlib.contextmanager + def transaction(self, *, required: bool = True): + """Iceberg has no connection-level cross-table transaction; declines. + + required=True raises; required=False warns ONCE and no-ops. (pyiceberg + offers table-scoped transactions — a future capability, not this one.) + """ + if required: + raise TransactionUnsupportedError( + "iceberg has no connection-level transaction; use table-scoped " + "pyiceberg transactions, or call transaction(required=False)." + ) + warn_once("iceberg", "iceberg has no transaction support; transaction() is a no-op.") + yield diff --git a/src/mountainash_data/core/_warn.py b/src/mountainash_data/core/_warn.py new file mode 100644 index 0000000..882d86a --- /dev/null +++ b/src/mountainash_data/core/_warn.py @@ -0,0 +1,23 @@ +"""Process-wide "warn at most once per key" helper (Gap 3, fable finding 6). + +Shared by the ibis transaction machinery and the iceberg backend so a no-op +transaction() on an unsupported backend warns once per dialect, not per call. +Lives in core/ so neither backend imports the other. +""" +from __future__ import annotations + +import threading +import warnings + +_WARNED: set[str] = set() +_LOCK = threading.Lock() + + +def warn_once(key: str, message: str) -> None: + """Emit ``message`` via ``warnings.warn`` the first time ``key`` is seen.""" + with _LOCK: + first = key not in _WARNED + if first: + _WARNED.add(key) + if first: + warnings.warn(message, stacklevel=3) diff --git a/src/mountainash_data/core/errors.py b/src/mountainash_data/core/errors.py new file mode 100644 index 0000000..e27eefd --- /dev/null +++ b/src/mountainash_data/core/errors.py @@ -0,0 +1,20 @@ +"""Shared backend exceptions.""" + +from __future__ import annotations + + +class TransactionError(RuntimeError): + """Base for transaction() failures.""" + + +class TransactionUnsupportedError(TransactionError): + """transaction() called on a backend with no transaction concept.""" + + +class TransactionPoisonedError(TransactionError): + """The unit of work was aborted by a caught nested failure; it cannot commit.""" + + +class TransactionIntegrityError(TransactionError): + """Atomicity cannot be guaranteed: the driver is autocommit-off at entry, or the + server-side transaction vanished (ibis interleaved a commit/rollback) before COMMIT.""" diff --git a/src/mountainash_data/core/protocol.py b/src/mountainash_data/core/protocol.py index 7b2659d..085f8a9 100644 --- a/src/mountainash_data/core/protocol.py +++ b/src/mountainash_data/core/protocol.py @@ -41,3 +41,32 @@ def inspect_table( def inspect_namespace(self, name: str) -> NamespaceInfo: ... def inspect_catalog(self, catalog: str | None = None) -> CatalogInfo: ... + + def raw_driver_connection(self) -> t.Any: + """Return the underlying native driver handle (escape hatch). + + For SQL backends this is a live PEP-249 / native connection + (duckdb.DuckDBPyConnection, psycopg.Connection, sqlite3.Connection, + ...) suitable for transactions, DDL, information_schema, and + driver-specific idioms. The handle *kind* varies by backend (DBAPI + connection / client object / session object) and is NOT guaranteed to + be DBAPI-conformant — callers must not assume DBAPI semantics without + first checking the concrete backend. Raises (never returns ``None`` as + a sentinel) if not connected or the backend exposes no driver handle. + """ + ... + + @property + def supports_transactions(self) -> bool: + """True if transaction() opens a real unit of work (transaction_support is not NONE).""" + ... + + def transaction(self, *, required: bool = True) -> t.ContextManager[None]: + """Reentrant unit of work. Outermost issues BEGIN, nested calls join, + outermost COMMITs, any exception ROLLBACKs the whole unit. required=True + raises TransactionUnsupportedError on a backend with no transaction + concept; required=False warns once and runs as a no-op. Statements run + through this backend/ibis participate only while the driver is autocommit + (an adopted autocommit-off connection is refused with + TransactionIntegrityError). See spec §5.1–5.3.""" + ... diff --git a/tests/test_integration/test_transaction_live.py b/tests/test_integration/test_transaction_live.py new file mode 100644 index 0000000..f435087 --- /dev/null +++ b/tests/test_integration/test_transaction_live.py @@ -0,0 +1,29 @@ +import os +import pytest +from mountainash_data import IbisBackend + +pytestmark = pytest.mark.integration + +REQUIRE = os.environ.get("MOUNTAINASH_REQUIRE_LIVE_DB") == "1" +PG_URL = os.environ.get("MOUNTAINASH_TEST_POSTGRES_URL") + + +def _skip_or_fail(reason): + if REQUIRE: + pytest.fail(reason) + pytest.skip(reason) + + +def test_postgres_transaction_rollback(): + if not PG_URL: + _skip_or_fail("MOUNTAINASH_TEST_POSTGRES_URL not set") + with IbisBackend(PG_URL) as be: + raw = be.raw_driver_connection() + cur = raw.cursor() + cur.execute("CREATE TEMP TABLE t_tx (x INT)") + with pytest.raises(ValueError): + with be.transaction(): + cur.execute("INSERT INTO t_tx VALUES (1)") + raise ValueError("boom") + cur.execute("SELECT count(*) FROM t_tx") + assert cur.fetchone()[0] == 0 diff --git a/tests/test_unit/backends/ibis/test_adoption.py b/tests/test_unit/backends/ibis/test_adoption.py new file mode 100644 index 0000000..61eb88f --- /dev/null +++ b/tests/test_unit/backends/ibis/test_adoption.py @@ -0,0 +1,78 @@ +from mountainash_data.backends.ibis._adoption import ( + snapshot_options, restore_options, apply_options, +) +from mountainash_data.backends.ibis.dialects._registry import SessionOption + + +class FakeResult: + def __init__(self, value): + self._value = value + def fetchone(self): + return (self._value,) + + +class FakeHandle: + def __init__(self, values=None): + self.values = values or {} + self.calls = [] + def execute(self, sql): + self.calls.append(sql) + # read SQL returns a canned value keyed by substring match + for k, v in self.values.items(): + if k in sql: + return FakeResult(v) + return FakeResult(None) + + +OPT = SessionOption( + "python_enable_replacements", + "SELECT current_setting('python_enable_replacements')", + lambda v: f"SET python_enable_replacements={'true' if v else 'false'}", +) + + +def test_snapshot_reads_values(): + h = FakeHandle({"python_enable_replacements": True}) + snap = snapshot_options(h, (OPT,)) + assert snap == {"python_enable_replacements": True} + + +def test_restore_replays_captured(): + h = FakeHandle() + restore_options(h, (OPT,), {"python_enable_replacements": True}) + assert "SET python_enable_replacements=true" in h.calls + + +def test_apply_renders_declared_values(): + h = FakeHandle() + apply_options(h, (OPT,), {"python_enable_replacements": True}) + assert "SET python_enable_replacements=true" in h.calls + + +def test_apply_ignores_unknown_option_names(): + import warnings + + h = FakeHandle() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + apply_options(h, (OPT,), {"not_a_real_option": 1}) + assert h.calls == [] # nothing rendered for unknown names + assert any("not_a_real_option" in str(w.message) for w in caught) + + +def test_snapshot_skips_options_without_read_sql(): + opt = SessionOption("x", None, lambda v: f"SET x={v}") + h = FakeHandle() + assert snapshot_options(h, (opt,)) == {} + + +def test_snapshot_warns_when_read_raises(): + import warnings as _w + class Boom: + def execute(self, sql): + raise RuntimeError("cannot read setting") + with _w.catch_warnings(record=True) as rec: + _w.simplefilter("always") + snap = snapshot_options(Boom(), (OPT,)) + assert snap == {} + assert any("python_enable_replacements" in str(x.message) for x in rec) diff --git a/tests/test_unit/backends/ibis/test_backend.py b/tests/test_unit/backends/ibis/test_backend.py index 9b9ddf2..e36f0d5 100644 --- a/tests/test_unit/backends/ibis/test_backend.py +++ b/tests/test_unit/backends/ibis/test_backend.py @@ -1,5 +1,7 @@ """Tests for IbisBackend factory.""" +import duckdb +import sqlite3 import pytest import polars as pl @@ -215,6 +217,36 @@ def test_get_connection_accessor(): assert isinstance(conn, IbisConnection) +def test_transaction_none_dialect_required_false_noops_without_connection(): + # required=False must no-op even when the NONE backend is never connected + import warnings + be = IbisBackend(dialect="clickhouse") # NONE support, not connected + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with be.transaction(required=False): + pass # must not raise RuntimeError("not connected") + + +def test_raw_driver_connection_duckdb_returns_native_handle(): + with IbisBackend(dialect="duckdb", database=":memory:") as be: + raw = be.raw_driver_connection() + assert isinstance(raw, duckdb.DuckDBPyConnection) + # usable as a real handle + assert raw.execute("SELECT 1").fetchone()[0] == 1 + + +def test_raw_driver_connection_sqlite_returns_native_handle(): + with IbisBackend(dialect="sqlite", database=":memory:") as be: + raw = be.raw_driver_connection() + assert isinstance(raw, sqlite3.Connection) + + +def test_raw_driver_connection_requires_connected(): + be = IbisBackend(dialect="duckdb", database=":memory:") + with pytest.raises(RuntimeError, match="not connected"): + be.raw_driver_connection() + + # --------------------------------------------------------------------------- # DialectSpec hooks # --------------------------------------------------------------------------- @@ -485,3 +517,58 @@ def test_list_namespaces_accepts_catalog_kwarg(): with IbisBackend(dialect="duckdb", database=":memory:") as backend: # catalog=None is the default; the kwarg must be accepted without error. assert isinstance(backend.list_namespaces(catalog=None), list) + + +# --------------------------------------------------------------------------- +# transaction() / supports_transactions (Gap 3 Task 3) +# --------------------------------------------------------------------------- + +def test_supports_transactions_introspection(): + assert IbisBackend(dialect="duckdb", database=":memory:").supports_transactions is True + assert IbisBackend(dialect="clickhouse").supports_transactions is False + + +def test_transaction_ibis_level_op_rolls_back(tmp_path): + # the consumer's REAL shape: an ibis-level op (create_table) inside transaction() + db = str(tmp_path / "t.db") + with IbisBackend(dialect="duckdb", database=db) as be: + import pandas as pd + with pytest.raises(ValueError): + with be.transaction(): + be.create_table("t", pd.DataFrame({"x": [1]})) + raise ValueError("boom") + assert "t" not in be.list_tables() # rolled back + + +def test_transaction_commits(tmp_path): + db = str(tmp_path / "t.db") + with IbisBackend(dialect="duckdb", database=db) as be: + raw = be.raw_driver_connection() + raw.execute("CREATE TABLE t (x INT)") + with be.transaction(): + raw.execute("INSERT INTO t VALUES (1)") + assert raw.execute("SELECT count(*) FROM t").fetchone()[0] == 1 + + +def test_transaction_rolls_back(tmp_path): + db = str(tmp_path / "t.db") + with IbisBackend(dialect="duckdb", database=db) as be: + raw = be.raw_driver_connection() + raw.execute("CREATE TABLE t (x INT)") + with pytest.raises(ValueError): + with be.transaction(): + raw.execute("INSERT INTO t VALUES (1)") + raise ValueError("boom") + assert raw.execute("SELECT count(*) FROM t").fetchone()[0] == 0 + + +def test_transaction_nested_joins(tmp_path): + db = str(tmp_path / "t.db") + with IbisBackend(dialect="duckdb", database=db) as be: + raw = be.raw_driver_connection() + raw.execute("CREATE TABLE t (x INT)") + # nested MUST NOT raise "cannot start a transaction within a transaction" + with be.transaction(): + with be.transaction(): + raw.execute("INSERT INTO t VALUES (1)") + assert raw.execute("SELECT count(*) FROM t").fetchone()[0] == 1 diff --git a/tests/test_unit/backends/ibis/test_backend_adopt.py b/tests/test_unit/backends/ibis/test_backend_adopt.py index da85407..b5352ab 100644 --- a/tests/test_unit/backends/ibis/test_backend_adopt.py +++ b/tests/test_unit/backends/ibis/test_backend_adopt.py @@ -113,3 +113,56 @@ def test_unknown_dialect_raises(): IbisBackend.from_ibis_connection( ibis.duckdb.from_connection(raw), dialect="not-a-dialect" ) + + +def test_from_raw_connection_preserves_python_enable_replacements(): + raw = duckdb.connect() + # caller's session default is True + before = raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] + assert before is True + be = IbisBackend.from_raw_connection(raw, dialect="duckdb", preserve_session=True) + after = raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] + assert after is True # restored despite ibis stomping it to False during adoption + be.close() + + +def test_from_raw_connection_without_preserve_leaves_ibis_default(): + raw = duckdb.connect() + be = IbisBackend.from_raw_connection(raw, dialect="duckdb", preserve_session=False) + after = raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] + assert after is False # ibis stomped it; we did not restore + be.close() + + +def test_from_raw_connection_returns_working_backend(): + raw = duckdb.connect() + be = IbisBackend.from_raw_connection(raw, dialect="duckdb") + assert be.raw_driver_connection() is raw + be.close() + + +def test_from_raw_connection_gated_on_unverified_dialect(): + # postgres has raw_adoption_verified=False -> clear error, not a cryptic ibis failure + with pytest.raises(NotImplementedError, match="raw adoption not yet verified"): + IbisBackend.from_raw_connection(object(), dialect="postgres") + + +def test_apply_session_options_reenables_replacements(): + raw = duckdb.connect() + adopted = ibis.duckdb.from_connection(raw) # ibis stomps replacements to False + assert raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] is False + be = IbisBackend.from_ibis_connection( + adopted, dialect="duckdb", + apply_session_options={"python_enable_replacements": True}, + ) + assert raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] is True + be.close() + + +def test_from_ibis_connection_default_unchanged(): + raw = duckdb.connect() + adopted = ibis.duckdb.from_connection(raw) + be = IbisBackend.from_ibis_connection(adopted, dialect="duckdb") + # no apply -> ibis default stands + assert raw.execute("SELECT current_setting('python_enable_replacements')").fetchone()[0] is False + be.close() diff --git a/tests/test_unit/backends/ibis/test_dialect_spec.py b/tests/test_unit/backends/ibis/test_dialect_spec.py index b9f3335..d51eaf6 100644 --- a/tests/test_unit/backends/ibis/test_dialect_spec.py +++ b/tests/test_unit/backends/ibis/test_dialect_spec.py @@ -4,6 +4,8 @@ from mountainash_data.backends.ibis.dialects._registry import ( DialectSpec, DIALECTS, + SessionOption, + TransactionSupport, ) @@ -47,3 +49,74 @@ def test_registry_entries_are_dialect_specs(): for name, spec in DIALECTS.items(): assert isinstance(spec, DialectSpec), f"{name} entry is not a DialectSpec" assert spec.ibis_backend_name, f"{name} missing ibis_backend_name" + + +def test_raw_handle_attr_defaults_to_con(): + assert DIALECTS["duckdb"].raw_handle_attr == "con" + assert DIALECTS["postgres"].raw_handle_attr == "con" + assert DIALECTS["sqlite"].raw_handle_attr == "con" + + +def test_raw_handle_attr_overrides(): + assert DIALECTS["bigquery"].raw_handle_attr == "client" + assert DIALECTS["pyspark"].raw_handle_attr == "_session" + + +def test_every_dialect_has_raw_handle_attr(): + for name, spec in DIALECTS.items(): + assert isinstance(spec.raw_handle_attr, str) and spec.raw_handle_attr, \ + f"{name} missing or empty raw_handle_attr" + + +def test_transaction_support_assignments(): + assert DIALECTS["duckdb"].transaction_support is TransactionSupport.FULL + assert DIALECTS["mssql"].transaction_support is TransactionSupport.FULL + assert DIALECTS["clickhouse"].transaction_support is TransactionSupport.NONE + assert DIALECTS["trino"].transaction_support is TransactionSupport.LIMITED + + +def test_begin_statement_assignments(): + assert DIALECTS["duckdb"].begin_statement == "BEGIN" + assert DIALECTS["mssql"].begin_statement == "BEGIN TRANSACTION" + assert DIALECTS["oracle"].begin_statement is None + assert DIALECTS["exasol"].begin_statement is None + + +def test_none_support_implies_no_begin_statement(): + for name, spec in DIALECTS.items(): + if spec.transaction_support is TransactionSupport.NONE: + assert spec.begin_statement is None, name + + +def test_duckdb_adoption_mutations_declared(): + names = {o.name for o in DIALECTS["duckdb"].adoption_mutations} + assert "python_enable_replacements" in names + assert "timezone" in names + + +def test_non_mutating_dialects_empty(): + for d in ("trino", "clickhouse", "druid", "bigquery"): + assert DIALECTS[d].adoption_mutations == () + + +def test_session_options_well_formed(): + for name, spec in DIALECTS.items(): + for opt in spec.adoption_mutations: + assert isinstance(opt, SessionOption) + assert opt.name + # render_set must produce a str statement + assert isinstance(opt.render_set(True), str) + + +def test_duckdb_timezone_render_is_injection_safe(): + tz_opt = next(o for o in DIALECTS["duckdb"].adoption_mutations if o.name == "timezone") + rendered = tz_opt.render_set("UTC'; DROP TABLE t; --") + # the malicious quote must be escaped inside the literal, not break out of it + assert "DROP TABLE" in rendered # value preserved as data + assert rendered.count("SET TimeZone=") == 1 + assert not rendered.rstrip().endswith("--") # not left as trailing raw SQL + + +def test_raw_adoption_verified_assignments(): + assert DIALECTS["duckdb"].raw_adoption_verified is True + assert DIALECTS["postgres"].raw_adoption_verified is False diff --git a/tests/test_unit/backends/ibis/test_raw.py b/tests/test_unit/backends/ibis/test_raw.py new file mode 100644 index 0000000..5104bb9 --- /dev/null +++ b/tests/test_unit/backends/ibis/test_raw.py @@ -0,0 +1,88 @@ +from mountainash_data.backends.ibis._raw import raw_execute, raw_fetch_scalar + + +class FakeHandleWithExecute: + def __init__(self, result=None): + self.calls = [] + self._result = result + + def execute(self, sql): + self.calls.append(sql) + return self._result + + +class RecordingCursor: + def __init__(self, log, fetch_result=None): + self.log = log + self._fetch_result = fetch_result + self.closed = False + + def execute(self, sql): + self.log.append(("cur", sql)) + + def fetchone(self): + return self._fetch_result + + def close(self): + self.closed = True + self.log.append(("close", None)) + + +class FakeHandleNoExecute: + def __init__(self, fetch_result=None): + self.log = [] + self._fetch_result = fetch_result + self.cursor_obj = None + + def cursor(self): + self.cursor_obj = RecordingCursor(self.log, self._fetch_result) + return self.cursor_obj + + +class FakeResult: + def __init__(self, row): + self._row = row + + def fetchone(self): + return self._row + + +def test_raw_execute_direct_path_uses_handle_execute(): + h = FakeHandleWithExecute() + raw_execute(h, "SELECT 1") + assert h.calls == ["SELECT 1"] + + +def test_raw_execute_cursor_path_when_no_execute(): + h = FakeHandleNoExecute() + raw_execute(h, "SELECT 1") + assert ("cur", "SELECT 1") in h.log + assert ("close", None) in h.log + assert h.cursor_obj.closed is True + + +def test_raw_execute_hook_override_skips_handle_execute(): + calls = [] + h = FakeHandleWithExecute() + + def hook(handle, sql): + calls.append(sql) + + raw_execute(h, "SELECT 1", hook=hook) + assert calls == ["SELECT 1"] + assert h.calls == [] + + +def test_raw_fetch_scalar_direct_path_returns_scalar(): + h = FakeHandleWithExecute(result=FakeResult((42,))) + assert raw_fetch_scalar(h, "SELECT 42") == 42 + + +def test_raw_fetch_scalar_cursor_path_returns_scalar(): + h = FakeHandleNoExecute(fetch_result=(7,)) + assert raw_fetch_scalar(h, "SELECT 7") == 7 + + +def test_raw_fetch_scalar_empty_result_returns_none(): + h = FakeHandleWithExecute(result=FakeResult(None)) + assert raw_fetch_scalar(h, "SELECT NULL") is None diff --git a/tests/test_unit/backends/ibis/test_transaction.py b/tests/test_unit/backends/ibis/test_transaction.py new file mode 100644 index 0000000..7233ece --- /dev/null +++ b/tests/test_unit/backends/ibis/test_transaction.py @@ -0,0 +1,144 @@ +import warnings +import pytest +from mountainash_data.backends.ibis._transaction import run_transaction, _ACTIVE +from mountainash_data.backends.ibis.dialects._registry import TransactionSupport +from mountainash_data.core.errors import ( + TransactionUnsupportedError, TransactionPoisonedError, TransactionIntegrityError, +) + + +class FakeHandle: + def __init__(self): + self.calls = [] + def execute(self, sql): + self.calls.append(sql) + + +def _tx(h, **kw): + kw.setdefault("support", TransactionSupport.FULL) + kw.setdefault("begin_statement", "BEGIN") + kw.setdefault("dialect", "duckdb") + kw.setdefault("required", True) + return run_transaction(h, **kw) + + +def test_autocommit_off_entry_raises(): + h = FakeHandle() + with pytest.raises(TransactionIntegrityError): + with _tx(h, autocommit_probe=lambda _c: False): + pass + assert h.calls == [] # refused before BEGIN + assert id(h) not in _ACTIVE + + +def test_commit_time_integrity_probe_raises_if_tx_vanished(): + h = FakeHandle() + with pytest.raises(TransactionIntegrityError): + with _tx(h, in_transaction_probe=lambda _c: False): + pass + assert "COMMIT" not in h.calls # integrity failure instead of a false commit + assert id(h) not in _ACTIVE + + +def test_outermost_commit(): + h = FakeHandle() + with _tx(h): + pass + assert h.calls == ["BEGIN", "COMMIT"] + assert id(h) not in _ACTIVE + + +def test_exception_rolls_back(): + h = FakeHandle() + with pytest.raises(ValueError): + with _tx(h): + raise ValueError("boom") + assert h.calls == ["BEGIN", "ROLLBACK"] + assert id(h) not in _ACTIVE + + +def test_nested_joins_no_second_begin(): + h = FakeHandle() + with _tx(h): + with _tx(h): + pass + assert h.calls == ["BEGIN", "COMMIT"] # inner joined; only one BEGIN/COMMIT + + +def test_nested_exception_rolls_back_whole_unit(): + h = FakeHandle() + with pytest.raises(ValueError): + with _tx(h): + with _tx(h): + raise ValueError("boom") + assert h.calls == ["BEGIN", "ROLLBACK"] + + +def test_none_support_required_raises(): + h = FakeHandle() + with pytest.raises(TransactionUnsupportedError): + with _tx(h, support=TransactionSupport.NONE, begin_statement=None): + pass + assert h.calls == [] + + +def test_none_support_not_required_warns_once_and_noops(): + from mountainash_data.core import _warn as _warnmod; _warnmod._WARNED.discard("clickhouse") + h = FakeHandle() + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with _tx(h, support=TransactionSupport.NONE, begin_statement=None, + required=False, dialect="clickhouse"): + pass + assert h.calls == [] + assert any("clickhouse" in str(x.message) for x in w) + + +def test_begin_statement_none_skips_begin(): + h = FakeHandle() + with _tx(h, begin_statement=None, dialect="oracle"): + pass + assert h.calls == ["COMMIT"] # implicit begin; commit still issued + + +def test_begin_failure_leaves_no_registry_entry(): + class Boom(FakeHandle): + def execute(self, sql): + if sql == "BEGIN": + raise RuntimeError("begin failed") + super().execute(sql) + h = Boom() + with pytest.raises(RuntimeError, match="begin failed"): + with _tx(h): + pass + assert id(h) not in _ACTIVE # register-after-begin: no stale entry + + +def test_poison_via_caught_nested_exception_does_not_commit(): + # caller CATCHES the nested failure inside the outer block; outer must NOT commit + h = FakeHandle() + with pytest.raises(TransactionPoisonedError): + with _tx(h): + try: + with _tx(h): + raise ValueError("inner") + except ValueError: + pass # swallow — but the unit of work is poisoned + assert h.calls == ["BEGIN", "ROLLBACK"] # rolled back, never committed + assert id(h) not in _ACTIVE + + +def test_transport_uses_cursor_when_no_execute(): + # a DBAPI connection without .execute() must go through .cursor().execute() + class Cursor: + def __init__(self, log): self.log = log + def execute(self, sql): self.log.append(("cur", sql)) + def close(self): self.log.append(("close", None)) + class ConnNoExecute: + def __init__(self): self.log = [] + def cursor(self): return Cursor(self.log) + h = ConnNoExecute() + with _tx(h): + pass + assert ("cur", "BEGIN") in h.log and ("cur", "COMMIT") in h.log + assert ("close", None) in h.log diff --git a/tests/test_unit/backends/iceberg/test_backend.py b/tests/test_unit/backends/iceberg/test_backend.py index fdf9513..9ef7a17 100644 --- a/tests/test_unit/backends/iceberg/test_backend.py +++ b/tests/test_unit/backends/iceberg/test_backend.py @@ -9,12 +9,15 @@ test environment. All tests in this module are skipped when it is absent. """ +import warnings import pytest +from unittest.mock import MagicMock pyiceberg = pytest.importorskip("pyiceberg", reason="pyiceberg not installed") from mountainash_data.backends.iceberg.backend import IcebergBackend # noqa: E402 from mountainash_data.core.protocol import Backend # noqa: E402 +from mountainash_data.core.errors import TransactionUnsupportedError # noqa: E402 def test_iceberg_backend_satisfies_protocol(): @@ -31,3 +34,36 @@ def test_unknown_catalog_raises(): def test_iceberg_backend_carries_config(): backend = IcebergBackend(catalog="rest", uri="http://localhost:8181", token="abc") assert backend._config == {"uri": "http://localhost:8181", "token": "abc"} + + +def test_raw_driver_connection_returns_catalog(monkeypatch): + be = IcebergBackend(catalog="rest", uri="http://localhost:8181") + fake_conn = MagicMock() + fake_catalog = object() + fake_conn.catalog_backend = fake_catalog + be._conn = fake_conn # simulate connected + assert be.raw_driver_connection() is fake_catalog + + +def test_raw_driver_connection_requires_connected(): + be = IcebergBackend(catalog="rest", uri="http://localhost:8181") + with pytest.raises(RuntimeError, match="not connected"): + be.raw_driver_connection() + + +def test_transaction_required_raises(): + be = IcebergBackend(catalog="rest", uri="http://localhost:8181") + with pytest.raises(TransactionUnsupportedError): + with be.transaction(): + pass + + +def test_transaction_not_required_noops(): + be = IcebergBackend(catalog="rest", uri="http://localhost:8181") + from mountainash_data.core import _warn as _warnmod + _warnmod._WARNED.discard("iceberg") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with be.transaction(required=False): + pass + assert any("iceberg" in str(x.message).lower() for x in w) diff --git a/tests/test_unit/core/test_protocol.py b/tests/test_unit/core/test_protocol.py index 13f1f54..0848790 100644 --- a/tests/test_unit/core/test_protocol.py +++ b/tests/test_unit/core/test_protocol.py @@ -62,3 +62,19 @@ def test_connection_close_idempotent_marker(): conn = _FakeConnection() conn.close() assert conn.closed is True + + +def test_protocol_declares_raw_driver_connection(): + assert hasattr(Backend, "raw_driver_connection") + + +def test_iceberg_backend_satisfies_widened_protocol(): + import pytest + pytest.importorskip("pyiceberg") + from mountainash_data.backends.iceberg.backend import IcebergBackend + assert hasattr(IcebergBackend, "raw_driver_connection") + + +def test_protocol_declares_transaction(): + from mountainash_data.core.protocol import Backend + assert hasattr(Backend, "transaction") diff --git a/tests/test_unit/core/test_warn.py b/tests/test_unit/core/test_warn.py new file mode 100644 index 0000000..108eecd --- /dev/null +++ b/tests/test_unit/core/test_warn.py @@ -0,0 +1,22 @@ +import warnings +from mountainash_data.core import _warn + + +def test_warn_once_emits_first_time_only(): + _warn._WARNED.discard("k-alpha") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _warn.warn_once("k-alpha", "first") + _warn.warn_once("k-alpha", "second") + assert len(w) == 1 + assert "first" in str(w[0].message) + + +def test_warn_once_distinct_keys_each_warn(): + _warn._WARNED.discard("k-beta") + _warn._WARNED.discard("k-gamma") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _warn.warn_once("k-beta", "b") + _warn.warn_once("k-gamma", "g") + assert len(w) == 2