fix(query): route transaction executor through the SSH tunnel

sqlit/domains/query/app/transaction.py

@@ -153,6 +153,8 @@ class TransactionExecutor:
 
     config: ConnectionConfig
     provider: DatabaseProvider
+    tunnel: Any | None = None  # Optional existing SSH tunnel to reuse (e.g. the session's)
+    _created_tunnel: Any | None = None
     _state: TransactionStateManager | None = None
     _transaction_connection: Any | None = None
     _analyzer: KeywordQueryAnalyzer | None = None
@@ -161,6 +163,24 @@ def __post_init__(self) -> None:
         self._state = TransactionStateManager()
         self._analyzer = KeywordQueryAnalyzer()
 
+    def _connect(self) -> Any:
+        """Open a connection, routed through the SSH tunnel when one is configured."""
+        connect_config = self.config
+        tunnel = self.tunnel or self._created_tunnel
+        if tunnel is None and self.config.tunnel and self.config.tunnel.enabled:
+            from sqlit.domains.connections.app.tunnel import create_ssh_tunnel
+
+            self._created_tunnel, _host, _port = create_ssh_tunnel(self.config)
+            tunnel = self._created_tunnel
+        if tunnel is not None:
+            connect_config = self.config.with_endpoint(host="127.0.0.1", port=str(tunnel.local_bind_port))
+        conn = self.provider.connection_factory.connect(connect_config)
+        try:
+            self.provider.post_connect(conn, self.config)
+        except Exception:
+            pass
+        return conn
+
     @property
     def in_transaction(self) -> bool:
         """Whether we're currently inside a transaction."""
@@ -200,19 +220,11 @@ def execute(self, sql: str, max_rows: int | None = None) -> QueryResult | NonQue
         if use_persistent:
             # Reuse or create a persistent connection for transaction scope
             if self._transaction_connection is None:
-                self._transaction_connection = self.provider.connection_factory.connect(self.config)
-                try:
-                    self.provider.post_connect(self._transaction_connection, self.config)
-                except Exception:
-                    pass
+                self._transaction_connection = self._connect()
             conn = self._transaction_connection
         else:
             # Not in transaction - use temporary connection
-            conn = self.provider.connection_factory.connect(self.config)
-            try:
-                self.provider.post_connect(conn, self.config)
-            except Exception:
-                pass
+            conn = self._connect()
             is_temp_connection = True
 
         try:
