Keep history for unsaved connections in session

Author: Maxteabag

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

@@ -125,6 +125,30 @@ def _get_history_store(self: QueryMixinHost) -> Any:
             return store
         return self.services.history_store
 
+    def _get_unsaved_history_store(self: QueryMixinHost) -> Any:
+        store = getattr(self, "_unsaved_history_store", None)
+        if store is None:
+            from sqlit.domains.query.store.memory import InMemoryHistoryStore
+
+            store = InMemoryHistoryStore()
+            self._unsaved_history_store = store
+        return store
+
+    def _should_save_query_history(self: QueryMixinHost, config: Any) -> bool:
+        """Return True if the connection is saved and history should be persisted."""
+        name = getattr(config, "name", "")
+        if not name:
+            return False
+        connections = getattr(self, "connections", None) or []
+        return any(getattr(conn, "name", None) == name for conn in connections)
+
+    def _save_query_history(self: QueryMixinHost, config: Any, query: str) -> None:
+        """Save query history only for saved connections."""
+        if self._should_save_query_history(config):
+            self._get_history_store().save_query(config.name, query)
+            return
+        self._get_unsaved_history_store().save_query(config.name, query)
+
     def _get_query_service(self: QueryMixinHost, provider: Any) -> QueryService:
         if self._query_service is None or (
             self._query_service_db_type is not None
@@ -262,8 +286,6 @@ async def _run_query_async(self: QueryMixinHost, query: str, keep_insert_mode: b
 
         # Use TransactionExecutor for transaction-aware query execution
         executor = self._get_transaction_executor(config, provider)
-        service = self._get_query_service(provider)
-
         # Check if this is a multi-statement query
         statements = split_statements(query)
         is_multi_statement = len(statements) > 1
@@ -305,7 +327,7 @@ async def _run_query_async(self: QueryMixinHost, query: str, keep_insert_mode: b
                         return
 
                     try:
-                        await asyncio.to_thread(service._save_to_history, config.name, query)
+                        await asyncio.to_thread(self._save_query_history, config, query)
                     except Exception:
                         pass
                     result = outcome.result
@@ -336,7 +358,7 @@ async def _run_query_async(self: QueryMixinHost, query: str, keep_insert_mode: b
                 elapsed_ms = (time.perf_counter() - start_time) * 1000
 
                 try:
-                    await asyncio.to_thread(service._save_to_history, config.name, query)
+                    await asyncio.to_thread(self._save_query_history, config, query)
                 except Exception:
                     pass
                 self._display_multi_statement_results(multi_result, elapsed_ms)
@@ -350,7 +372,7 @@ async def _run_query_async(self: QueryMixinHost, query: str, keep_insert_mode: b
                 elapsed_ms = (time.perf_counter() - start_time) * 1000
 
                 try:
-                    await asyncio.to_thread(service._save_to_history, config.name, query)
+                    await asyncio.to_thread(self._save_query_history, config, query)
                 except Exception:
                     pass
 
@@ -401,8 +423,6 @@ async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
 
         # Create a dedicated executor for atomic execution
         executor = TransactionExecutor(config=config, provider=provider)
-        service = self._get_query_service(provider)
-
         try:
             start_time = time.perf_counter()
             max_rows = self.services.runtime.max_rows or MAX_FETCH_ROWS
@@ -414,7 +434,7 @@ async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
             elapsed_ms = (time.perf_counter() - start_time) * 1000
 
             try:
-                await asyncio.to_thread(service._save_to_history, config.name, query)
+                await asyncio.to_thread(self._save_query_history, config, query)
             except Exception:
                 pass
 
@@ -551,9 +571,11 @@ def action_show_history(self: QueryMixinHost) -> None:
 
         from ..screens import QueryHistoryScreen
 
-        history_store = self._get_history_store()
         starred_store = self.services.starred_store
-        history = history_store.load_for_connection(self.current_config.name)
+        if self._should_save_query_history(self.current_config):
+            history = self._get_history_store().load_for_connection(self.current_config.name)
+        else:
+            history = self._get_unsaved_history_store().load_for_connection(self.current_config.name)
         starred = starred_store.load_for_connection(self.current_config.name)
         self.push_screen(
             QueryHistoryScreen(history, self.current_config.name, starred),
@@ -587,16 +609,28 @@ def _show_telescope(self: QueryMixinHost, *, auto_open_filter: bool) -> None:
         """Open telescope with optional filter preset."""
         from ..screens import QueryHistoryScreen
 
+        connection_map = self._get_telescope_connection_map()
+        available_connections = set(connection_map.keys())
+
         history_store = self._get_history_store()
         if hasattr(history_store, "load_all"):
             history = history_store.load_all()
         else:
             history = []
-            for config in self._get_telescope_connection_map().values():
+            for config in connection_map.values():
                 history.extend(history_store.load_for_connection(config.name))
             history.sort(key=lambda entry: entry.timestamp, reverse=True)
 
-        connection_map = self._get_telescope_connection_map()
+        unsaved_store = getattr(self, "_unsaved_history_store", None)
+        if unsaved_store is not None and hasattr(unsaved_store, "load_all"):
+            history.extend(unsaved_store.load_all())
+
+        if available_connections:
+            history = [
+                entry for entry in history
+                if getattr(entry, "connection_name", None) in available_connections
+            ]
+        history.sort(key=lambda entry: entry.timestamp, reverse=True)
         connection_labels = {
             name: self._format_telescope_connection_label(config)
             for name, config in connection_map.items()

tests/ui/test_query_history.py

@@ -4,9 +4,19 @@
 
 import pytest
 
+from textual.widgets import OptionList
+
+from sqlit.domains.query.store.history import QueryHistoryEntry
+from sqlit.domains.query.ui.screens.query_history import QueryHistoryScreen
 from sqlit.domains.shell.app.main import SSMSTUI
 
-from .mocks import MockConnectionStore, MockSettingsStore, build_test_services, create_test_connection
+from .mocks import (
+    MockConnectionStore,
+    MockHistoryStore,
+    MockSettingsStore,
+    build_test_services,
+    create_test_connection,
+)
 
 
 class TestQueryHistoryCursorMemory:
@@ -153,3 +163,105 @@ async def test_cursor_cache_handles_same_query_text(self):
 
             # Cursor should be at the remembered position
             assert app.query_input.cursor_location == (0, 5)
+
+
+class TestQueryHistorySavePolicy:
+    """Tests for query history behavior across saved and unsaved connections."""
+
+    @pytest.mark.asyncio
+    async def test_show_history_for_unsaved_connection_uses_session_history(self) -> None:
+        unsaved_conn = create_test_connection("temp-db", "sqlite")
+        history_store = MockHistoryStore()
+        services = build_test_services(
+            connection_store=MockConnectionStore([]),
+            settings_store=MockSettingsStore({"theme": "tokyo-night"}),
+            history_store=history_store,
+        )
+        app = SSMSTUI(services=services)
+
+        async with app.run_test(size=(100, 35)) as pilot:
+            app.current_config = unsaved_conn
+            app._save_query_history(unsaved_conn, "SELECT 1")
+
+            app.action_show_history()
+            await pilot.pause(0.2)
+
+            screen = next(
+                (s for s in app.screen_stack if isinstance(s, QueryHistoryScreen)),
+                None,
+            )
+            assert screen is not None, "History screen should be present"
+
+            option_list = screen.query_one("#history-list", OptionList)
+            assert option_list.option_count == 1
+
+    def test_saved_connection_queries_saved(self) -> None:
+        saved_conn = create_test_connection("saved-db", "sqlite")
+        history_store = MockHistoryStore()
+        services = build_test_services(
+            connection_store=MockConnectionStore([saved_conn]),
+            settings_store=MockSettingsStore({"theme": "tokyo-night"}),
+            history_store=history_store,
+        )
+        app = SSMSTUI(services=services)
+        app.connections = [saved_conn]
+
+        app._save_query_history(saved_conn, "SELECT 1")
+
+        assert history_store.entries["saved-db"][0]["query"] == "SELECT 1"
+
+    @pytest.mark.asyncio
+    async def test_telescope_hides_unavailable_unsaved_history(self) -> None:
+        saved_conn = create_test_connection("saved-db", "sqlite")
+        saved_entry = QueryHistoryEntry(
+            query="select 1",
+            timestamp="2026-01-01T00:00:00",
+            connection_name="saved-db",
+        )
+        unsaved_entry = QueryHistoryEntry(
+            query="select 2",
+            timestamp="2026-01-02T00:00:00",
+            connection_name="temp-db",
+        )
+
+        class StubHistoryStore:
+            def __init__(self, entries):
+                self._entries = entries
+
+            def load_all(self):
+                return list(self._entries)
+
+            def load_for_connection(self, connection_name):
+                return [e for e in self._entries if e.connection_name == connection_name]
+
+            def delete_entry(self, connection_name, timestamp):
+                _ = connection_name
+                _ = timestamp
+                return False
+
+            def save_query(self, connection_name, query):
+                _ = connection_name
+                _ = query
+
+        history_store = StubHistoryStore([saved_entry, unsaved_entry])
+        services = build_test_services(
+            connection_store=MockConnectionStore([saved_conn]),
+            settings_store=MockSettingsStore({"theme": "tokyo-night"}),
+            history_store=history_store,
+        )
+        app = SSMSTUI(services=services)
+
+        async with app.run_test(size=(100, 35)) as pilot:
+            app.connections = [saved_conn]
+            app.action_telescope()
+            await pilot.pause(0.2)
+
+            screen = next(
+                (s for s in app.screen_stack if isinstance(s, QueryHistoryScreen)),
+                None,
+            )
+            assert screen is not None, "Telescope screen should be present"
+
+            option_list = screen.query_one("#history-list", OptionList)
+            assert option_list.option_count == 1
+            assert all(entry.connection_name == "saved-db" for entry in screen._merged_entries)