Make cursor-based adapter compatible with psycopg v3 transactions and multi-statement

Running tests/integration/test_transactions.py against a real PostgreSQL
container with the psycopg3 adapter surfaced two driver-behavior
differences that need handling in the shared CursorBasedAdapter:

1. autocommit-aware execute_non_query
   psycopg2 silently makes conn.commit() a no-op when conn.autocommit is
   True; psycopg v3 actively COMMITs any open transaction. The adapter's
   execute_non_query unconditionally called conn.commit(), which under
   psycopg3 meant `executor.execute("BEGIN")` immediately ended the
   transaction it just started — so the subsequent INSERT was committed
   and ROLLBACK had nothing to undo. Skipping the commit when
   conn.autocommit is True preserves the original intent on both
   drivers.

2. nextset() walk in execute_query
   psycopg v3 exposes each result set in a multi-statement string and
   leaves the cursor positioned at the first (e.g. on BEGIN, which has
   no columns). psycopg2 only ever surfaced the last set. Walking
   cursor.nextset() until it returns falsy lands on the last set; the
   loop is a no-op on drivers that don't implement nextset() (a number
   raise NotSupportedError, hence the try/except).

The integration tests in tests/integration/test_transactions.py needed
three fixes that are independent of the driver swap — verified to fail
or hang on upstream main with psycopg2 as well:

- test_multi_statement_as_single_query_works
  The "BEGIN; INSERT; SELECT" payload left an open transaction on the
  executor's persistent connection; the test's subsequent DROP TABLE
  (on a different connection) then blocked waiting for the lock. Added
  reset_tui_executor() before the cleanup so the persistent connection
  rolls back and releases the lock.

- test_atomic_execute_rolls_back_on_error
  The test expected atomic_execute to raise on a failing statement.
  Per its docstring and current implementation, atomic_execute reports
  the failure via a MultiStatementResult with completed=False and the
  failing index in error_index. Updated the assertions accordingly.

- test_atomic_execute_commits_on_success
  Multi-statement atomic_execute is documented to return a
  MultiStatementResult (one StatementResult per statement), not a
  single QueryResult. Asserting on result.results[-1].result picks up
  the trailing SELECT's QueryResult exactly as the test intended.

Result against the docker-compose postgres:16 fixture:
  - tests/integration/test_transactions.py: 13 passed
  - tests/integration/test_database_browsing_flow.py::TestPostgreSQLDatabaseBrowsing: 1 passed
  - tests/integration/test_stale_connection_reconnect.py -k postgres: 5 passed, 1 idle-variant skipped (pre-existing)
  - tests/unit: 842 passed

Author: brentrager

sqlit/domains/connections/providers/adapters/base.py

@@ -492,6 +492,19 @@ def execute_query(self, conn: Any, query: str, max_rows: int | None = None) -> t
         """Execute a query using cursor-based approach with optional row limit."""
         cursor = conn.cursor()
         cursor.execute(query)
+        # Multi-statement queries (e.g. "BEGIN; INSERT; SELECT") produce
+        # several result sets on drivers that expose them — psycopg v3
+        # leaves the cursor positioned on the first set (the BEGIN, which
+        # has no columns), so advance to the last one. psycopg2 and other
+        # drivers either don't implement nextset() or return None for
+        # single-statement queries; either way this loop terminates.
+        while True:
+            try:
+                more = cursor.nextset()
+            except Exception:
+                more = None
+            if not more:
+                break
         if cursor.description:
             columns = [col[0] for col in cursor.description]
             if max_rows is not None:
@@ -511,7 +524,14 @@ def execute_non_query(self, conn: Any, query: str) -> int:
         cursor = conn.cursor()
         cursor.execute(query)
         rowcount = int(cursor.rowcount)
-        conn.commit()
+        # When the connection is in autocommit mode, an explicit conn.commit()
+        # has driver-dependent semantics: psycopg2 makes it a no-op, but
+        # psycopg v3 actively ends any open transaction. Calling commit()
+        # here would unconditionally close out an explicit BEGIN before the
+        # caller's INSERT/ROLLBACK gets a chance to run. Skipping it when
+        # autocommit is on preserves the intent in both worlds.
+        if not getattr(conn, "autocommit", False):
+            conn.commit()
         return rowcount
 
 

