feat: cancel queries; impl for duckdb and sqlite (#629)

* feat: cancel queries; impl for duckdb and sqlite

* fix: update keys app snapshots

CHANGELOG.md

@@ -6,7 +6,11 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
+- For adapters that support canceling queries, Harlequin will now display a "Cancel Query" button while queries are running.
+- Two new Actions are available for key bindings: `run_query`  (with a global app scope, in addition to the existing `code_editor.run_query` action) and `cancel_query`.
 - Adapters may now implement a `connection_id` property to improve Harlequin's ability to persist the data catalog and query history across Harlequin invocations ([#410](https://github.com/tconbeer/harlequin/issues/410)).
+- Adapters may now implement a `HarlequinConnection.cancel()` method to cancel all in-flight queries ([#333](https://github.com/tconbeer/harlequin/issues/333)).
+- Queries can now be canceled when using the DuckDB or SQLite adapters.
 
 ### Bug Fixes
 

src/harlequin/actions.py

@@ -79,6 +79,10 @@ class Action:
     "refresh_catalog": Action(
         target=None, action="refresh_catalog", description="Refresh Data Catalog"
     ),
+    "run_query": Action(target=None, action="run_query", description="Run Query"),
+    "cancel_query": Action(
+        target=None, action="cancel_query", description="Cancel Query"
+    ),
     #######################################################
     # CodeEditor ACTIONS
     #######################################################

src/harlequin/adapter.py

@@ -91,6 +91,15 @@ def execute(self, query: str) -> HarlequinCursor | None:
         """
         pass
 
+    def cancel(self) -> None:
+        """
+        Cancels/interrupts all in-flight queries previously executed with `execute`.
+
+        After implementing this method, set the adapter class variable
+        IMPLEMENTS_CANCEL to True to show the cancel button in the Harlequin UI.
+        """
+        return None
+
     @abstractmethod
     def get_catalog(self) -> Catalog:
         """
@@ -203,6 +212,7 @@ class variable.
     ADAPTER_OPTIONS: list[HarlequinAdapterOption] | None = None
     COPY_FORMATS: list[HarlequinCopyFormat] | None = None
     """DEPRECATED. Adapter Copy formats are now ignored by Harlequin."""
+    IMPLEMENTS_CANCEL = False
 
     @abstractmethod
     def __init__(self, conn_str: Sequence[str], **options: Any) -> None:

src/harlequin/app.py

@@ -122,6 +122,10 @@ def __init__(
         self.ddl_queries = ddl_queries
 
 
+class QueriesCanceled(Message):
+    pass
+
+
 class ResultsFetched(Message):
     def __init__(
         self,
@@ -229,7 +233,9 @@ def compose(self) -> ComposeResult:
             type_color=self.app_colors.gray,
         )
         self.run_query_bar = RunQueryBar(
-            max_results=self.max_results, classes="non-responsive"
+            max_results=self.max_results,
+            classes="non-responsive",
+            show_cancel_button=self.adapter.IMPLEMENTS_CANCEL,
         )
         self.footer = Footer()
 
@@ -294,6 +300,11 @@ def submit_query_from_run_query_bar(self, message: Button.Pressed) -> None:
             )
         )
 
+    @on(Button.Pressed, "#cancel_query")
+    def cancel_query(self, message: Button.Pressed) -> None:
+        message.stop()
+        self.action_cancel_query()
+
     @on(Button.Pressed, "#transaction_button")
     def handle_transaction_button_press(self, message: Button.Pressed) -> None:
         message.stop()
@@ -534,6 +545,12 @@ def fetch_data_or_reset_table(self, message: QueriesExecuted) -> None:
             )
             self.update_schema_data()
 
+    @on(QueriesCanceled)
+    def reset_after_cancel(self) -> None:
+        self.run_query_bar.set_responsive()
+        self.results_viewer.show_table(did_run=False)
+        self.notify("Queries canceled.", severity="error")
+
     @on(ResultsFetched)
     async def load_tables(self, message: ResultsFetched) -> None:
         for id_, (cols, data, query_text) in message.data.items():
@@ -725,6 +742,14 @@ def action_bind_keymaps(self, *keymap_names: str) -> None:
                 pretty_print_error(e)
                 self.exit(return_code=2)
 
+    async def action_run_query(self) -> None:
+        if self.editor is None:
+            return
+        await self.editor.action_submit()
+
+    def action_cancel_query(self) -> None:
+        self._cancel_query()
+
     def action_export(self) -> None:
         show_export_error = partial(
             self._push_error_modal,
@@ -896,6 +921,19 @@ def _execute_query(self, message: QuerySubmitted) -> None:
             )
         )
 
