Merge pull request #122 from Maxteabag/feature/auto-reconnect-after-driver-install

Auto-reconnect after driver installation

Author: Maxteabag

sqlit/domains/connections/ui/connection_error_handlers.py

@@ -36,9 +36,19 @@ def can_handle(self, error: Exception) -> bool:
 
     def handle(self, app: ConnectionErrorApp, error: Exception, config: ConnectionConfig) -> None:
         from sqlit.domains.connections.providers.exceptions import MissingDriverError
+        from sqlit.shared.core.debug_events import emit_debug_event
 
+        from .restart_cache import write_pending_connection_cache
         from .screens import PackageSetupScreen
 
+        # Save pending connection for auto-reconnect after driver install restart
+        if config.name:
+            write_pending_connection_cache(config.name)
+            emit_debug_event(
+                "driver_install.pending_connection_saved",
+                connection_name=config.name,
+            )
+
         # No on_success callback - uses default "Restart to apply" behavior
         app.push_screen(PackageSetupScreen(cast(MissingDriverError, error)))
 

sqlit/domains/connections/ui/restart_cache.py

@@ -28,3 +28,18 @@ def clear_restart_cache() -> None:
         get_restart_cache_path().unlink(missing_ok=True)
     except Exception:
         pass
+
+
+def write_pending_connection_cache(connection_name: str) -> None:
+    """Cache a pending connection name for auto-reconnect after driver install restart.
+
+    This is used when a user tries to connect to a server but the driver is missing.
+    After the driver is installed and the app restarts, it can auto-connect to this
+    connection.
+    """
+    payload = {
+        "version": 2,
+        "type": "pending_connection",
+        "connection_name": connection_name,
+    }
+    write_restart_cache(payload)

sqlit/domains/explorer/ui/mixins/tree.py

@@ -216,6 +216,18 @@ def action_refresh_tree(self: TreeMixinHost) -> None:
         if hasattr(self, "_loading_nodes"):
             self._loading_nodes.clear()
         self._schema_service = None
+
+        # Reload saved connections from disk (in case added via CLI)
+        try:
+            services = getattr(self, "services", None)
+            if services:
+                store = getattr(services, "connection_store", None)
+                if store:
+                    reloaded = store.load_all(load_credentials=False)
+                    self.connections = reloaded
+        except Exception:
+            pass  # Keep existing connections if reload fails
+
         self.refresh_tree()
         loader = getattr(self, "_load_schema_cache", None)
         if callable(loader):

sqlit/domains/shell/app/commands/debug.py

@@ -41,6 +41,16 @@ def _set_debug_enabled(app: Any, enabled: bool) -> None:
     else:
         app._debug_events_enabled = bool(enabled)
 
+    # Persist the setting across sessions
+    try:
+        services = getattr(app, "services", None)
+        if services:
+            store = getattr(services, "settings_store", None)
+            if store:
+                store.set("debug_events_enabled", enabled)
+    except Exception:
+        pass
+
     path = getattr(app, "_debug_event_log_path", None)
     suffix = f" (log: {path})" if path else ""
     state = "enabled" if enabled else "disabled"

sqlit/domains/shell/app/startup_flow.py

@@ -36,6 +36,10 @@ def run_on_mount(app: AppProtocol) -> None:
     app._startup_stamp("settings_loaded")
 
     app._expanded_paths = set(settings.get("expanded_nodes", []))
+    if settings.get("debug_events_enabled"):
+        setter = getattr(app, "_set_debug_events_enabled", None)
+        if callable(setter):
+            setter(True)
     if "process_worker" in settings:
         app.services.runtime.process_worker = bool(settings.get("process_worker"))
     if "process_worker_warm_on_idle" in settings:
@@ -83,6 +87,9 @@ def run_on_mount(app: AppProtocol) -> None:
         app.object_tree.cursor_line = 0
     app._update_section_labels()
     maybe_restore_connection_screen(app)
+    # Auto-connect to pending connection after driver install (if not already connecting)
+    if app._startup_connect_config is None:
+        maybe_auto_connect_pending(app)
     app._startup_stamp("restore_checked")
     if app._debug_mode:
         app.call_after_refresh(app._record_launch_ms)
@@ -224,6 +231,83 @@ def _get_restart_cache_path() -> Path:
     return Path(tempfile.gettempdir()) / "sqlit-driver-install-restore.json"
 
 