tests/integration/test_transactions.py

@@ -300,6 +300,11 @@ def test_multi_statement_as_single_query_works(self, postgres_db: str):
         assert len(result.columns) > 0
         assert result.row_count > 0
 
+        # The query opened a transaction (BEGIN) but never closed it. Close
+        # the executor so its persistent connection rolls back, otherwise
+        # the DROP TABLE below blocks waiting for the lock.
+        reset_tui_executor()
+
         # Cleanup
         conn = adapter.connect(config)
         try:
@@ -509,13 +514,19 @@ def test_atomic_execute_rolls_back_on_error(self, postgres_db: str):
             initial_count = result.rows[0][0]
 
             # This should fail on the second INSERT (NULL not allowed)
-            # and rollback the first INSERT too
+            # and rollback the first INSERT too. atomic_execute reports
+            # the failure via a MultiStatementResult with completed=False
+            # rather than raising.
+            from sqlit.domains.query.app.multi_statement import MultiStatementResult
+
             query = """
             INSERT INTO atomic_test (name) VALUES ('row1');
             INSERT INTO atomic_test (name) VALUES (NULL);
             """
-            with pytest.raises(Exception):
-                executor.atomic_execute(query)
+            atomic_result = executor.atomic_execute(query)
+            assert isinstance(atomic_result, MultiStatementResult)
+            assert atomic_result.completed is False
+            assert atomic_result.error_index is not None
 
             # Count should be unchanged (both inserts rolled back)
             result = executor.execute("SELECT COUNT(*) FROM atomic_test")
@@ -556,16 +567,23 @@ def test_atomic_execute_commits_on_success(self, postgres_db: str):
 
         executor = TransactionExecutor(config, provider)
         try:
+            from sqlit.domains.query.app.multi_statement import MultiStatementResult
+
             query = """
             INSERT INTO atomic_test (name) VALUES ('row1');
             INSERT INTO atomic_test (name) VALUES ('row2');
             SELECT * FROM atomic_test;
             """
             result = executor.atomic_execute(query)
 
-            # Should return the SELECT result
-            assert isinstance(result, QueryResult)
-            assert result.row_count == 2
+            # Multi-statement atomic_execute returns a MultiStatementResult
+            # carrying each statement's outcome. The SELECT lands as the
+            # last entry and exposes the rows.
+            assert isinstance(result, MultiStatementResult)
+            assert result.completed is True
+            select_result = result.results[-1].result
+            assert isinstance(select_result, QueryResult)
+            assert select_result.row_count == 2
         finally:
             executor.close()
             conn = adapter.connect(config)

tests/unit/test_postgresql_adapter.py

@@ -48,8 +48,8 @@ def test_postgresql_uses_custom_port_when_server_left_blank() -> None:
     psycopg2.connect so the connection actually reaches the docker
     container.
     """
-    mock_psycopg2 = MagicMock()
-    with patch.dict("sys.modules", {"psycopg2": mock_psycopg2}):
+    mock_psycopg = MagicMock()
+    with patch.dict("sys.modules", {"psycopg": mock_psycopg}):
         from sqlit.domains.connections.providers.postgresql.adapter import PostgreSQLAdapter
 
         adapter = PostgreSQLAdapter()
@@ -65,9 +65,9 @@ def test_postgresql_uses_custom_port_when_server_left_blank() -> None:
 
         adapter.connect(config)
 
-        kwargs = mock_psycopg2.connect.call_args.kwargs
+        kwargs = mock_psycopg.connect.call_args.kwargs
         assert kwargs.get("port") == 5433, (
-            "port=5433 must reach psycopg2.connect, but the adapter "
+            "port=5433 must reach psycopg.connect, but the adapter "
             f"silently drops it when server is blank. kwargs={kwargs!r}"
         )
         assert kwargs.get("host") == "localhost", (