From 396dca1811a554dbbafc7f259d0e076104896403 Mon Sep 17 00:00:00 2001 From: Peter Adams <18162810+Maxteabag@users.noreply.github.com> Date: Mon, 28 Sep 2026 10:59:20 +0200 Subject: [PATCH] fix(query): route transaction executor through the SSH tunnel TransactionExecutor connected with the session's original config, so multi-statement scripts, gt, BEGIN and (with the process worker off) every query went straight to the database host and skipped the tunnel. It now reuses the session's tunnel, or opens and closes its own. Also fix the SSH test container: the init script edited the template sshd_config instead of /config/sshd/sshd_config, and the tunnel mod is skipped when lscr.io is unreachable, leaving AllowTcpForwarding off. Refs #337 --- sqlit/domains/query/app/transaction.py | 44 +++++++---- .../query/ui/mixins/query_execution.py | 19 ++++- sqlit/shared/ui/protocols/query.py | 4 + tests/fixtures/99-fix-sshd-config.sh | 17 ++++- ...est_executor_reset_on_connection_switch.py | 20 +++++ tests/unit/test_transaction_executor.py | 74 +++++++++++++++++++ 6 files changed, 156 insertions(+), 22 deletions(-) diff --git a/sqlit/domains/query/app/transaction.py b/sqlit/domains/query/app/transaction.py index 6565fad9..31284b44 100644 --- a/sqlit/domains/query/app/transaction.py +++ b/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() diff --git a/sqlit/domains/query/ui/mixins/query_execution.py b/sqlit/domains/query/ui/mixins/query_execution.py index b6ea1d32..dc559ad4 100644 --- a/sqlit/domains/query/ui/mixins/query_execution.py +++ b/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 diff --git a/sqlit/shared/ui/protocols/query.py b/sqlit/shared/ui/protocols/query.py index cc6d5b7d..6e790e25 100644 --- a/sqlit/shared/ui/protocols/query.py +++ b/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: ... diff --git a/tests/fixtures/99-fix-sshd-config.sh b/tests/fixtures/99-fix-sshd-config.sh index ab3a6586..ef49e22a 100755 --- a/tests/fixtures/99-fix-sshd-config.sh +++ b/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 diff --git a/tests/unit/test_executor_reset_on_connection_switch.py b/tests/unit/test_executor_reset_on_connection_switch.py index 499d83d0..1b819465 100644 --- a/tests/unit/test_executor_reset_on_connection_switch.py +++ b/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" diff --git a/tests/unit/test_transaction_executor.py b/tests/unit/test_transaction_executor.py index 549fed2a..43357e4e 100644 --- a/tests/unit/test_transaction_executor.py +++ b/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")]