+    @work(
+        thread=True,
+        exclusive=True,
+        exit_on_error=True,
+        group="query_cancellers",
+        description="Cancelling queries.",
+    )
+    def _cancel_query(self) -> None:
+        if self.connection is None or not self.adapter.IMPLEMENTS_CANCEL:
+            return
+        self.connection.cancel()
+        self.post_message(QueriesCanceled())
+
     def _get_query_text(self) -> str:
         if self.editor is None:
             return ""
@@ -935,7 +973,7 @@ def _fetch_data(
         for id_, (cur, q) in cursors.items():
             try:
                 cur_data = cur.fetchall()
-            except HarlequinQueryError as e:
+            except BaseException as e:
                 errors.append((e, q))
             else:
                 data[id_] = (cur.columns(), cur_data, q)

src/harlequin/app.tcss

@@ -160,14 +160,30 @@ RunQueryBar Button {
 RunQueryBar Button#run_query {
     background: $primary;
     margin: 0 0 0 4;
-}
-
-RunQueryBar Button#run_query:hover {
-    background: $secondary;
-}
-
-RunQueryBar Button#run_query:focus {
-    text-style: reverse;
+    &:hover {
+        background: $secondary;
+    }
+    &:focus {
+        text-style: reverse;
+    }
+    &.hidden {
+        display: none;
+    }
+}
+
+RunQueryBar Button#cancel_query {
+    background: $error;
+    margin: 0 0 0 4;
+    text-style: bold;
+    &:hover {
+        background: $secondary;
+    }
+    &:focus {
+        text-style: reverse;
+    }
+    &.hidden {
+        display: none;
+    }
 }
 
 RunQueryBar Button#transaction_button,
@@ -204,10 +220,6 @@ RunQueryBar Button#rollback_button:hover, {
     color: $secondary;
 }
 
-RunQueryBar Button#run_query:focus {
-    text-style: reverse;
-}
-
 RunQueryBar Checkbox {
     border: none;
     padding: 0;

src/harlequin/components/run_query_bar.py

@@ -19,8 +19,10 @@ def __init__(
         classes: Union[str, None] = None,
         disabled: bool = False,
         max_results: int = 10_000,
+        show_cancel_button: bool = False,
     ) -> None:
         self.max_results = max_results
+        self.show_cancel_button = show_cancel_button
         super().__init__(
             *children, name=name, id=id, classes=classes, disabled=disabled
         )
@@ -49,6 +51,8 @@ def compose(self) -> ComposeResult:
             ),
         )
         self.run_button = Button("Run Query", id="run_query")
+        self.cancel_button = Button("Cancel Query", id="cancel_query")
+        self.cancel_button.add_class("hidden")
         with Horizontal(id="transaction_buttons"):
             yield self.transaction_button
             yield self.commit_button
@@ -57,6 +61,7 @@ def compose(self) -> ComposeResult:
             yield self.limit_checkbox
             yield self.limit_input
             yield self.run_button
+            yield self.cancel_button
 
     def on_mount(self) -> None:
         if self.app.is_headless:
@@ -87,6 +92,14 @@ def limit_value(self) -> int | None:
 
     def set_not_responsive(self) -> None:
         self.add_class("non-responsive")
+        if self.show_cancel_button:
+            with self.app.batch_update():
+                self.run_button.add_class("hidden")
+                self.cancel_button.remove_class("hidden")
 
     def set_responsive(self) -> None:
         self.remove_class("non-responsive")
