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")]