From 2a244c25c14179d78b0cbf04365b4ec74a561878 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Wed, 28 Oct 2026 16:06:52 -0700 Subject: [PATCH] fix: commit a transaction the pool hands a connection out in before using it A pool hook that runs a statement through the driver (a SQLAlchemy `checkout` listener, a psycopg-pool `check`) leaves the connection in a transaction. Bare reads then crashed on psycopg switching to autocommit, silently rolled the hook's settings back on asyncpg, a SQLAlchemy+asyncpg session tripped an assertion, a psycopg session's writes were rolled back, and isolation levels could not be set after the hook's SELECT. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/adapters/psycopg.py | 13 ++ src/zero_server/server/adapters/sqlalchemy.py | 27 +++- tests/wire/conftest.py | 12 +- tests/wire/test_mutate.py | 117 ++++++++++++++++++ 4 files changed, 165 insertions(+), 4 deletions(-) diff --git a/src/zero_server/server/adapters/psycopg.py b/src/zero_server/server/adapters/psycopg.py index 63466e0..ef05ab3 100644 --- a/src/zero_server/server/adapters/psycopg.py +++ b/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,6 +98,7 @@ 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: @@ -104,6 +106,17 @@ async def _transaction[R]( 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()`).""" diff --git a/src/zero_server/server/adapters/sqlalchemy.py b/src/zero_server/server/adapters/sqlalchemy.py index 51ab111..784abf1 100644 --- a/src/zero_server/server/adapters/sqlalchemy.py +++ b/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: diff --git a/tests/wire/conftest.py b/tests/wire/conftest.py index 9d0de4a..9d8d736 100644 --- a/tests/wire/conftest.py +++ b/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") diff --git a/tests/wire/test_mutate.py b/tests/wire/test_mutate.py index 32e5689..2fc477a 100644 --- a/tests/wire/test_mutate.py +++ b/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