+        if self.show_cancel_button:
+            with self.app.batch_update():
+                self.run_button.remove_class("hidden")
+                self.cancel_button.add_class("hidden")

src/harlequin_duckdb/adapter.py

@@ -47,6 +47,8 @@ def set_limit(self, limit: int) -> HarlequinCursor:
     def fetchall(self) -> AutoBackendType | None:
         try:
             result = self.relation.fetch_arrow_table()
+        except duckdb.InterruptException:
+            return None
         except duckdb.Error as e:
             raise HarlequinQueryError(
                 msg=str(e), title="DuckDB raised an error when running your query:"
@@ -105,6 +107,8 @@ def __init__(self, conn: duckdb.DuckDBPyConnection, init_message: str = "") -> N
     def execute(self, query: str) -> DuckDbCursor | None:
         try:
             rel = self.conn.sql(query)
+        except duckdb.InterruptException:
+            return None
         except duckdb.Error as e:
             raise HarlequinQueryError(
                 msg=str(e),
@@ -116,6 +120,9 @@ def execute(self, query: str) -> DuckDbCursor | None:
         else:
             return None
 
+    def cancel(self) -> None:
+        self.conn.interrupt()
+
     def get_catalog(self) -> Catalog:
         catalog_items: list[CatalogItem] = []
         databases = self._get_databases()
@@ -267,6 +274,7 @@ def _short_column_type(cls, native_type: DuckDBPyType | str) -> str:
 class DuckDbAdapter(HarlequinAdapter):
     ADAPTER_OPTIONS = DUCKDB_OPTIONS
     COPY_FORMATS = None
+    IMPLEMENTS_CANCEL = True
 
     def __init__(
         self,

src/harlequin_sqlite/adapter.py

@@ -30,7 +30,10 @@ def __init__(self, conn: HarlequinSqliteConnection, cur: sqlite3.Cursor) -> None
         self.conn = conn
         self.cur = cur
         self._limit: int | None = None
-        _first_row = cur.fetchone()
+        try:
+            _first_row = cur.fetchone()
+        except sqlite3.Error:  # maybe canceled query here
+            _first_row = None
         self.has_records = _first_row is not None
         self._first_row: tuple[Any, ...] = _first_row or tuple(
             [None] * len(cur.description)
@@ -50,11 +53,21 @@ def set_limit(self, limit: int) -> "HarlequinSqliteCursor":
 
     def fetchall(self) -> AutoBackendType | None:
         if self.has_records:
-            remaining_rows = (
-                self.cur.fetchall()
-                if self._limit is None
-                else self.cur.fetchmany(self._limit - 1)
-            )
+            try:
+                remaining_rows = (
+                    self.cur.fetchall()
+                    if self._limit is None
+                    else self.cur.fetchmany(self._limit - 1)
+                )
+            except sqlite3.OperationalError:  # maybe canceled here
+                return None
+            except sqlite3.Error as e:
+                raise HarlequinQueryError(
+                    msg=str(e),
+                    title=(
+                        "SQLite raised an error when fetching results for your query:"
+                    ),
+                ) from e
             return [self._first_row, *remaining_rows]
         else:
             return None
@@ -109,6 +122,9 @@ def execute(self, query: str) -> HarlequinSqliteCursor | None:
         else:
             return None
 
+    def cancel(self) -> None:
+        self.conn.interrupt()
+
     def get_catalog(self) -> Catalog:
         catalog_items: list[CatalogItem] = []
         databases = self._get_databases()
@@ -251,6 +267,7 @@ def validate_sql(self, text: str) -> str:
 class HarlequinSqliteAdapter(HarlequinAdapter):
     ADAPTER_OPTIONS: list[HarlequinAdapterOption] | None = SQLITE_OPTIONS
     COPY_FORMATS: list[HarlequinCopyFormat] | None = None
+    IMPLEMENTS_CANCEL = True
 
     def __init__(
         self,