+def maybe_auto_connect_pending(app: AppProtocol) -> bool:
+    """Auto-connect to a pending connection after driver install restart.
+
+    Returns True if a connection was initiated, False otherwise.
+    """
+    from sqlit.shared.core.debug_events import emit_debug_event
+
+    from sqlit.domains.connections.ui.restart_cache import (
+        clear_restart_cache,
+        get_restart_cache_path,
+    )
+
+    cache_path = get_restart_cache_path()
+    emit_debug_event(
+        "startup.pending_connection_check",
+        cache_path=str(cache_path),
+        exists=cache_path.exists(),
+    )
+    if not cache_path.exists():
+        return False
+
+    emit_debug_event(
+        "startup.pending_connection_found",
+        contents=cache_path.read_text(),
+    )
+
+    try:
+        payload = json.loads(cache_path.read_text(encoding="utf-8"))
+    except Exception as e:
+        emit_debug_event("startup.pending_connection_parse_error", error=str(e))
+        clear_restart_cache()
+        return False
+
+    # Always clear cache after reading
+    clear_restart_cache()
+
+    # Check for version 2 pending_connection type
+    if not isinstance(payload, dict):
+        emit_debug_event("startup.pending_connection_invalid", reason="not a dict")
+        return False
+    if payload.get("version") != 2:
+        emit_debug_event("startup.pending_connection_invalid", reason="wrong version", version=payload.get("version"))
+        return False
+    if payload.get("type") != "pending_connection":
+        emit_debug_event("startup.pending_connection_invalid", reason="wrong type", type=payload.get("type"))
+        return False
+
+    connection_name = payload.get("connection_name")
+    if not connection_name:
+        emit_debug_event("startup.pending_connection_invalid", reason="no connection_name")
+        return False
+
+    emit_debug_event(
+        "startup.pending_connection_lookup",
+        connection_name=connection_name,
+        available_connections=[getattr(c, "name", None) for c in app.connections],
+    )
+
+    # Find the connection by name
+    config = next(
+        (c for c in app.connections if getattr(c, "name", None) == connection_name),
+        None,
+    )
+    if config is None:
+        emit_debug_event("startup.pending_connection_not_found", connection_name=connection_name)
+        return False
+
+    emit_debug_event("startup.pending_connection_connecting", connection_name=connection_name)
+
+    # Auto-connect after refresh (same pattern as startup_connect_config)
+    def _connect_pending() -> None:
+        app.connect_to_server(config)
+
+    app.call_after_refresh(_connect_pending)
+    return True
+
+
 def maybe_restore_connection_screen(app: AppProtocol) -> None:
     """Restore an in-progress connection form after a driver-install restart."""
     cache_path = _get_restart_cache_path()
@@ -239,14 +323,16 @@ def maybe_restore_connection_screen(app: AppProtocol) -> None:
             pass
         return
 
+    # Only handle version 1 (connection form restore), leave version 2 for maybe_auto_connect_pending
+    if not isinstance(payload, dict) or payload.get("version") != 1:
+        return
+
+    # Clear cache only for version 1
     try:
         cache_path.unlink(missing_ok=True)
     except Exception:
         pass
 
-    if not isinstance(payload, dict) or payload.get("version") != 1:
-        return
-
     values = payload.get("values")
     if not isinstance(values, dict):
         return

tests/unit/test_auto_reconnect_after_driver_install.py

@@ -0,0 +1,100 @@
+"""Test auto-reconnect after driver installation restart."""
+
+from __future__ import annotations
+
+import json
+from unittest.mock import MagicMock
+
+from sqlit.domains.connections.domain.config import ConnectionConfig
+
+
+class TestAutoReconnectAfterDriverInstall:
+    """Test that app auto-connects after driver install restart."""
+
+    def test_pending_connection_cache_written_on_missing_driver(self):
+        """
+        When user tries to connect but driver is missing,
+        the connection name should be cached for auto-reconnect after restart.
+        """
+        from sqlit.domains.connections.ui.restart_cache import (
+            get_restart_cache_path,
+            write_pending_connection_cache,
+        )
+
+        config = ConnectionConfig(name="my-mssql-server", db_type="mssql")
+
+        # Write the pending connection cache
+        write_pending_connection_cache(config.name)
+
+        # Verify cache was written
+        cache_path = get_restart_cache_path()
+        assert cache_path.exists()
+
+        payload = json.loads(cache_path.read_text())
+        assert payload["version"] == 2
+        assert payload["type"] == "pending_connection"
+        assert payload["connection_name"] == "my-mssql-server"
+
+        # Cleanup
+        cache_path.unlink(missing_ok=True)
+
+    def test_startup_reads_pending_connection_and_connects(self):
+        """
+        On startup, if pending_connection cache exists,
+        app should auto-connect to that connection.
+        """
+        from sqlit.domains.connections.ui.restart_cache import (
+            get_restart_cache_path,
+            write_pending_connection_cache,
+        )
+        from sqlit.domains.shell.app.startup_flow import maybe_auto_connect_pending
+
+        # Setup: Write pending connection cache
+        write_pending_connection_cache("my-mssql-server")
+
+        # Mock app with the saved connection
+        mock_app = MagicMock()
+        saved_config = ConnectionConfig(name="my-mssql-server", db_type="mssql")
+        mock_app.connections = [saved_config]
+        mock_app.connect_to_server = MagicMock()
+        mock_app.call_after_refresh = MagicMock()
+
+        # Call the startup function
+        result = maybe_auto_connect_pending(mock_app)
+
+        # Should have scheduled a connection via call_after_refresh
+        assert result is True
+        mock_app.call_after_refresh.assert_called_once()
+
+        # Execute the callback to verify it calls connect_to_server
+        callback = mock_app.call_after_refresh.call_args[0][0]
+        callback()
+        mock_app.connect_to_server.assert_called_once_with(saved_config)
+
+        # Cache should be cleared
+        assert not get_restart_cache_path().exists()
+
+    def test_startup_ignores_missing_connection(self):
+        """
+        If the cached connection no longer exists, don't crash.
+        """
+        from sqlit.domains.connections.ui.restart_cache import (
+            get_restart_cache_path,
+            write_pending_connection_cache,
+        )
+        from sqlit.domains.shell.app.startup_flow import maybe_auto_connect_pending
+
+        write_pending_connection_cache("deleted-connection")
+
+        mock_app = MagicMock()
+        mock_app.connections = []  # No connections
+        mock_app.connect_to_server = MagicMock()
+
+        result = maybe_auto_connect_pending(mock_app)
+
+        # Should return False (no connection made)
+        assert result is False
+        mock_app.connect_to_server.assert_not_called()
+
+        # Cache should still be cleared
+        assert not get_restart_cache_path().exists()