@@ -291,11 +303,7 @@ def atomic_execute(
         statements = [s for s in statements if not is_comment_only_statement(s)]
 
         # Create a dedicated connection for this atomic operation
-        conn = self.provider.connection_factory.connect(self.config)
-        try:
-            self.provider.post_connect(conn, self.config)
-        except Exception:
-            pass
+        conn = self._connect()
 
         try:
             # Start transaction
@@ -378,5 +386,11 @@ def close(self) -> None:
         Always call this when done with the executor.
         """
         self._close_transaction_connection()
+        if self._created_tunnel is not None:
+            try:
+                self._created_tunnel.stop()
+            except Exception:
+                pass
+            self._created_tunnel = None
         if self._state:
             self._state.reset()

sqlit/domains/query/ui/mixins/query_execution.py

@@ -47,6 +47,7 @@ class QueryExecutionMixin(ProcessWorkerLifecycleMixin):
     _cancellable_query: CancellableQuery | None = None
     _transaction_executor: TransactionExecutor | None = None
     _transaction_executor_config: Any | None = None
+    _transaction_executor_tunnel: Any | None = None
     _query_spinner: Spinner | None = None
     _query_cursor_cache: dict[str, tuple[int, int]] | None = None
     _query_target_database: str | None = None
@@ -355,19 +356,31 @@ def _get_transaction_executor(self: QueryMixinHost, config: Any, provider: Any)
         from sqlit.domains.query.app.transaction import TransactionExecutor
 
         # Create new executor if none exists or if config changed
-        if self._transaction_executor is None or self._transaction_executor_config != config:
+        tunnel = self._session_tunnel()
+        if (
+            self._transaction_executor is None
+            or self._transaction_executor_config != config
+            or self._transaction_executor_tunnel is not tunnel
+        ):
             if self._transaction_executor is not None:
                 self._transaction_executor.close()
-            self._transaction_executor = TransactionExecutor(config=config, provider=provider)
+            self._transaction_executor = TransactionExecutor(config=config, provider=provider, tunnel=tunnel)
             self._transaction_executor_config = config
+            self._transaction_executor_tunnel = tunnel
         return self._transaction_executor
 
+    def _session_tunnel(self: QueryMixinHost) -> Any | None:
+        """The SSH tunnel of the active session, so executors reuse it instead of the raw host."""
+        session = getattr(self, "_session", None)
+        return getattr(session, "tunnel", None) if session is not None else None
+
     def _reset_transaction_executor(self: QueryMixinHost) -> None:
         """Reset the transaction executor (e.g., on disconnect)."""
         if self._transaction_executor is not None:
             self._transaction_executor.close()
             self._transaction_executor = None
         self._transaction_executor_config = None
+        self._transaction_executor_tunnel = None
 
     def _on_disconnect(self: QueryMixinHost) -> None:
         """Handle disconnect lifecycle event."""
@@ -635,7 +648,7 @@ async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
             config = provider.apply_database_override(config, active_db)
 
         # Create a dedicated executor for atomic execution
-        executor = TransactionExecutor(config=config, provider=provider)
+        executor = TransactionExecutor(config=config, provider=provider, tunnel=self._session_tunnel())
         try:
             start_time = time.perf_counter()
             max_rows = self.services.runtime.max_rows or MAX_FETCH_ROWS

sqlit/shared/ui/protocols/query.py

@@ -33,6 +33,7 @@ class QueryStateProtocol(Protocol):
     _undo_history: Any | None
     _transaction_executor: Any | None
     _transaction_executor_config: Any | None
+    _transaction_executor_tunnel: Any | None
     _results_render_worker: Worker[Any] | None
     _results_render_token: int
     _query_document: Any | None
@@ -259,6 +260,9 @@ def _reset_transaction_executor(self) -> None:
     def _get_transaction_executor(self, config: Any, provider: Any) -> Any:
         ...
 
+    def _session_tunnel(self) -> Any | None:
+        ...
+
     def _display_multi_statement_results(self, multi_result: Any, elapsed_ms: float) -> None:
         ...
 

tests/fixtures/99-fix-sshd-config.sh

@@ -1,6 +1,15 @@
 #!/bin/bash
-# Increase MaxStartups and MaxSessions for testing
-echo "MaxStartups 100:30:200" >> /etc/ssh/sshd_config
-echo "MaxSessions 100" >> /etc/ssh/sshd_config
-# Reload sshd to pick up the changes
+# sshd runs with -f /config/sshd/sshd_config; /etc/ssh/sshd_config is only the template.
+for cfg in /config/sshd/sshd_config /etc/ssh/sshd_config; do
+    [ -f "$cfg" ] || continue
+    # Enable TCP forwarding for SSH tunnel tests. The image ships "AllowTcpForwarding no",
+    # and the openssh-server-ssh-tunnel mod is downloaded at startup, so it silently
+    # does nothing when lscr.io can't be reached.
+    sed -i 's/^AllowTcpForwarding no/AllowTcpForwarding yes/' "$cfg"
+    grep -q '^AllowTcpForwarding yes' "$cfg" || echo "AllowTcpForwarding yes" >> "$cfg"
+    # Increase MaxStartups and MaxSessions for testing
+    grep -q '^MaxStartups 100:30:200' "$cfg" || echo "MaxStartups 100:30:200" >> "$cfg"
+    grep -q '^MaxSessions 100' "$cfg" || echo "MaxSessions 100" >> "$cfg"
+done
+# Reload sshd to pick up the changes if it is already running
 pkill -HUP sshd 2>/dev/null || true

tests/unit/test_executor_reset_on_connection_switch.py

@@ -48,6 +48,7 @@ def _make_mixin_instance(self) -> MagicMock:
             instance, config, provider
         )
         instance._reset_transaction_executor = lambda: QueryExecutionMixin._reset_transaction_executor(instance)
+        instance._session_tunnel = lambda: QueryExecutionMixin._session_tunnel(instance)
 
         return instance
 
@@ -251,3 +252,22 @@ def test_full_connection_switch_flow(self, mock_executor_class: MagicMock) -> No
         # THIS IS THE BUG FIX: should get server B's executor, not server A's
         assert exec_b is executor_b
         assert exec_b is not executor_a
+
+
+@patch(EXECUTOR_PATCH_TARGET)
+def test_new_session_tunnel_creates_new_executor(mock_executor_class: MagicMock) -> None:
+    """A reconnect opens a new SSH tunnel; the executor must not keep the stale one (#337)."""
+    instance = TestExecutorResetOnConfigChange()._make_mixin_instance()
+    config = TestExecutorResetOnConfigChange()._make_config("server-a", host="db.internal")
+    provider = MagicMock()
+
+    instance._session = MagicMock(tunnel="tunnel-1")
+    instance._get_transaction_executor(config, provider)
+    instance._get_transaction_executor(config, provider)
+    assert mock_executor_class.call_count == 1
+    assert mock_executor_class.call_args.kwargs["tunnel"] == "tunnel-1"
+
+    instance._session = MagicMock(tunnel="tunnel-2")
+    instance._get_transaction_executor(config, provider)
+    assert mock_executor_class.call_count == 2
+    assert mock_executor_class.call_args.kwargs["tunnel"] == "tunnel-2"

tests/unit/test_transaction_executor.py

@@ -18,8 +18,10 @@ def close(self) -> None:
 class FakeConnectionFactory:
     def __init__(self) -> None:
         self.created: list[FakeConnection] = []
+        self.configs: list[ConnectionConfig] = []
 
     def connect(self, config: ConnectionConfig) -> FakeConnection:
+        self.configs.append(config)
         conn = FakeConnection(f"conn-{len(self.created)}")
         self.created.append(conn)
         return conn
@@ -76,3 +78,75 @@ def test_execute_begin_commit_closes_connection() -> None:
     assert executor._transaction_connection is None
     assert len(provider.connection_factory.created) == 1
     assert provider.connection_factory.created[0].closed is True
+
+
+class FakeTunnel:
+    local_bind_port = 40123
+
+    def __init__(self) -> None:
+        self.stopped = False
+
+    def stop(self) -> None:
+        self.stopped = True
+
+
+def _ssh_config() -> ConnectionConfig:
+    return ConnectionConfig.from_dict(
+        {
+            "name": "ssh",
+            "db_type": "postgresql",
+            "server": "db.internal.example",
+            "port": "5432",
+            "ssh_enabled": True,
+            "ssh_host": "bastion",
+            "ssh_username": "me",
+        }
+    )
+
+
+def _endpoints(provider: FakeProvider) -> list[tuple[str, str]]:
+    return [(c.tcp_endpoint.host, c.tcp_endpoint.port) for c in provider.connection_factory.configs]
+
+
+def test_ssh_executor_reuses_session_tunnel_for_every_path() -> None:
+    """Regression #337: transaction paths connected to the private DB host, bypassing the tunnel."""
+    provider = FakeProvider()
+    tunnel = FakeTunnel()
+    executor = TransactionExecutor(_ssh_config(), provider, tunnel=tunnel)
+
+    executor.execute("SELECT 1")
+    executor.execute("BEGIN")
+    executor.execute("SELECT 2")
+    executor.execute("COMMIT")
+    executor.atomic_execute("SET x = 1; SELECT 3")
+    executor.close()
+
+    assert _endpoints(provider) == [("127.0.0.1", "40123")] * 3
+    assert tunnel.stopped is False  # the session owns it
+
+
+def test_ssh_executor_without_session_tunnel_opens_and_stops_its_own(monkeypatch) -> None:
+    created: list[FakeTunnel] = []
+
+    def fake_create_ssh_tunnel(config: ConnectionConfig):
+        tunnel = FakeTunnel()
+        created.append(tunnel)
+        return tunnel, "127.0.0.1", tunnel.local_bind_port
+
+    monkeypatch.setattr("sqlit.domains.connections.app.tunnel.create_ssh_tunnel", fake_create_ssh_tunnel)
+    provider = FakeProvider()
+    executor = TransactionExecutor(_ssh_config(), provider)
+
+    executor.execute("SELECT 1")
+    executor.atomic_execute("SELECT 2")
+    executor.close()
+
+    assert _endpoints(provider) == [("127.0.0.1", "40123")] * 2
+    assert len(created) == 1 and created[0].stopped
+
+
+def test_executor_without_ssh_connects_to_configured_host() -> None:
+    provider = FakeProvider()
+    config = ConnectionConfig.from_dict({"name": "direct", "db_type": "postgresql", "server": "db", "port": "5432"})
+    TransactionExecutor(config, provider).execute("SELECT 1")
+    assert _endpoints(provider) == [("db", "5432")]