From cd8742ee42578e0071d645c808beddb5af2de10f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 22 Jul 2026 05:08:18 +0000 Subject: [PATCH] fix(psycopg): restore autocommit after transactions --- docs/changelog.rst | 4 + sqlspec/adapters/psycopg/driver.py | 56 +++- .../psycopg/test_transaction_ownership.py | 226 +++++++++++++ .../adapters/postgres/test_cockroach_retry.py | 14 + .../unit/adapters/test_psycopg/test_driver.py | 302 ++++++++++++++++-- 5 files changed, 562 insertions(+), 40 deletions(-) create mode 100644 tests/integration/adapters/postgres/psycopg/test_transaction_ownership.py diff --git a/docs/changelog.rst b/docs/changelog.rst index a8ad163ed..3b6454409 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -42,6 +42,8 @@ v0.56.0 **Fixed:** +* Psycopg sync and async transactions now restore the connection's original + autocommit mode after SQLSpec-owned commit or rollback operations (`#648`_). * Builder caching now reuses value-independent expression templates and binds each call's current parameters and statement configuration. This also isolates CTE bodies and returned ASTs instead of sharing mutable cached @@ -207,6 +209,8 @@ v0.54.0 - SQL processing correctness and cleanup ``QueryParams`` guarded-import warning through the custom Sphinx tooling instead of changing adapter runtime code. +.. _#648: https://github.com/litestar-org/sqlspec/issues/648 + v0.52.0 - SQL Server adapters, ADK profiles, and cloud connectors ------------------------------------------------------------------------------ diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 0fa3a7471..7b41666f6 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -172,7 +172,7 @@ class PsycopgSyncDriver(PsycopgPipelineMixin, SyncDriverAdapterBase): bulk data transfer, and PostgreSQL-specific error handling. """ - __slots__ = ("_data_dictionary",) + __slots__ = ("_data_dictionary", "_restore_autocommit", "_transaction_active") dialect = "postgres" def __init__( @@ -188,6 +188,8 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsycopgSyncDataDictionary | None = None + self._restore_autocommit = False + self._transaction_active = False # ───────────────────────────────────────────────────────────────────────────── # CORE DISPATCH METHODS @@ -337,12 +339,17 @@ def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult def begin(self) -> None: """Begin a database transaction on the current connection.""" + if self._connection_in_transaction(): + return try: - if self.connection.autocommit: + restore_autocommit = self.connection.autocommit + if restore_autocommit: self.connection.autocommit = False except psycopg.Error as e: msg = f"Failed to begin transaction: {e}" raise SQLSpecError(msg) from e + self._restore_autocommit = restore_autocommit + self._transaction_active = True def commit(self) -> None: """Commit the current transaction on the current connection.""" @@ -351,6 +358,8 @@ def commit(self) -> None: except psycopg.Error as e: msg = f"Failed to commit transaction: {e}" raise SQLSpecError(msg) from e + self._transaction_active = False + self._restore_original_autocommit() def rollback(self) -> None: """Rollback the current transaction on the current connection.""" @@ -359,6 +368,8 @@ def rollback(self) -> None: except psycopg.Error as e: msg = f"Failed to rollback transaction: {e}" raise SQLSpecError(msg) from e + self._transaction_active = False + self._restore_original_autocommit() def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" @@ -634,7 +645,18 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return bool(self.connection.info.transaction_status != TRANSACTION_STATUS_IDLE) + return self._transaction_active or self.connection.info.transaction_status != TRANSACTION_STATUS_IDLE + + def _restore_original_autocommit(self) -> None: + """Restore autocommit after a completed SQLSpec-owned transaction.""" + if not self._restore_autocommit: + return + self._restore_autocommit = False + try: + self.connection.autocommit = True + except psycopg.Error as e: + msg = f"Failed to restore autocommit: {e}" + raise SQLSpecError(msg) from e class PsycopgAsyncExceptionHandler(BaseAsyncExceptionHandler): @@ -670,7 +692,7 @@ class PsycopgAsyncDriver(PsycopgPipelineMixin, AsyncDriverAdapterBase): and async pub/sub support. """ - __slots__ = ("_data_dictionary",) + __slots__ = ("_data_dictionary", "_restore_autocommit", "_transaction_active") dialect = "postgres" def __init__( @@ -686,6 +708,8 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsycopgAsyncDataDictionary | None = None + self._restore_autocommit = False + self._transaction_active = False # ───────────────────────────────────────────────────────────────────────────── # CORE DISPATCH METHODS @@ -835,12 +859,17 @@ async def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQL async def begin(self) -> None: """Begin a database transaction on the current connection.""" + if self._connection_in_transaction(): + return try: - if self.connection.autocommit: + restore_autocommit = self.connection.autocommit + if restore_autocommit: await self.connection.set_autocommit(False) except psycopg.Error as e: msg = f"Failed to begin transaction: {e}" raise SQLSpecError(msg) from e + self._restore_autocommit = restore_autocommit + self._transaction_active = True async def commit(self) -> None: """Commit the current transaction on the current connection.""" @@ -849,6 +878,8 @@ async def commit(self) -> None: except psycopg.Error as e: msg = f"Failed to commit transaction: {e}" raise SQLSpecError(msg) from e + self._transaction_active = False + await self._restore_original_autocommit() async def rollback(self) -> None: """Rollback the current transaction on the current connection.""" @@ -857,6 +888,8 @@ async def rollback(self) -> None: except psycopg.Error as e: msg = f"Failed to rollback transaction: {e}" raise SQLSpecError(msg) from e + self._transaction_active = False + await self._restore_original_autocommit() async def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" @@ -1139,7 +1172,18 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return bool(self.connection.info.transaction_status != TRANSACTION_STATUS_IDLE) + return self._transaction_active or self.connection.info.transaction_status != TRANSACTION_STATUS_IDLE + + async def _restore_original_autocommit(self) -> None: + """Restore autocommit after a completed SQLSpec-owned transaction.""" + if not self._restore_autocommit: + return + self._restore_autocommit = False + try: + await self.connection.set_autocommit(True) + except psycopg.Error as e: + msg = f"Failed to restore autocommit: {e}" + raise SQLSpecError(msg) from e def _attribute_pipeline_failure( diff --git a/tests/integration/adapters/postgres/psycopg/test_transaction_ownership.py b/tests/integration/adapters/postgres/psycopg/test_transaction_ownership.py new file mode 100644 index 000000000..99b2622b2 --- /dev/null +++ b/tests/integration/adapters/postgres/psycopg/test_transaction_ownership.py @@ -0,0 +1,226 @@ +"""Live psycopg transaction ownership and autocommit restoration tests.""" + +from typing import TYPE_CHECKING + +import psycopg +import pytest + +from sqlspec import StatementStack +from sqlspec.adapters.psycopg import PsycopgAsyncConfig, PsycopgSyncConfig + +if TYPE_CHECKING: + from pytest_databases.docker.postgres import PostgresService + +pytestmark = pytest.mark.xdist_group("postgres") + + +def _connection_config(postgres_service: "PostgresService") -> dict[str, object]: + return { + "conninfo": ( + f"postgresql://{postgres_service.user}:{postgres_service.password}@" + f"{postgres_service.host}:{postgres_service.port}/{postgres_service.database}" + ), + "autocommit": True, + "min_size": 1, + "max_size": 1, + } + + +def test_sync_transaction_ownership_and_autocommit_restoration(postgres_service: "PostgresService") -> None: + subject_config = PsycopgSyncConfig(connection_config=_connection_config(postgres_service)) + observer_config = PsycopgSyncConfig(connection_config=_connection_config(postgres_service)) + + try: + subject_config.connection_instance = subject_config.create_pool() + observer_config.connection_instance = observer_config.create_pool() + + with observer_config.provide_session() as observer: + observer.execute_script("DROP TABLE IF EXISTS test_psycopg_sync_transaction_648") + observer.execute_script( + "CREATE TABLE test_psycopg_sync_transaction_648 (id INTEGER PRIMARY KEY, value TEXT NOT NULL)" + ) + + with subject_config.provide_session() as session: + physical_connection = session.connection + session.begin() + session.execute("INSERT INTO test_psycopg_sync_transaction_648 (id, value) VALUES (%s, %s)", 1, "committed") + session.commit() + + assert session.connection.autocommit is True + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + assert session.select_value("SELECT 1") == 1 + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + + with subject_config.provide_session() as reacquired: + assert reacquired.connection is physical_connection + assert reacquired.connection.autocommit is True + + with observer_config.provide_session() as observer: + assert observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 1) == 1 + + with subject_config.provide_session() as session: + session.begin() + session.execute( + "INSERT INTO test_psycopg_sync_transaction_648 (id, value) VALUES (%s, %s)", 2, "rolled-back" + ) + session.rollback() + assert session.connection.autocommit is True + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + + with observer_config.provide_session() as observer: + assert observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 2) == 0 + + with subject_config.provide_session() as session: + session.begin() + results = session.execute_stack( + StatementStack().push_execute( + "INSERT INTO test_psycopg_sync_transaction_648 (id, value) VALUES (%s, %s)", (3, "caller-owned") + ) + ) + assert len(results) == 1 + + with observer_config.provide_session() as observer: + assert ( + observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 3) + == 0 + ) + + session.commit() + + with observer_config.provide_session() as observer: + assert observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 3) == 1 + + with subject_config.provide_session() as session: + session.driver_features["stack_native_disabled"] = True + session.begin() + results = session.execute_stack( + StatementStack().push_execute( + "INSERT INTO test_psycopg_sync_transaction_648 (id, value) VALUES (%s, %s)", + (4, "caller-owned-fallback"), + ) + ) + assert len(results) == 1 + + with observer_config.provide_session() as observer: + assert ( + observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 4) + == 0 + ) + + session.commit() + + with observer_config.provide_session() as observer: + assert observer.select_value("SELECT COUNT(*) FROM test_psycopg_sync_transaction_648 WHERE id = %s", 4) == 1 + observer.execute_script("DROP TABLE IF EXISTS test_psycopg_sync_transaction_648") + finally: + subject_config.close_pool() + observer_config.close_pool() + + +async def test_async_transaction_ownership_and_autocommit_restoration(postgres_service: "PostgresService") -> None: + subject_config = PsycopgAsyncConfig(connection_config=_connection_config(postgres_service)) + observer_config = PsycopgAsyncConfig(connection_config=_connection_config(postgres_service)) + + try: + subject_config.connection_instance = await subject_config.create_pool() + observer_config.connection_instance = await observer_config.create_pool() + + async with observer_config.provide_session() as observer: + await observer.execute_script("DROP TABLE IF EXISTS test_psycopg_async_transaction_648") + await observer.execute_script( + "CREATE TABLE test_psycopg_async_transaction_648 (id INTEGER PRIMARY KEY, value TEXT NOT NULL)" + ) + + async with subject_config.provide_session() as session: + physical_connection = session.connection + await session.begin() + await session.execute( + "INSERT INTO test_psycopg_async_transaction_648 (id, value) VALUES (%s, %s)", 1, "committed" + ) + await session.commit() + + assert session.connection.autocommit is True + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + assert await session.select_value("SELECT 1") == 1 + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + + async with subject_config.provide_session() as reacquired: + assert reacquired.connection is physical_connection + assert reacquired.connection.autocommit is True + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value("SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 1) + == 1 + ) + + async with subject_config.provide_session() as session: + await session.begin() + await session.execute( + "INSERT INTO test_psycopg_async_transaction_648 (id, value) VALUES (%s, %s)", 2, "rolled-back" + ) + await session.rollback() + assert session.connection.autocommit is True + assert session.connection.info.transaction_status is psycopg.pq.TransactionStatus.IDLE + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value("SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 2) + == 0 + ) + + async with subject_config.provide_session() as session: + await session.begin() + results = await session.execute_stack( + StatementStack().push_execute( + "INSERT INTO test_psycopg_async_transaction_648 (id, value) VALUES (%s, %s)", (3, "caller-owned") + ) + ) + assert len(results) == 1 + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value( + "SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 3 + ) + == 0 + ) + + await session.commit() + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value("SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 3) + == 1 + ) + + async with subject_config.provide_session() as session: + session.driver_features["stack_native_disabled"] = True + await session.begin() + results = await session.execute_stack( + StatementStack().push_execute( + "INSERT INTO test_psycopg_async_transaction_648 (id, value) VALUES (%s, %s)", + (4, "caller-owned-fallback"), + ) + ) + assert len(results) == 1 + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value( + "SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 4 + ) + == 0 + ) + + await session.commit() + + async with observer_config.provide_session() as observer: + assert ( + await observer.select_value("SELECT COUNT(*) FROM test_psycopg_async_transaction_648 WHERE id = %s", 4) + == 1 + ) + await observer.execute_script("DROP TABLE IF EXISTS test_psycopg_async_transaction_648") + finally: + await subject_config.close_pool() + await observer_config.close_pool() diff --git a/tests/integration/adapters/postgres/test_cockroach_retry.py b/tests/integration/adapters/postgres/test_cockroach_retry.py index 8df26b776..550581eb5 100644 --- a/tests/integration/adapters/postgres/test_cockroach_retry.py +++ b/tests/integration/adapters/postgres/test_cockroach_retry.py @@ -4,6 +4,8 @@ import pytest +from sqlspec import StatementStack + if TYPE_CHECKING: from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgDriver from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver @@ -23,6 +25,12 @@ def test_cockroach_psycopg_sync_retries_whole_transaction( def operation() -> str: nonlocal calls calls += 1 + results = contract_cockroach_psycopg_sync_driver.execute_stack( + StatementStack().push_execute("SELECT 1 AS transaction_probe") + ) + assert len(results) == 1 + assert contract_cockroach_psycopg_sync_driver._transaction_active is True # pyright: ignore[reportPrivateUsage] + assert contract_cockroach_psycopg_sync_driver.connection.autocommit is False if calls == 1: raise _RetryableCockroachError("restart transaction") return "ok" @@ -55,6 +63,12 @@ async def test_cockroach_psycopg_async_retries_whole_transaction( async def operation() -> str: nonlocal calls calls += 1 + results = await contract_cockroach_psycopg_async_driver.execute_stack( + StatementStack().push_execute("SELECT 1 AS transaction_probe") + ) + assert len(results) == 1 + assert contract_cockroach_psycopg_async_driver._transaction_active is True # pyright: ignore[reportPrivateUsage] + assert contract_cockroach_psycopg_async_driver.connection.autocommit is False if calls == 1: raise _RetryableCockroachError("restart transaction") return "ok" diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 67faf2258..09ecb6a6f 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -1,5 +1,6 @@ """Unit tests for psycopg driver transaction behavior.""" +from types import SimpleNamespace from typing import TYPE_CHECKING, cast import psycopg @@ -13,10 +14,25 @@ class _SyncTransactionConnection: - def __init__(self, method_name: str, error: Exception | None, *, autocommit: bool = True) -> None: - self.method_name = method_name - self.error = error + def __init__( + self, + *, + autocommit: bool = True, + begin_error: Exception | None = None, + commit_error: Exception | None = None, + rollback_error: Exception | None = None, + restore_error: Exception | None = None, + transaction_status: int = 0, + ) -> None: self._autocommit = autocommit + self.autocommit_calls: list[bool] = [] + self.begin_error = begin_error + self.commit_calls = 0 + self.commit_error = commit_error + self.info = SimpleNamespace(transaction_status=transaction_status) + self.restore_error = restore_error + self.rollback_calls = 0 + self.rollback_error = rollback_error @property def autocommit(self) -> bool: @@ -24,45 +40,88 @@ def autocommit(self) -> bool: @autocommit.setter def autocommit(self, value: bool) -> None: - if self.method_name == "begin" and self.error is not None: - raise self.error + self.autocommit_calls.append(value) + error = self.begin_error if value is False else self.restore_error + if error is not None: + raise error self._autocommit = value def commit(self) -> None: - if self.method_name == "commit" and self.error is not None: - raise self.error + self.commit_calls += 1 + if self.commit_error is not None: + raise self.commit_error def rollback(self) -> None: - if self.method_name == "rollback" and self.error is not None: - raise self.error + self.rollback_calls += 1 + if self.rollback_error is not None: + raise self.rollback_error class _AsyncTransactionConnection: - def __init__(self, method_name: str, error: Exception | None, *, autocommit: bool = True) -> None: - self.method_name = method_name - self.error = error + def __init__( + self, + *, + autocommit: bool = True, + begin_error: Exception | None = None, + commit_error: Exception | None = None, + rollback_error: Exception | None = None, + restore_error: Exception | None = None, + transaction_status: int = 0, + ) -> None: self.autocommit = autocommit self.autocommit_calls: list[bool] = [] + self.begin_error = begin_error + self.commit_calls = 0 + self.commit_error = commit_error + self.info = SimpleNamespace(transaction_status=transaction_status) + self.restore_error = restore_error + self.rollback_calls = 0 + self.rollback_error = rollback_error async def set_autocommit(self, value: bool) -> None: self.autocommit_calls.append(value) - if self.method_name == "begin" and self.error is not None: - raise self.error + error = self.begin_error if value is False else self.restore_error + if error is not None: + raise error self.autocommit = value async def commit(self) -> None: - if self.method_name == "commit" and self.error is not None: - raise self.error + self.commit_calls += 1 + if self.commit_error is not None: + raise self.commit_error async def rollback(self) -> None: - if self.method_name == "rollback" and self.error is not None: - raise self.error + self.rollback_calls += 1 + if self.rollback_error is not None: + raise self.rollback_error + + +def _sync_driver(connection: _SyncTransactionConnection) -> PsycopgSyncDriver: + return PsycopgSyncDriver(cast("PsycopgSyncConnection", connection)) + + +def _async_driver(connection: _AsyncTransactionConnection) -> PsycopgAsyncDriver: + return PsycopgAsyncDriver(cast("PsycopgAsyncConnection", connection)) + + +def test_sync_transaction_state_is_inactive_initially() -> None: + driver = _sync_driver(_SyncTransactionConnection()) + + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] + + +async def test_async_transaction_state_is_inactive_initially() -> None: + driver = _async_driver(_AsyncTransactionConnection()) + + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("method_name", ["begin", "commit", "rollback"]) def test_sync_transaction_control_propagates_non_native_errors(method_name: str) -> None: - connection = _SyncTransactionConnection(method_name, RuntimeError("internal bug")) - driver = PsycopgSyncDriver(cast("PsycopgSyncConnection", connection)) + error_kwargs = {f"{method_name}_error": RuntimeError("internal bug")} + driver = _sync_driver(_SyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] with pytest.raises(RuntimeError, match="internal bug"): getattr(driver, method_name)() @@ -70,39 +129,214 @@ def test_sync_transaction_control_propagates_non_native_errors(method_name: str) @pytest.mark.parametrize("method_name", ["begin", "commit", "rollback"]) def test_sync_transaction_control_wraps_native_errors(method_name: str) -> None: - connection = _SyncTransactionConnection(method_name, psycopg.Error("native failure")) - driver = PsycopgSyncDriver(cast("PsycopgSyncConnection", connection)) + error_kwargs = {f"{method_name}_error": psycopg.Error("native failure")} + driver = _sync_driver(_SyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] with pytest.raises(SQLSpecError, match=f"Failed to {method_name} transaction"): getattr(driver, method_name)() -@pytest.mark.anyio @pytest.mark.parametrize("method_name", ["begin", "commit", "rollback"]) async def test_async_transaction_control_propagates_non_native_errors(method_name: str) -> None: - connection = _AsyncTransactionConnection(method_name, RuntimeError("internal bug")) - driver = PsycopgAsyncDriver(cast("PsycopgAsyncConnection", connection)) + error_kwargs = {f"{method_name}_error": RuntimeError("internal bug")} + driver = _async_driver(_AsyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] with pytest.raises(RuntimeError, match="internal bug"): await getattr(driver, method_name)() -@pytest.mark.anyio @pytest.mark.parametrize("method_name", ["begin", "commit", "rollback"]) async def test_async_transaction_control_wraps_native_errors(method_name: str) -> None: - connection = _AsyncTransactionConnection(method_name, psycopg.Error("native failure")) - driver = PsycopgAsyncDriver(cast("PsycopgAsyncConnection", connection)) + error_kwargs = {f"{method_name}_error": psycopg.Error("native failure")} + driver = _async_driver(_AsyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] with pytest.raises(SQLSpecError, match=f"Failed to {method_name} transaction"): await getattr(driver, method_name)() -@pytest.mark.anyio -@pytest.mark.parametrize(("autocommit", "expected_calls"), [(True, [False]), (False, [])]) -async def test_async_begin_uses_autocommit_property_before_setter(autocommit: bool, expected_calls: list[bool]) -> None: - connection = _AsyncTransactionConnection("begin", None, autocommit=autocommit) - driver = PsycopgAsyncDriver(cast("PsycopgAsyncConnection", connection)) +@pytest.mark.parametrize("autocommit", [True, False]) +def test_sync_commit_restores_only_changed_autocommit(autocommit: bool) -> None: + connection = _SyncTransactionConnection(autocommit=autocommit) + driver = _sync_driver(connection) + + driver.begin() + driver.commit() + + assert connection.autocommit_calls == ([False, True] if autocommit else []) + assert connection.autocommit is autocommit + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("autocommit", [True, False]) +async def test_async_commit_restores_only_changed_autocommit(autocommit: bool) -> None: + connection = _AsyncTransactionConnection(autocommit=autocommit) + driver = _async_driver(connection) + + await driver.begin() + await driver.commit() + + assert connection.autocommit_calls == ([False, True] if autocommit else []) + assert connection.autocommit is autocommit + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +def test_sync_rollback_restores_changed_autocommit() -> None: + connection = _SyncTransactionConnection() + driver = _sync_driver(connection) + + driver.begin() + driver.rollback() + + assert connection.autocommit_calls == [False, True] + assert connection.autocommit is True + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +async def test_async_rollback_restores_changed_autocommit() -> None: + connection = _AsyncTransactionConnection() + driver = _async_driver(connection) + + await driver.begin() + await driver.rollback() + + assert connection.autocommit_calls == [False, True] + assert connection.autocommit is True + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +def test_sync_begin_is_idempotent_for_logically_owned_transaction() -> None: + connection = _SyncTransactionConnection() + driver = _sync_driver(connection) + + driver.begin() + driver.begin() + + assert connection.autocommit_calls == [False] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] + + +async def test_async_begin_is_idempotent_for_logically_owned_transaction() -> None: + connection = _AsyncTransactionConnection() + driver = _async_driver(connection) + + await driver.begin() + await driver.begin() + + assert connection.autocommit_calls == [False] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] + + +def test_sync_begin_is_idempotent_for_libpq_owned_transaction() -> None: + connection = _SyncTransactionConnection(transaction_status=2) + driver = _sync_driver(connection) + + driver.begin() + + assert connection.autocommit_calls == [] + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + + +async def test_async_begin_is_idempotent_for_libpq_owned_transaction() -> None: + connection = _AsyncTransactionConnection(transaction_status=2) + driver = _async_driver(connection) + + await driver.begin() + + assert connection.autocommit_calls == [] + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + + +def test_sync_begin_failure_preserves_inactive_state() -> None: + driver = _sync_driver(_SyncTransactionConnection(begin_error=psycopg.Error("setter failed"))) + + with pytest.raises(SQLSpecError, match="Failed to begin transaction"): + driver.begin() + + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] + + +async def test_async_begin_failure_preserves_inactive_state() -> None: + driver = _async_driver(_AsyncTransactionConnection(begin_error=psycopg.Error("setter failed"))) + + with pytest.raises(SQLSpecError, match="Failed to begin transaction"): + await driver.begin() + + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("method_name", ["commit", "rollback"]) +def test_sync_native_completion_failure_preserves_transaction_state(method_name: str) -> None: + error_kwargs = {f"{method_name}_error": psycopg.Error("native failure")} + driver = _sync_driver(_SyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] + driver.begin() + + with pytest.raises(SQLSpecError, match=f"Failed to {method_name} transaction"): + getattr(driver, method_name)() + + assert driver._transaction_active is True # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is True # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("method_name", ["commit", "rollback"]) +async def test_async_native_completion_failure_preserves_transaction_state(method_name: str) -> None: + error_kwargs = {f"{method_name}_error": psycopg.Error("native failure")} + driver = _async_driver(_AsyncTransactionConnection(**error_kwargs)) # type: ignore[arg-type] + await driver.begin() + + with pytest.raises(SQLSpecError, match=f"Failed to {method_name} transaction"): + await getattr(driver, method_name)() + + assert driver._transaction_active is True # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is True # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("method_name", ["commit", "rollback"]) +def test_sync_restoration_failure_reports_completed_transaction(method_name: str) -> None: + connection = _SyncTransactionConnection(restore_error=psycopg.Error("restore failed")) + driver = _sync_driver(connection) + driver.begin() + + with pytest.raises(SQLSpecError, match="Failed to restore autocommit: restore failed"): + getattr(driver, method_name)() + + assert connection.commit_calls == (1 if method_name == "commit" else 0) + assert connection.rollback_calls == (1 if method_name == "rollback" else 0) + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("method_name", ["commit", "rollback"]) +async def test_async_restoration_failure_reports_completed_transaction(method_name: str) -> None: + connection = _AsyncTransactionConnection(restore_error=psycopg.Error("restore failed")) + driver = _async_driver(connection) + await driver.begin() + + with pytest.raises(SQLSpecError, match="Failed to restore autocommit: restore failed"): + await getattr(driver, method_name)() + + assert connection.commit_calls == (1 if method_name == "commit" else 0) + assert connection.rollback_calls == (1 if method_name == "rollback" else 0) + assert driver._transaction_active is False # pyright: ignore[reportPrivateUsage] + assert driver._restore_autocommit is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + + +def test_sync_logical_transaction_is_reported_while_libpq_is_idle() -> None: + driver = _sync_driver(_SyncTransactionConnection(transaction_status=0)) + + driver.begin() + + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] + + +async def test_async_logical_transaction_is_reported_while_libpq_is_idle() -> None: + driver = _async_driver(_AsyncTransactionConnection(transaction_status=0)) await driver.begin() - assert connection.autocommit_calls == expected_calls + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage]