fix: stop crashing and losing writes when a pool hook leaves a transaction open

src/zero_server/server/adapters/psycopg.py

@@ -66,6 +66,7 @@ def __init__(self, pool: "AsyncConnectionPool", isolation_level: IsolationLevel
 
     async def query(self, sql: str, params: Sequence[Any] = ()) -> list[Row]:
         async with self.pool.connection() as connection:
+            await _settle(connection)
             # In autocommit, so psycopg sends no `BEGIN` of its own.
             autocommit = connection.autocommit
             await connection.set_autocommit(True)
@@ -97,13 +98,25 @@ async def _transaction[R](
     fn: Callable[[DBTransaction[psycopg.AsyncConnection[Any]]], Awaitable[R]],
     isolation_level: IsolationLevel | None,
 ) -> R:
+    await _settle(connection)
     async with connection.transaction():
         tx = PsycopgTransaction(connection)
         if isolation_level is not None:
             await tx.query(set_isolation_statement(isolation_level))
         return await fn(tx)
 
 
+async def _settle(connection: psycopg.AsyncConnection[Any]) -> None:
+    """Commit a transaction the pool handed the connection out in, keeping what it set.
+
+    A pool `check` that runs a statement opens one, and the pool takes the connection as it is. Left open,
+    psycopg refuses to switch to autocommit, and `transaction()` nests a savepoint in it that the pool
+    rolls back with the rest when the connection returns.
+    """
+    if connection.info.transaction_status == pq.TransactionStatus.INTRANS:
+        await connection.commit()
+
+
 class PsycopgDatabase(ZQLDatabase[psycopg.AsyncConnection[Any]]):
     """A `ZQLDatabase` on a psycopg pool (`AsyncConnectionPool(conninfo, open=False)`, then `await pool.open()`)."""
 

src/zero_server/server/adapters/sqlalchemy.py

@@ -82,17 +82,21 @@ def __init__(self, engine: AsyncEngine, isolation_level: IsolationLevel | None =
 
     async def query(self, sql: str, params: Sequence[Any] = ()) -> list[Row]:
         async with self.engine.connect() as connection:
+            await _settle(connection)
             # Autocommit, so SQLAlchemy begins no transaction; the pool resets it when the connection returns.
             connection = await connection.execution_options(isolation_level="AUTOCOMMIT")
             return await (await sqlalchemy_transaction(connection)).query(sql, params)
 
     async def transaction[R](self, fn: Callable[[DBTransaction[AsyncConnection]], Awaitable[R]]) -> R:
-        async with self.engine.begin() as connection:
-            return await fn(await _begun(connection, self._isolation_level))
+        async with self.engine.connect() as connection:
+            await _settle(connection)
+            async with connection.begin():
+                return await fn(await _begun(connection, self._isolation_level))
 
     async def _transaction[R](
         self, connection: AsyncConnection, fn: Callable[[DBTransaction[AsyncConnection]], Awaitable[R]]
     ) -> R:
+        await _settle(connection)
         async with connection.begin():
             if connection.dialect.driver == "asyncpg":
                 await _send_begin(await connection.get_raw_connection())
@@ -103,6 +107,25 @@ def session(self) -> AbstractAsyncContextManager[SessionConnection[AsyncConnecti
         return session(self.engine.connect, _close, self._transaction)
 
 
+async def _settle(connection: AsyncConnection) -> None:
+    """Commit a transaction the pool handed the connection out in, keeping what it set.
+
+    A pool `checkout` hook that runs a statement through the driver opens one. Left open, switching to
+    autocommit fails on psycopg and rolls it back on asyncpg, and a transaction's isolation level can't be set.
+    """
+    driver_connection = (await connection.get_raw_connection()).driver_connection
+    assert driver_connection is not None
+    if connection.dialect.driver == "asyncpg":
+        in_transaction = driver_connection.is_in_transaction()
+    else:
+        from psycopg.pq import TransactionStatus
+
+        in_transaction = driver_connection.info.transaction_status == TransactionStatus.INTRANS
+    if in_transaction:
+        # SQLAlchemy sends no `BEGIN` of its own here, so this commits the transaction already open.
+        await (await connection.begin()).commit()
+
+
 async def _begun(connection: AsyncConnection, isolation_level: IsolationLevel | None) -> SQLAlchemyTransaction:
     tx = await sqlalchemy_transaction(connection)
     if isolation_level is not None:

tests/wire/conftest.py

@@ -265,11 +265,19 @@ async def query(self, sql: str, params: Any = ()) -> Any:
 
 
 async def recorded_run(
-    wire_db: WireDB, adapter: str, schema: Schema, *, bare: bool, failFirst: bool, concurrent: bool, runs: int
+    wire_db: WireDB,
+    adapter: str,
+    schema: Schema,
+    *,
+    bare: bool,
+    failFirst: bool,
+    concurrent: bool,
+    runs: int,
+    inner: Any = None,
 ) -> dict[str, Any]:
     from zero_server.query.builder import create_builder
 
-    inner = wire_db.py[adapter](schema).connection
+    inner = inner if inner is not None else wire_db.py[adapter](schema).connection
     connection = _RecordedBare(inner, failFirst) if bare else _Recorded(inner, failFirst)
     db = ZQLDatabase(connection, schema)
     query = create_builder(schema).item.where("value", "!=", "x").order_by("id", "asc")

tests/wire/test_mutate.py

@@ -3,6 +3,7 @@
 
 import contextlib
 from typing import Any
+from uuid import uuid4
 
 import pytest
 from hypothesis import example, given
@@ -532,3 +533,119 @@ async def _insert_items(engine: Any) -> None:
 
     async with engine.begin() as conn:
         await conn.execute(text("INSERT INTO item (id, value) VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')"))
+
+
+# What a pool hook runs on each connection it hands out. `SELECT 1` takes a snapshot, after which a
+# transaction's isolation level can no longer be set.
+CHECKOUT_STATEMENTS = ["SELECT 1", "SET statement_timeout = '4321ms'"]
+
+
+async def _sqlalchemy_checked_out_in_a_transaction(wire_db: WireDB, driver: str, level: Any) -> tuple[Any, Any]:
+    """An engine whose `checkout` hook runs statements through the driver, which opens a transaction."""
+    from sqlalchemy import event
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    from zero_server.server.adapters.sqlalchemy import SQLAlchemyConnection
+
+    engine = create_async_engine(wire_db.py_engine.url.set(drivername=f"postgresql+{driver}"), pool_size=1)
+
+    @event.listens_for(engine.sync_engine, "checkout")
+    def _checkout(dbapi_connection: Any, _record: Any, _proxy: Any) -> None:
+        cursor = dbapi_connection.cursor()
+        for statement in CHECKOUT_STATEMENTS:
+            cursor.execute(statement)
+        cursor.close()
+
+    return SQLAlchemyConnection(engine, level), engine.dispose
+
+
+async def _psycopg_checked_out_in_a_transaction(wire_db: WireDB, _driver: str, level: Any) -> tuple[Any, Any]:
+    """A pool whose `check` runs statements, which opens a transaction; the pool checks only that they ran."""
+    from psycopg_pool import AsyncConnectionPool
+
+    from zero_server.server.adapters.psycopg import PsycopgConnection
+
+    async def check(connection: Any) -> None:
+        for statement in CHECKOUT_STATEMENTS:
+            await connection.execute(statement)
+
+    url = wire_db.py_engine.url.set(drivername="postgresql").render_as_string(hide_password=False)
+    pool = AsyncConnectionPool(url, min_size=1, max_size=1, check=check, open=False)
+    await pool.open()
+    return PsycopgConnection(pool, level), pool.close
+
+
+CHECKED_OUT_IN_A_TRANSACTION = {
+    "sqlalchemy": (_sqlalchemy_checked_out_in_a_transaction, "asyncpg"),
+    "sqlalchemy-psycopg": (_sqlalchemy_checked_out_in_a_transaction, "psycopg"),
+    "psycopg": (_psycopg_checked_out_in_a_transaction, "psycopg"),
+}
+"""The adapters whose pool can hand a connection out in a transaction. asyncpg opens none implicitly."""
+
+
+@pytest.mark.parametrize("adapter", CHECKED_OUT_IN_A_TRANSACTION)
+def test_run_matches_upstream_on_a_connection_checked_out_in_a_transaction(
+    harness: Harness, wire_db: WireDB, adapter: str
+) -> None:
+    """A bare read on a connection a pool hook left in a transaction reads as upstream does, and keeps
+    what the hook set."""
+    from tests.wire.conftest import recorded_run
+
+    wire_db.reset()
+    for engine in (wire_db.ts_engine, wire_db.py_engine):
+        wire_db.run(_insert_items(engine))
+    args: dict[str, Any] = {"bare": True, "failFirst": False, "concurrent": False, "runs": 2}
+    expected = harness.call("recordedRun", url=wire_db.ts_url, **args)
+    schema = load_schema(harness.call("wireSchema"))
+    open_connection, driver = CHECKED_OUT_IN_A_TRANSACTION[adapter]
+
+    async def check() -> tuple[Any, Any]:
+        connection, close = await open_connection(wire_db, driver, None)
+        try:
+            actual = await recorded_run(wire_db, adapter, schema, inner=connection, **args)
+            return actual, await connection.query("SHOW statement_timeout")
+        finally:
+            await close()
+
+    actual, timeout = wire_db.run(check())
+    assert as_js(actual) == as_js(expected)
+    assert timeout == [{"statement_timeout": "4321ms"}]
+
+
+@pytest.mark.parametrize("adapter", CHECKED_OUT_IN_A_TRANSACTION)
+@pytest.mark.parametrize("level", [None, "serializable"])
+def test_transactions_match_upstream_on_a_connection_checked_out_in_a_transaction(
+    harness: Harness, wire_db: WireDB, adapter: str, level: Any
+) -> None:
+    """A transaction on a connection a pool hook left in a transaction runs at upstream's isolation level,
+    keeps what the hook set, and commits, alone and in a push's session."""
+    from sqlalchemy import text
+
+    from zero_server.server.zql_database import ZQLDatabase
+
+    wire_db.reset()
+    expected = harness.call("isolation", url=wire_db.ts_url, level=level)
+    open_connection, driver = CHECKED_OUT_IN_A_TRANSACTION[adapter]
+
+    async def callback(tx: Any, _hooks: Any) -> tuple[str, str]:
+        await tx.db_transaction.query("INSERT INTO item (id, value) VALUES ($1, 'x')", [str(uuid4())])
+        rows = await tx.db_transaction.query("SHOW transaction_isolation")
+        timeout = await tx.db_transaction.query("SHOW statement_timeout")
+        return rows[0]["transaction_isolation"], timeout[0]["statement_timeout"]
+
+    async def check() -> tuple[Any, Any, int]:
+        connection, close = await open_connection(wire_db, driver, level)
+        try:
+            db = ZQLDatabase(connection)
+            alone = await db.transaction(callback)
+            async with db.session() as session:
+                in_session = await session.transaction(callback)
+            async with wire_db.py_engine.connect() as conn:
+                committed = (await conn.execute(text("SELECT count(*) FROM item"))).scalar_one()
+            return alone, in_session, committed
+        finally:
+            await close()
+
+    alone, in_session, committed = wire_db.run(check())
+    assert alone == in_session == (expected, "4321ms")
+    assert committed == 2