diff --git a/sqlit/domains/query/app/multi_statement.py b/sqlit/domains/query/app/multi_statement.py index 26077d32..2c9479ce 100644 --- a/sqlit/domains/query/app/multi_statement.py +++ b/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): diff --git a/sqlit/domains/query/app/transaction.py b/sqlit/domains/query/app/transaction.py index 31284b44..7b7e3455 100644 --- a/sqlit/domains/query/app/transaction.py +++ b/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: diff --git a/tests/unit/test_transaction_executor.py b/tests/unit/test_transaction_executor.py index 43357e4e..720d96ff 100644 --- a/tests/unit/test_transaction_executor.py +++ b/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)