diff --git a/sqlit/domains/explorer/ui/mixins/tree_filter.py b/sqlit/domains/explorer/ui/mixins/tree_filter.py index 6cd7de4e..f7e864ed 100644 --- a/sqlit/domains/explorer/ui/mixins/tree_filter.py +++ b/sqlit/domains/explorer/ui/mixins/tree_filter.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from rich.markup import escape as escape_markup @@ -13,6 +14,49 @@ pass +@dataclass +class _NodeSnapshot: + """Frozen capture of one tree node's state, used to restore the tree + between filter keystrokes without calling refresh_tree. + + Calling refresh_tree drops lazy-loaded children (e.g. tables under a + database in multi-DB browse mode) because the re-expand triggers an + async reload that doesn't complete before the next filter pass — see + issue #141. + """ + + label: Any + data: Any + allow_expand: bool + is_expanded: bool + children: list["_NodeSnapshot"] = field(default_factory=list) + + +def _snapshot_node(node: Any) -> _NodeSnapshot: + return _NodeSnapshot( + label=node.label, + data=node.data, + allow_expand=getattr(node, "allow_expand", False), + is_expanded=getattr(node, "is_expanded", False), + children=[_snapshot_node(c) for c in node.children], + ) + + +def _restore_node_under(parent: Any, snap: _NodeSnapshot) -> None: + child = parent.add(snap.label, data=snap.data) + try: + child.allow_expand = snap.allow_expand + except Exception: + pass + if snap.is_expanded: + try: + child.expand() + except Exception: + pass + for grandchild in snap.children: + _restore_node_under(child, grandchild) + + class TreeFilterMixin: """Mixin providing tree filter functionality.""" @@ -24,6 +68,7 @@ class TreeFilterMixin: _tree_filter_matches: list[Any] = [] _tree_filter_match_index: int = 0 _tree_original_labels: dict[int, str] = {} + _tree_snapshot: list[_NodeSnapshot] | None = None def action_tree_filter(self: TreeFilterMixinHost) -> None: """Open the tree filter.""" @@ -38,6 +83,12 @@ def action_tree_filter(self: TreeFilterMixinHost) -> None: self._tree_filter_matches = [] self._tree_filter_match_index = 0 self._tree_original_labels = {} + # Freeze the currently loaded tree (incl. lazy-loaded children) + # so we can restore it between keystrokes without calling + # refresh_tree, which would lose async-loaded folder contents. + self._tree_snapshot = [ + _snapshot_node(c) for c in self.object_tree.root.children + ] self.tree_filter_input.show() self._update_tree_filter() @@ -52,7 +103,8 @@ def action_tree_filter_close(self: TreeFilterMixinHost) -> None: self._tree_filter_typing = False self.tree_filter_input.hide() self._restore_tree_labels() - self._show_all_tree_nodes() + self._restore_tree_from_snapshot() + self._tree_snapshot = None self._update_footer_bindings() def action_tree_filter_accept(self: TreeFilterMixinHost) -> None: @@ -182,11 +234,12 @@ def _update_tree_filter(self: TreeFilterMixinHost) -> None: self._tree_filter_fuzzy = raw_text.startswith("~") self._tree_filter_query = raw_text[1:] if self._tree_filter_fuzzy else raw_text - # Rebuild the full tree first so each filter pass searches every node, - # not just the survivors of the previous (narrower) filter. Without this, - # backspacing from a no-match query like "tt" back to "t" would leave the - # tree empty because the "t"-matching nodes were physically removed. - self._show_all_tree_nodes() + # Restore from the snapshot taken when the filter opened, so each + # filter pass searches every node (not just the survivors of the + # previous narrower filter — see PR #211 for the backspace case) + # while preserving lazy-loaded children that refresh_tree would + # have dropped (issue #141). + self._restore_tree_from_snapshot() self._tree_original_labels = {} total = self._count_all_nodes() @@ -325,7 +378,27 @@ def _set_node_visibility( def _show_all_tree_nodes(self: TreeFilterMixinHost) -> None: """Rebuild the tree to restore all nodes after filtering.""" - self.refresh_tree() + if self._tree_snapshot is not None: + self._restore_tree_from_snapshot() + else: + # Fallback for paths that aren't inside an open filter session. + self.refresh_tree() + + def _restore_tree_from_snapshot(self: TreeFilterMixinHost) -> None: + """Rebuild the root's children from the snapshot taken at filter open.""" + snapshot = self._tree_snapshot + if snapshot is None: + return + root = self.object_tree.root + # Clear existing children (works for both Textual TreeNode and the + # test mock — both implement child.remove()). + for child in list(root.children): + try: + child.remove() + except Exception: + pass + for snap in snapshot: + _restore_node_under(root, snap) def _restore_tree_labels(self: TreeFilterMixinHost) -> None: """Restore original labels for all modified nodes.""" diff --git a/tests/ui/explorer/test_tree_filter_search.py b/tests/ui/explorer/test_tree_filter_search.py index 4366ec8e..5e3e6c41 100644 --- a/tests/ui/explorer/test_tree_filter_search.py +++ b/tests/ui/explorer/test_tree_filter_search.py @@ -14,7 +14,12 @@ from unittest.mock import MagicMock -from sqlit.domains.explorer.domain.tree_nodes import ConnectionNode +from sqlit.domains.explorer.domain.tree_nodes import ( + ConnectionNode, + DatabaseNode, + FolderNode, + TableNode, +) from sqlit.domains.explorer.ui.mixins.tree_filter import TreeFilterMixin @@ -237,3 +242,145 @@ def test_backspace_to_empty_restores_all(self): visible = _visible_node_names(host) assert set(visible) == set(self.CONNECTION_NAMES) + + +class _MultiDbFilterHost(TreeFilterMixin): + """Filter host that models multi-database 'browse all' mode. + + The tree starts shaped like a real session after the user has expanded + into one database's Tables folder: + + connection + └── Databases + ├── CS (expanded — tables already loaded) + │ └── Tables + │ ├── cs_user + │ ├── cs_session + │ └── cs_ticket + └── Sales (collapsed — lazy-loaded if expanded) + + `refresh_tree` here mirrors the real `refresh_tree_incremental`: it + rebuilds the connection + Databases + database nodes synchronously, + but the contents of the Tables folder are reloaded asynchronously and + are *not* present when refresh_tree returns. This is what produces + issue #141 — opening the filter calls refresh_tree, the lazy-loaded + tables vanish, and the subsequent filter search finds nothing. + """ + + def __init__( + self, + connection_name: str, + databases: list[str], + tables_by_db: dict[str, list[str]], + ): + self._connection_name = connection_name + self._databases = databases + self._tables_by_db = tables_by_db + self.object_tree = MockTree() + self.tree_filter_input = MockFilterInput() + self._populate(include_lazy_children=True) + + def _populate(self, *, include_lazy_children: bool) -> None: + self.object_tree.root.children = [] + config = MagicMock() + config.name = self._connection_name + conn_node = object.__new__(ConnectionNode) + object.__setattr__(conn_node, "config", config) + conn = self.object_tree.root.add(self._connection_name, data=conn_node) + conn.allow_expand = True + conn.is_expanded = True + + dbs_folder = conn.add("Databases", data=FolderNode(folder_type="databases")) + dbs_folder.allow_expand = True + dbs_folder.is_expanded = True + + for db_name in self._databases: + db_node = dbs_folder.add(db_name, data=DatabaseNode(name=db_name)) + db_node.allow_expand = True + if db_name not in self._tables_by_db: + continue + db_node.is_expanded = True + tables_folder = db_node.add( + "Tables", + data=FolderNode(folder_type="tables", database=db_name), + ) + tables_folder.allow_expand = True + tables_folder.is_expanded = True + + # Real refresh_tree only restores the shell here. The Tables + # folder gets re-expanded asynchronously and its children + # don't materialize before _update_tree_filter runs. + if include_lazy_children: + for table in self._tables_by_db[db_name]: + leaf = tables_folder.add( + table, + data=TableNode(database=db_name, schema="", name=table), + ) + leaf.allow_expand = True + + def refresh_tree(self) -> None: + # Match the real behavior: shell only, lazy children absent. + self._populate(include_lazy_children=False) + + def _update_footer_bindings(self) -> None: + pass + + def _activate_tree_node(self, _node) -> None: + pass + + +class TestMultiDbFilterIssue141: + """Regression tests for issue #141. + + When sqlit is connected in multi-database 'browse all' mode and the + user has expanded into a database's Tables folder, opening `/` and + typing a substring of a table name should find that table. It + currently returns zero matches because `_update_tree_filter` calls + `refresh_tree`, which tears down the tree and re-expands lazy folders + asynchronously — the search runs before the tables are reloaded. + """ + + def _open_filter(self, host: _MultiDbFilterHost) -> None: + TreeFilterMixin.action_tree_filter(host) # type: ignore[arg-type] + + def _type(self, host: _MultiDbFilterHost, text: str) -> None: + for ch in text: + host._tree_filter_text += ch + TreeFilterMixin._update_tree_filter(host) # type: ignore[arg-type] + + def _matched_names(self, host: _MultiDbFilterHost) -> list[str]: + names: list[str] = [] + for node in host._tree_filter_matches: + data = node.data + if data is not None and hasattr(data, "get_label_text"): + names.append(data.get_label_text()) + return sorted(names) + + def test_filter_finds_lazy_loaded_tables_in_multi_db_mode(self): + host = _MultiDbFilterHost( + connection_name="prod", + databases=["CS", "Sales"], + tables_by_db={"CS": ["cs_user", "cs_session", "cs_ticket"]}, + ) + + # Sanity check: before we open the filter, the tables are present. + cs_node = host.object_tree.root.children[0].children[0].children[0] + tables_folder = cs_node.children[0] + assert sorted(c.label for c in tables_folder.children) == [ + "cs_session", + "cs_ticket", + "cs_user", + ] + + self._open_filter(host) + self._type(host, "cs") + + matched = self._matched_names(host) + # "cs" matches the CS database node itself plus its tables. + # Issue #141: without the fix, only "CS" would be in the list + # because refresh_tree had wiped the lazy-loaded table nodes. + assert "cs_user" in matched + assert "cs_session" in matched + assert "cs_ticket" in matched, ( + f"Issue #141: filter must find lazy-loaded tables; got {matched}" + )