feat/catalog improvements

.pre-commit-config.yaml

@@ -25,6 +25,7 @@ repos:
           - questionary
           - tomlkit
           - pandas-stubs
+          - importlib_metadata
         args:
           - "--disallow-untyped-calls"
           - "--disallow-untyped-defs"

CHANGELOG.md

@@ -6,7 +6,9 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
--   Harlequin now shows a more helpful error message when attempting to open a sqlite file with the duckdb adapter (or vice versa) ([#401](https://github.com/tconbeer/harlequin/issues/401)).
+-   Harlequin now shows a more helpful error message when attempting to open a sqlite file with the duckdb adapter or vice versa ([#401](https://github.com/tconbeer/harlequin/issues/401)).
+-   <kbd>ctrl+r</kbd> forces a refresh of the Data Catalog (the catalog is automatically refreshed after DDL queries are executed in Harlequin) ([#375](https://github.com/tconbeer/harlequin/issues/375)).
+-   At startup, Harlequin attempts to load a cached version of the Data Catalog. The Data Catalog will be updated in the background. A loading indicator will be displayed if there is no cached catalog for the connection parameters ([#397](https://github.com/tconbeer/harlequin/issues/397)).
 
 ### Bug Fixes
 

src/harlequin/app.py

@@ -26,8 +26,8 @@
 from harlequin import HarlequinConnection
 from harlequin.adapter import HarlequinAdapter, HarlequinCursor
 from harlequin.autocomplete import completer_factory
-from harlequin.cache import BufferState, Cache, write_cache
 from harlequin.catalog import Catalog, CatalogItem, NewCatalog
+from harlequin.catalog_cache import get_cached_catalog, update_cache_with_catalog
 from harlequin.colors import HarlequinColors
 from harlequin.components import (
     CodeEditor,
@@ -40,6 +40,8 @@
     RunQueryBar,
     export_callback,
 )
+from harlequin.editor_cache import BufferState, Cache
+from harlequin.editor_cache import write_cache as write_editor_cache
 from harlequin.exception import (
     HarlequinConfigError,
     HarlequinConnectionError,
@@ -107,6 +109,7 @@ class Harlequin(App, inherit_bindings=False):
         Binding("f9", "toggle_sidebar", "Toggle Sidebar", show=False),
         Binding("f10", "toggle_full_screen", "Toggle Full Screen Mode", show=False),
         Binding("ctrl+e", "export", "Export Data", show=False),
+        Binding("ctrl+r", "refresh_catalog", "Refresh Data Catalog", show=False),
     ]
 
     full_screen: reactive[bool] = reactive(False)
@@ -115,6 +118,8 @@ class Harlequin(App, inherit_bindings=False):
     def __init__(
         self,
         adapter: HarlequinAdapter,
+        *,
+        connection_hash: str | None = None,
         theme: str = "harlequin",
         max_results: int | str = 100_000,
         driver_class: Union[Type[Driver], None] = None,
@@ -123,6 +128,8 @@ def __init__(
     ):
         super().__init__(driver_class, css_path, watch_css)
         self.adapter = adapter
+        self.connection_hash = connection_hash
+        self.catalog: Catalog | None = None
         self.theme = theme
         try:
             self.max_results = int(max_results)
@@ -339,6 +346,7 @@ async def _handle_worker_error(self, message: Worker.StateChanged) -> None:
 
     @on(NewCatalog)
     def update_tree_and_completers(self, message: NewCatalog) -> None:
+        self.catalog = message.catalog
         self.data_catalog.update_tree(message.catalog)
         self.data_catalog.loading = False
         self.update_completers(message.catalog)
@@ -488,7 +496,9 @@ async def action_quit(self) -> None:
             buffers.append(
                 BufferState(editor.cursor, editor.selection_anchor, editor.text)
             )
-        write_cache(Cache(focus_index=focus_index, buffers=buffers))
+        write_editor_cache(Cache(focus_index=focus_index, buffers=buffers))
+        if self.catalog is not None and self.connection_hash is not None:
+            update_cache_with_catalog(self.connection_hash, self.catalog)
         await super().action_quit()
 
     def action_show_help_screen(self) -> None:
@@ -509,6 +519,10 @@ def action_toggle_sidebar(self) -> None:
         else:
             self.sidebar_hidden = not self.sidebar_hidden
 
+    def action_refresh_catalog(self) -> None:
+        self.data_catalog.loading = True
+        self.update_schema_data()
+
     @work(
         thread=True,
         exclusive=True,
@@ -518,6 +532,10 @@ def action_toggle_sidebar(self) -> None:
     )
     def _connect(self) -> None:
         connection = self.adapter.connect()
+        if self.connection_hash is not None and (
+            cached_catalog := get_cached_catalog(self.connection_hash)
+        ):
+            self.post_message(NewCatalog(catalog=cached_catalog))
         self.post_message(DatabaseConnected(connection=connection))
 
     @work(
@@ -591,7 +609,13 @@ def _fetch_data(
             ResultsFetched(cursors=cursors, data=data, errors=errors, elapsed=elapsed)
         )
 
-    @work(thread=True, exclusive=True, exit_on_error=True, group="completer_builders")
+    @work(
+        thread=True,
+        exclusive=True,
+        exit_on_error=True,
+        group="completer_builders",
+        description="building completers",
+    )
     def update_completers(self, catalog: Catalog) -> None:
         if self.connection is None:
             return

src/harlequin/cache.py

@@ -1,63 +1,8 @@
-import pickle
-from dataclasses import dataclass
-from pathlib import Path
-from typing import List, Union
+"""
+This is required for backward-compatibility for previously-pickled caches.
 
-from platformdirs import user_cache_dir
-from textual_textarea.key_handlers import Cursor as Cursor
+cache.py was renamed to editor_cache.py when the Data Catalog cache was created.
+"""
+from harlequin.editor_cache import BufferState, Cache
 
-CACHE_VERSION = 0
-
-
-@dataclass
-class BufferState:
-    cursor: Cursor
-    selection_anchor: Union[Cursor, None]
-    text: str
-
-
-@dataclass
-class Cache:
-    focus_index: int  # currently doesn't impact focus on load
-    buffers: List[BufferState]
-
-
-def get_cache_file() -> Path:
-    """
-    Returns the path to the cache file on disk
-    """
-    cache_dir = Path(user_cache_dir(appname="harlequin"))
-    cache_file = cache_dir / f"cache-{CACHE_VERSION}.pickle"
-    return cache_file
-
-
-def load_cache() -> Union[Cache, None]:
-    """
-    Returns a Cache (a list of strings) by loading
-    from a pickle saved to disk
-    """
-    cache_file = get_cache_file()
-    try:
-        with cache_file.open("rb") as f:
-            cache: Cache = pickle.load(f)
-            assert isinstance(cache, Cache)
-    except (
-        pickle.UnpicklingError,
-        ValueError,
-        IndexError,
-        FileNotFoundError,
-        AssertionError,
-    ):
-        return None
-    else:
-        return cache
-
-
-def write_cache(cache: Cache) -> None:
-    """
-    Updates dumps buffer contents to to disk
-    """
-    cache_file = get_cache_file()
-    cache_file.parent.mkdir(parents=True, exist_ok=True)
-    with open(cache_file, "wb") as f:
-        pickle.dump(cache, f)
+__all__ = ["BufferState", "Cache"]

src/harlequin/catalog_cache.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+from platformdirs import user_cache_dir
+from textual_textarea.key_handlers import Cursor as Cursor
+
+from harlequin.catalog import Catalog
+
+CACHE_VERSION = 0
+
+
+class PermissiveEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, Path):
+            return str(obj)
+        # Never raise a TypeError, just use the repr
+        try:
+            return str(obj)
+        except TypeError:
+            return ""
+
+
+@dataclass
+class CatalogCache:
+    databases: dict[str, Catalog]
+
+
+def get_connection_hash(conn_str: Sequence[str], config: dict[str, Any]) -> str:
+    return (
+        hashlib.md5(
+            json.dumps(
+                {"conn_str": tuple(conn_str), **config},
+                cls=PermissiveEncoder,
+            ).encode("utf-8")
+        )
+        .digest()
+        .hex()
+    )
+
+
+def get_cached_catalog(connection_hash: str) -> Catalog | None:
+    cache = _load_cache()
+    if cache is None:
+        return None
+    return cache.databases.get(connection_hash, None)
+
+
+def update_cache_with_catalog(connection_hash: str, catalog: Catalog) -> None:
+    cache = _load_cache()
+    if cache is None:
+        cache = CatalogCache(databases={})
+    cache.databases[connection_hash] = catalog
+    _write_cache(cache)
+
+
+def _get_cache_file() -> Path:
+    """
+    Returns the path to the cache file on disk
+    """
+    cache_dir = Path(user_cache_dir(appname="harlequin"))
+    cache_file = cache_dir / f"catalog-cache-{CACHE_VERSION}.pickle"
+    return cache_file
+
+
+def _load_cache() -> CatalogCache | None:
+    """
+    Returns a Cache by loading from a pickle saved to disk
+    """
+    cache_file = _get_cache_file()
+    try:
+        with cache_file.open("rb") as f:
+            cache: CatalogCache = pickle.load(f)
+            assert isinstance(cache, CatalogCache)
+    except (
+        pickle.UnpicklingError,
+        ValueError,
+        IndexError,
+        FileNotFoundError,
+        AssertionError,
+    ):
+        return None
+    else:
+        return cache
+
+
+def _write_cache(cache: CatalogCache) -> None:
+    """
+    Updates cache with current data catalog
+    """
+    cache_file = _get_cache_file()
+    cache_file.parent.mkdir(parents=True, exist_ok=True)
+    with open(cache_file, "wb") as f:
+        pickle.dump(cache, f)

src/harlequin/cli.py

@@ -8,6 +8,7 @@
 
 from harlequin import Harlequin
 from harlequin.adapter import HarlequinAdapter
+from harlequin.catalog_cache import get_connection_hash
 from harlequin.colors import GREEN, PINK, PURPLE, YELLOW
 from harlequin.config import get_config_for_profile
 from harlequin.config_wizard import wizard
@@ -248,11 +249,15 @@ def inner_cli(
                 pretty_print_error(e)
                 ctx.exit(2)
 
-        # load and instantiate the adapter
-        adapter = config.pop("adapter", DEFAULT_ADAPTER)
+        # remove the harlequin config from the options passed to the adapter
         conn_str: Sequence[str] = config.pop("conn_str", tuple())  # type: ignore
         if isinstance(conn_str, str):
             conn_str = (conn_str,)
+        max_results: str | int = config.pop("limit", DEFAULT_LIMIT)  # type: ignore
+        theme: str = config.pop("theme", DEFAULT_THEME)  # type: ignore
+
+        # load and instantiate the adapter
+        adapter = config.pop("adapter", DEFAULT_ADAPTER)
         adapter_cls: type[HarlequinAdapter] = adapters[adapter]  # type: ignore
         try:
             adapter_instance = adapter_cls(conn_str=conn_str, **config)
@@ -262,8 +267,9 @@ def inner_cli(
 
         tui = Harlequin(
             adapter=adapter_instance,
-            max_results=config.get("limit", DEFAULT_LIMIT),  # type: ignore
-            theme=config.get("theme", DEFAULT_THEME),  # type: ignore
+            connection_hash=get_connection_hash(conn_str, config),
+            max_results=max_results,
+            theme=theme,
         )
         tui.run()
 

src/harlequin/components/code_editor.py

@@ -14,8 +14,8 @@
 from textual_textarea.key_handlers import Cursor
 
 from harlequin.autocomplete import MemberCompleter, WordCompleter
-from harlequin.cache import BufferState, load_cache
 from harlequin.components.error_modal import ErrorModal
+from harlequin.editor_cache import BufferState, load_cache
 
 
 class CodeEditor(TextArea):

src/harlequin/components/help_screen.md

@@ -8,6 +8,7 @@
 - F9, ctrl+b: Toggle the sidebar.
 - F10: Toggle full screen mode for the current widget.
 - ctrl+e: Write the returned data to a CSV, Parquet, or JSON file.
+- ctrl+r: Refresh the Data Catalog.
 
 
 ### Query Editor Bindings