Loading state fix (#173)

* fix: refactor use of workers, fix loading state

* fix: use pilot.pause() to await cpu idle in tests

* chore: update changelog

CHANGELOG.md

@@ -4,8 +4,14 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### New Features
+
 -   Harlequin now returns the result of multiple select queries to different tabs in the Results Viewer. To run multiple queries, type them into the Query Editor (separated by semicolons), then press <kbd>ctrl+a</kbd> to select all, and then <kbd>ctrl+enter</kbd> to run the selection ([#34](https://github.com/tconbeer/harlequin/issues/34)).
 -   If there are multiple results tabs, you can switch between them with <kbd>j</kbd> and <kbd>k</kbd>.
+
+### Bug Fixes
+-   Fixes issues with the loading state when loading large result sets.
+
 ## [0.0.24] - 2023-08-04
 
 -   Adds a new CLI option, `--extension` or `-e`, which will install and load a named DuckDB extension.

src/harlequin/tui/app.py

@@ -383,28 +383,41 @@ async def watch_relations(
         """
         # invalidate results so watch_data runs even if the results are the same
         self.results_viewer.clear_all_tables()
+        self.set_result_viewer_data(relations)
+
+    @work(
+        exclusive=True,
+        exit_on_error=True,
+        group="duck_query_runners",
+        description="fetching data from duckdb.",
+    )
+    async def set_result_viewer_data(
+        self, relations: Dict[str, duckdb.DuckDBPyRelation]
+    ) -> None:
         data: Dict[str, List[Tuple]] = {}
         for id_, rel in relations.items():
             self.results_viewer.push_table(table_id=id_, relation=rel)
-
             try:
-                worker = self.fetch_relation_data(rel)
-                await worker.wait()
-            except WorkerFailed as e:
+                rel_data = rel.fetchall()
+            except duckdb.DataError as e:
                 self.push_screen(
                     ErrorModal(
                         title="DuckDB Error",
                         header=("DuckDB raised an error when running your query:"),
-                        error=e.error,
+                        error=e,
                     )
                 )
                 self.results_viewer.show_table()
             else:
-                if worker.result is not None:
-                    data[id_] = worker.result
+                data[id_] = rel_data
         self.results_viewer.data = data
 
-    @work(exclusive=True, exit_on_error=False)
+    @work(
+        exclusive=True,
+        exit_on_error=False,
+        group="duck_query_runners",
+        description="Building relation.",
+    )
     async def _build_relation(
         self, query_text: str
     ) -> Dict[str, duckdb.DuckDBPyRelation]:
@@ -418,18 +431,7 @@ async def _build_relation(
                 relations[table_id] = rel
         return relations
 
-    @work(exclusive=True, exit_on_error=False)
-    async def fetch_relation_data(
-        self, relation: duckdb.DuckDBPyRelation
-    ) -> List[Tuple]:
-        data = relation.fetchall()
-        worker = get_current_worker()
-        if not worker.is_cancelled:
-            return data
-        else:
-            return []
-
-    @work(exclusive=True)
+    @work(exclusive=True, group="duck_schema_updaters")
     async def update_schema_data(self) -> None:
         catalog = get_catalog(self.connection)
         worker = get_current_worker()

src/harlequin/tui/components/results_viewer.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Dict, Iterator, List, Tuple, Union
 
 import duckdb
@@ -15,7 +16,7 @@
     TabbedContent,
     TabPane,
 )
-from textual.worker import Worker, get_current_worker
+from textual.worker import Worker, WorkerState
 
 from harlequin.tui.utils import short_type
 
@@ -78,7 +79,13 @@ def _focus_on_visible_table(self) -> None:
         if maybe_table is not None:
             maybe_table.focus()
 
-    def on_tabbed_content_tab_activated(self) -> None:
+    def on_tabbed_content_tab_activated(
+        self, message: TabbedContent.TabActivated
+    ) -> None:
+        message.stop()
+        # Don't update the border if we're still loading the table.
+        if self.border_title and str(self.border_title).startswith("Loading"):
+            return
         maybe_table = self.get_visible_table()
         if maybe_table is not None and self.data:
             id_ = maybe_table.id
@@ -117,37 +124,41 @@ def get_visible_table(self) -> Union[ResultsTable, None]:
 
     async def watch_data(self, data: Dict[str, List[Tuple]]) -> None:
         if data:
-            self.set_not_responsive(data=data)
-            for table_id, result in data.items():
-                table = self.tab_switcher.query_one(f"#{table_id}", ResultsTable)
-                for i, chunk in self.chunk(result[: self.MAX_RESULTS]):
-                    worker = self.add_data_to_table(table, chunk)
-                    await worker.wait()
-                    self.increment_progress_bar()
-                    if i == 0:
-                        self.show_table()
-            self.set_responsive(data=data)
+            await self.set_not_responsive(data=data)
+            self.load_data(data=data)
+
+    @work(exclusive=True, group="data_loaders", description="Loading data.")
+    async def load_data(self, data: Dict[str, List[Tuple]]) -> Dict[str, List[Tuple]]:
+        for table_id, result in data.items():
+            table = self.tab_switcher.query_one(f"#{table_id}", ResultsTable)
+            for i, chunk in self.chunk(result[: self.MAX_RESULTS]):
+                table.add_rows(chunk)
+                self.increment_progress_bar()
+                await asyncio.sleep(0)
+                if i == 0:
+                    self.show_table()
+        return data
+
+    def on_worker_state_changed(self, event: Worker.StateChanged) -> None:
+        if (
+            event.worker.name == "load_data"
+            and event.worker.state == WorkerState.SUCCESS
+        ):
+            self.set_responsive(data=event.worker.result)
             self.post_message(self.Ready())
             self.focus()
 
     @staticmethod
     def chunk(
-        data: List[Tuple], chunksize: int = 2000
+        data: List[Tuple], chunksize: int = 1000
     ) -> Iterator[Tuple[int, List[Tuple]]]:
         for i in range(len(data) // chunksize + 1):
             yield i, data[i * chunksize : (i + 1) * chunksize]
 
-    @work(exclusive=False)
-    async def add_data_to_table(self, table: ResultsTable, data: List[Tuple]) -> Worker:
-        worker = get_current_worker()
-        if not worker.is_cancelled:
-            table.add_rows(data)
-        return worker
-
-    def set_not_responsive(self, data: Dict[str, List[Tuple]]) -> None:
+    async def set_not_responsive(self, data: Dict[str, List[Tuple]]) -> None:
         if len(data) > 1:
             self.border_title = f"Loading Data from {len(data):,} Queries."
-        else:
+        elif data:
             self.border_title = (
                 f"Loading Data {self._human_row_count(next(iter(data.values())))}."
             )

tests/conftest.py

@@ -1,25 +1,9 @@
 from pathlib import Path
 
 import pytest
-from textual.app import App
 
 
 @pytest.fixture
 def data_dir() -> Path:
     here = Path(__file__)
     return here.parent / "data"
-
-
-class TestHelpers:
-    @staticmethod
-    async def await_data_loaded(app: App) -> None:
-        # when the query is submitted, it should update app.query_text, app.relation,
-        # and app.data using three different workers.
-        await app.workers.wait_for_complete()
-        await app.workers.wait_for_complete()
-        await app.workers.wait_for_complete()
-
-
-@pytest.fixture
-def helpers() -> TestHelpers:
-    return TestHelpers()

tests/functional_tests/test_app.py

@@ -5,8 +5,6 @@
 from harlequin.tui.components import ExportScreen
 from harlequin.tui.components.results_viewer import ResultsTable
 
-from ..conftest import TestHelpers
-
 
 @pytest.mark.asyncio
 async def test_select_1(app: Harlequin) -> None:
@@ -19,13 +17,9 @@ async def test_select_1(app: Harlequin) -> None:
             await pilot.press(key)
         await pilot.press("ctrl+j")  # alias for ctrl+enter
 
-        # when the query is submitted, it should update app.query_text, app.relation,
-        # and app.data using three different workers.
-        await app.workers.wait_for_complete()
+        await pilot.pause()
         assert app.query_text == q
-        await app.workers.wait_for_complete()
         assert app.relations
-        await app.workers.wait_for_complete()
         assert len(app.results_viewer.data) == 1
         assert app.results_viewer.data[next(iter(app.results_viewer.data))] == [(1,)]
 
@@ -38,7 +32,7 @@ async def test_multiple_queries(app: Harlequin) -> None:
         await pilot.press("ctrl+j")
 
         # should only run one query
-        await TestHelpers.await_data_loaded(app=app)
+        await pilot.pause()
         assert app.query_text == "select 1;"
         assert len(app.results_viewer.data) == 1
         assert app.results_viewer.data[next(iter(app.results_viewer.data))] == [(1,)]
@@ -48,7 +42,7 @@ async def test_multiple_queries(app: Harlequin) -> None:
         await pilot.press("ctrl+a")
         await pilot.press("ctrl+j")
         # should run both queries
-        await TestHelpers.await_data_loaded(app=app)
+        await pilot.pause()
         assert app.query_text == "select 1; select 2"
         assert len(app.results_viewer.data) == 2
         assert "hide-tabs" not in app.results_viewer.classes
@@ -76,7 +70,7 @@ async def test_query_formatting(app: Harlequin) -> None:
 
 
 @pytest.mark.asyncio
-async def test_run_query_bar(app_small_db: Harlequin, helpers: TestHelpers) -> None:
+async def test_run_query_bar(app_small_db: Harlequin) -> None:
     app = app_small_db
     async with app.run_test() as pilot:
         # initialization
@@ -89,14 +83,14 @@ async def test_run_query_bar(app_small_db: Harlequin, helpers: TestHelpers) -> N
         # dataset has 857 records
         app.editor.text = "select * from drivers"
         await pilot.click(bar.button.__class__)
-        await helpers.await_data_loaded(app)
+        await pilot.pause()
         assert len(app.results_viewer.data[next(iter(app.results_viewer.data))]) > 500
 
         # apply a limit by clicking the limit checkbox
         await pilot.click(bar.checkbox.__class__)
         assert bar.checkbox.value is True
         await pilot.click(bar.button.__class__)
-        await helpers.await_data_loaded(app)
+        await pilot.pause()
         assert len(app.results_viewer.data[next(iter(app.results_viewer.data))]) == 500
 
         # type an invalid limit, checkbox should be unchecked
@@ -119,7 +113,7 @@ async def test_run_query_bar(app_small_db: Harlequin, helpers: TestHelpers) -> N
 
         # run the query with a smaller limit
         await pilot.click(bar.button.__class__)
-        await helpers.await_data_loaded(app)
+        await pilot.pause()
         assert len(app.results_viewer.data[next(iter(app.results_viewer.data))]) == 100