feat: connect in a worker, show catalog loading (#399)

tconbeer · web-flow · commit 699847ce1d07 · 2024-01-10T14:27:15.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Features
+
+-   Harlequin now loads immediately and connects to your database in the background ([#393](https://github.com/tconbeer/harlequin/issues/393)).
+-   Harlequin shows a loading indicator before the Data Catalog is hydrated for the first time ([#396](https://github.com/tconbeer/harlequin/issues/396)).
+
 ## [1.9.2] - 2024-01-10
 
 ### Features
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -50,6 +50,7 @@ harlequin-trino = { version = "^0.1", optional = true }
 pre-commit = "^3.3.1"
 textual-dev = "^1.0.1"
 harlequin-postgres = "^0.2"
+harlequin-mysql = "^0.1.1"
 
 [tool.poetry.group.static.dependencies]
 black = "^23.3.0"
diff --git a/src/harlequin/app.py b/src/harlequin/app.py
@@ -23,6 +23,7 @@
 from textual_fastdatatable import DataTable
 from textual_fastdatatable.backend import AutoBackendType
 
+from harlequin import HarlequinConnection
 from harlequin.adapter import HarlequinAdapter, HarlequinCursor
 from harlequin.autocomplete import completer_factory
 from harlequin.cache import BufferState, Cache, write_cache
@@ -42,12 +43,20 @@
 from harlequin.exception import (
     HarlequinConfigError,
     HarlequinConnectionError,
+    HarlequinError,
     HarlequinQueryError,
     HarlequinThemeError,
+    pretty_error_message,
     pretty_print_error,
 )
 
 
+class DatabaseConnected(Message):
+    def __init__(self, connection: HarlequinConnection) -> None:
+        super().__init__()
+        self.connection = connection
+
+
 class QuerySubmitted(Message):
     def __init__(self, query_text: str, limit: int | None) -> None:
         super().__init__()
@@ -126,14 +135,7 @@ def __init__(
             )
             self.exit(return_code=2)
         self.query_timer: Union[float, None] = None
-        try:
-            self.connection = self.adapter.connect()
-        except HarlequinConnectionError as e:
-            pretty_print_error(e)
-            self.exit(return_code=2)
-        else:
-            if self.connection.init_message:
-                self.notify(self.connection.init_message)
+        self.connection: HarlequinConnection | None = None
 
         try:
             self.app_colors = HarlequinColors.from_theme(theme)
@@ -150,9 +152,12 @@ def compose(self) -> ComposeResult:
             yield DataCatalog(type_color=self.app_colors.gray)
             with Vertical(id="main_panel"):
                 yield EditorCollection(language="sql", theme=self.theme)
-                yield RunQueryBar(max_results=self.max_results)
+                yield RunQueryBar(
+                    max_results=self.max_results, classes="non-responsive"
+                )
                 yield ResultsViewer(
-                    max_results=self.max_results, type_color=self.app_colors.gray
+                    max_results=self.max_results,
+                    type_color=self.app_colors.gray,
                 )
         yield Footer()
 
@@ -186,8 +191,9 @@ async def on_mount(self) -> None:
 
         self.editor.focus()
         self.run_query_bar.checkbox.value = False
+        self.data_catalog.loading = True
 
-        self.update_schema_data()
+        self._connect()
 
     @on(Button.Pressed, "#run_query")
     def submit_query_from_run_query_bar(self, message: Button.Pressed) -> None:
@@ -208,6 +214,17 @@ def submit_query_from_editor(self, message: CodeEditor.Submitted) -> None:
             )
         )
 
+    @on(DatabaseConnected)
+    def initialize_app(self, message: DatabaseConnected) -> None:
+        self.connection = message.connection
+        self.run_query_bar.set_responsive()
+        self.results_viewer.show_table(did_run=False)
+        if message.connection.init_message:
+            self.notify(message.connection.init_message, title="Database Connected.")
+        else:
+            self.notify("Database Connected.")
+        self.update_schema_data()
+
     @on(DataCatalog.NodeSubmitted)
     def insert_node_into_editor(
         self, message: DataCatalog.NodeSubmitted[CatalogItem]
@@ -277,11 +294,11 @@ def copy_data_to_clipboard(self, message: DataTable.SelectionCopied) -> None:
                 self.notify("Selected data copied to clipboard.")
 
     @on(Worker.StateChanged)
-    def handle_worker_error(self, message: Worker.StateChanged) -> None:
+    async def handle_worker_error(self, message: Worker.StateChanged) -> None:
         if message.state == WorkerState.ERROR:
-            self._handle_worker_error(message)
+            await self._handle_worker_error(message)
 
-    def _handle_worker_error(self, message: Worker.StateChanged) -> None:
+    async def _handle_worker_error(self, message: Worker.StateChanged) -> None:
         if (
             message.worker.name == "update_schema_data"
             and message.worker.error is not None
@@ -296,16 +313,33 @@ def _handle_worker_error(self, message: Worker.StateChanged) -> None:
         ):
             self.run_query_bar.set_responsive()
             self.results_viewer.show_table()
-            header = getattr(message.worker.error, "title", "Query Error")
+            header = getattr(
+                message.worker.error, "title", message.worker.error.__class__.__name__
+            )
             self._push_error_modal(
                 title="Query Error",
                 header=header,
                 error=message.worker.error,
             )
+        elif message.worker.name == "_connect" and message.worker.error is not None:
+            title = getattr(
+                message.worker.error,
+                "title",
+                "Harlequin could not connect to your database.",
+            )
+            error = (
+                message.worker.error
+                if isinstance(message.worker.error, HarlequinError)
+                else HarlequinConnectionError(
+                    msg=str(message.worker.error), title=title
+                )
+            )
+            self.exit(return_code=2, message=pretty_error_message(error))
 
     @on(NewCatalog)
     def update_tree_and_completers(self, message: NewCatalog) -> None:
         self.data_catalog.update_tree(message.catalog)
+        self.data_catalog.loading = False
         self.update_completers(message.catalog)
 
     @on(QueriesExecuted)
@@ -383,6 +417,8 @@ def watch_full_screen(self, full_screen: bool) -> None:
 
     @on(QuerySubmitted)
     def execute_query(self, message: QuerySubmitted) -> None:
+        if self.connection is None:
+            return
         if message.query_text:
             self.full_screen = False
             self.run_query_bar.set_not_responsive()
@@ -472,6 +508,17 @@ def action_toggle_sidebar(self) -> None:
         else:
             self.sidebar_hidden = not self.sidebar_hidden
 
+    @work(
+        thread=True,
+        exclusive=True,
+        exit_on_error=False,
+        group="connect",
+        description="Connecting to DB",
+    )
+    def _connect(self) -> None:
+        connection = self.adapter.connect()
+        self.post_message(DatabaseConnected(connection=connection))
+
     @work(
         thread=True,
         exclusive=True,
@@ -480,6 +527,8 @@ def action_toggle_sidebar(self) -> None:
         description="Executing queries.",
     )
     def _execute_query(self, message: QuerySubmitted) -> None:
+        if self.connection is None:
+            return
         cursors: Dict[str, HarlequinCursor] = {}
         queries = self._split_query_text(message.query_text)
         for q in queries:
@@ -514,7 +563,7 @@ def _push_error_modal(self, title: str, header: str, error: BaseException) -> No
                 title=title,
                 header=header,
                 error=error,
-            )
+            ),
         )
 
     @work(
@@ -543,6 +592,8 @@ def _fetch_data(
 
     @work(thread=True, exclusive=True, exit_on_error=True, group="completer_builders")
     def update_completers(self, catalog: Catalog) -> None:
+        if self.connection is None:
+            return
         if (
             self.editor_collection.word_completer is not None
             and self.editor_collection.member_completer is not None
@@ -561,6 +612,8 @@ def update_completers(self, catalog: Catalog) -> None:
 
     @work(thread=True, exclusive=True, exit_on_error=False, group="schema_updaters")
     def update_schema_data(self) -> None:
+        if self.connection is None:
+            return
         catalog = self.connection.get_catalog()
         self.post_message(NewCatalog(catalog=catalog))
 
@@ -570,6 +623,8 @@ def _validate_selection(self) -> str:
         return the empty string.
         """
         selection = self.editor.selected_text
+        if self.connection is None:
+            return selection
         if selection:
             try:
                 return self.connection.validate_sql(selection)
diff --git a/src/harlequin/exception.py b/src/harlequin/exception.py
@@ -1,3 +1,6 @@
+from rich.panel import Panel
+
+
 class HarlequinExit(Exception):
     pass
 
@@ -39,13 +42,14 @@ class HarlequinTzDataError(HarlequinError):
 
 def pretty_print_error(error: HarlequinError) -> None:
     from rich import print
-    from rich.panel import Panel
-
-    print(
-        Panel.fit(
-            str(error),
-            title=error.title if error.title else ("Harlequin encountered an error."),
-            title_align="left",
-            border_style="red",
-        )
+
+    print(pretty_error_message(error))
+
+
+def pretty_error_message(error: HarlequinError) -> Panel:
+    return Panel.fit(
+        str(error),
+        title=error.title if error.title else ("Harlequin encountered an error."),
+        title_align="left",
+        border_style="red",
     )
diff --git a/src/harlequin/global.tcss b/src/harlequin/global.tcss
@@ -77,6 +77,10 @@ Toast.-error {
     border-left: wide $error;
 }
 
+Toast .toast--title {
+    color: $primary;
+}
+
 /* DATA CATALOG */
 DataCatalog {
     height: 1fr;
