fix(#648): restore psycopg autocommit after transactions

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
 ------------------------------------------------------------------------------
 

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(

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()