fix(query): run multi-statement scripts on one connection

sqlit/domains/query/app/multi_statement.py

@@ -10,6 +10,7 @@
 
 import re
 from collections.abc import Iterator
+from contextlib import nullcontext
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -457,6 +458,13 @@ def execute(self, sql: str, max_rows: int | None = None) -> MultiStatementResult
         if not statements:
             return MultiStatementResult(results=[], completed=True, error_index=None)
 
+        # Run the script on one connection when the executor supports it, so
+        # session state carries from one statement to the next.
+        batch = getattr(self._executor, "batch", None)
+        with batch() if callable(batch) else nullcontext():
+            return self._execute_statements(statements, max_rows)
+
+    def _execute_statements(self, statements: list[str], max_rows: int | None) -> MultiStatementResult:
         results: list[StatementResult] = []
 
         for i, statement in enumerate(statements):

sqlit/domains/query/app/transaction.py

@@ -10,6 +10,8 @@
 from __future__ import annotations
 
 import re
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
@@ -157,6 +159,8 @@ class TransactionExecutor:
     _created_tunnel: Any | None = None
     _state: TransactionStateManager | None = None
     _transaction_connection: Any | None = None
+    _batch_connection: Any | None = None
+    _batch_active: bool = False
     _analyzer: KeywordQueryAnalyzer | None = None
 
     def __post_init__(self) -> None:
@@ -186,6 +190,29 @@ def in_transaction(self) -> bool:
         """Whether we're currently inside a transaction."""
         return self._state.in_transaction if self._state else False
 
+    @contextmanager
+    def batch(self) -> Iterator[None]:
+        """Run every execute() inside this block on one connection.
+
+        Outside a transaction each execute() otherwise gets a fresh connection, so
+        session state (SET @var, temp tables, SET search_path) would not carry from
+        one statement of a script to the next.
+        """
+        self._batch_active = True
+        try:
+            yield
+        finally:
+            self._batch_active = False
+            conn = self._batch_connection
+            self._batch_connection = None
+            # A script that opened a transaction without ending it keeps the
+            # connection as the transaction connection.
+            if conn is not None and conn is not self._transaction_connection:
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+
     def execute(self, sql: str, max_rows: int | None = None) -> QueryResult | NonQueryResult:
         """Execute a query with transaction awareness.
 
@@ -218,10 +245,16 @@ def execute(self, sql: str, max_rows: int | None = None) -> QueryResult | NonQue
         )
 
         if use_persistent:
-            # Reuse or create a persistent connection for transaction scope
+            # Reuse or create a persistent connection for transaction scope.
+            # Inside a batch, the batch's connection carries the transaction so
+            # session state set earlier in the script is kept.
             if self._transaction_connection is None:
-                self._transaction_connection = self._connect()
+                self._transaction_connection = self._batch_connection or self._connect()
             conn = self._transaction_connection
+        elif self._batch_active:
+            if self._batch_connection is None:
+                self._batch_connection = self._connect()
+            conn = self._batch_connection
         else:
             # Not in transaction - use temporary connection
             conn = self._connect()
@@ -374,10 +407,12 @@ def atomic_execute(
     def _close_transaction_connection(self) -> None:
         """Close the persistent transaction connection."""
         if self._transaction_connection is not None:
-            try:
-                self._transaction_connection.close()
-            except Exception:
-                pass
+            # The batch still needs its connection once the transaction ends.
+            if self._transaction_connection is not self._batch_connection:
+                try:
+                    self._transaction_connection.close()
+                except Exception:
+                    pass
             self._transaction_connection = None
 
     def close(self) -> None:

tests/unit/test_transaction_executor.py

@@ -150,3 +150,62 @@ def test_executor_without_ssh_connects_to_configured_host() -> None:
     config = ConnectionConfig.from_dict({"name": "direct", "db_type": "postgresql", "server": "db", "port": "5432"})
     TransactionExecutor(config, provider).execute("SELECT 1")
     assert _endpoints(provider) == [("db", "5432")]
+
+
+def test_batch_runs_script_on_one_connection() -> None:
+    """Regression #337: session state (SET @var) must carry between statements of a script."""
+    executor, provider = _make_executor()
+
+    with executor.batch():
+        executor.execute("SET @x = 1")
+        executor.execute("SELECT @x")
+        executor.execute("UPDATE t SET a = 1")
+
+    created = provider.connection_factory.created
+    assert len(created) == 1
+    assert created[0].closed is True
+    assert {conn for _kind, conn, _sql in provider.query_executor.calls} == {created[0]}
+
+
+def test_batch_transaction_uses_batch_connection_and_keeps_it_after_commit() -> None:
+    executor, provider = _make_executor()
+
+    with executor.batch():
+        executor.execute("SET @x = 1")
+        executor.execute("BEGIN")
+        executor.execute("INSERT INTO t VALUES (@x)")
+        executor.execute("COMMIT")
+        assert provider.connection_factory.created[0].closed is False
+        executor.execute("SELECT @x")
+
+    created = provider.connection_factory.created
+    assert len(created) == 1
+    assert created[0].closed is True
+    assert executor.in_transaction is False
+
+
+def test_batch_ending_inside_transaction_keeps_connection_open() -> None:
+    executor, provider = _make_executor()
+
+    with executor.batch():
+        executor.execute("SET @x = 1")
+        executor.execute("BEGIN")
+        executor.execute("INSERT INTO t VALUES (@x)")
+
+    conn = provider.connection_factory.created[0]
+    assert executor.in_transaction is True
+    assert conn.closed is False
+
+    executor.execute("COMMIT")
+    assert len(provider.connection_factory.created) == 1
+    assert conn.closed is True
+
+
+def test_without_batch_each_statement_gets_its_own_connection() -> None:
+    executor, provider = _make_executor()
+
+    executor.execute("SELECT 1")
+    executor.execute("SELECT 2")
+
+    assert len(provider.connection_factory.created) == 2
+    assert all(conn.closed for conn in provider.connection_factory.created)