Add query cancellation controls (#41)

Co-authored-by: SuLe <sule@SuLedeMacBook-Air.local>
Co-authored-by: skyler <77960507+t8y2@users.noreply.github.com>

src/App.vue

@@ -1,7 +1,7 @@
 <script setup lang="ts">
 import { ref, computed, watch, onMounted, onUnmounted } from "vue";
 import { useI18n } from "vue-i18n";
-import { DatabaseZap, FilePlus2, Play, Loader2, X, Globe, Moon, Sun, Upload, Download, Plus, History, Server, Table2, Database, Search, ShieldCheck, Bot, Pin, AlignLeft, CloudDownload } from "lucide-vue-next";
+import { DatabaseZap, FilePlus2, Play, Loader2, Square, X, Globe, Moon, Sun, Upload, Download, Plus, History, Server, Table2, Database, Search, ShieldCheck, Bot, Pin, AlignLeft, CloudDownload } from "lucide-vue-next";
 import { Splitpanes, Pane } from "splitpanes";
 import "splitpanes/dist/splitpanes.css";
 import { Button } from "@/components/ui/button";
@@ -44,6 +44,7 @@ import { setLocale, currentLocale, type Locale } from "@/i18n";
 import { getCurrentWindow, type Theme } from "@tauri-apps/api/window";
 import { getCurrentWebview } from "@tauri-apps/api/webview";
 import * as api from "@/lib/tauri";
+import { canCancelQueryExecution, queryExecutionLabelKey } from "@/lib/queryExecutionState";
 import { resolveExecutableSql } from "@/lib/sqlExecutionTarget";
 import type { SqlFormatDialect } from "@/lib/sqlFormatter";
 import { isCloseTabShortcut, isExecuteSqlShortcut } from "@/lib/keyboardShortcuts";
@@ -345,6 +346,11 @@ async function doExecute(sql = executableSql.value) {
   });
 }
 
+function cancelActiveExecution() {
+  const tab = activeTab.value;
+  if (tab) void queryStore.cancelTabExecution(tab.id);
+}
+
 function onDangerConfirm() {
   const sql = pendingDangerSql.value || executableSql.value;
   pendingDangerSql.value = "";
@@ -809,12 +815,19 @@ async function setupFileDrop() {
               <div class="flex items-center gap-0.5">
                 <Tooltip>
                   <TooltipTrigger as-child>
-                    <Button variant="ghost" size="icon" class="h-6 w-6" :disabled="activeTab.isExecuting" @click="tryExecute()">
-                      <Loader2 v-if="activeTab.isExecuting" class="h-3.5 w-3.5 animate-spin" />
+                    <Button
+                      :variant="activeTab.isExecuting ? 'destructive' : 'ghost'"
+                      size="icon"
+                      class="h-6 w-6"
+                      :disabled="activeTab.isCancelling || (!activeTab.isExecuting && !executableSql.trim())"
+                      @click="activeTab.isExecuting ? cancelActiveExecution() : tryExecute()"
+                    >
+                      <Loader2 v-if="activeTab.isCancelling" class="h-3.5 w-3.5 animate-spin" />
+                      <Square v-else-if="activeTab.isExecuting" class="h-3.5 w-3.5 fill-current" />
                       <Play v-else class="h-3.5 w-3.5" />
                     </Button>
                   </TooltipTrigger>
-                  <TooltipContent>{{ t('toolbar.executeShortcut') }}</TooltipContent>
+                  <TooltipContent>{{ activeTab.isExecuting ? t('toolbar.stopQuery') : t('toolbar.executeShortcut') }}</TooltipContent>
                 </Tooltip>
                 <Tooltip>
                   <TooltipTrigger as-child>
@@ -908,6 +921,12 @@ async function setupFileDrop() {
                 <Pane :size="60" :min-size="20">
                   <div class="h-full flex flex-col">
                     <DataGrid v-if="activeTab.result" :key="activeTab.id" class="flex-1 min-h-0" :result="activeTab.result" :sql="activeTab.lastExecutedSql || activeTab.sql" />
+                    <div v-else-if="activeTab.isExecuting" class="flex-1 min-h-0 flex flex-col items-center justify-center gap-3 text-muted-foreground text-sm">
+                      <div class="flex items-center">
+                        <Loader2 class="h-5 w-5 animate-spin mr-2" />
+                        {{ t(queryExecutionLabelKey(activeTab)) }}
+                      </div>
+                    </div>
                     <div v-else class="flex-1 min-h-0 flex items-center justify-center text-muted-foreground text-sm">
                       {{ t('editor.pressToExecute') }}
                     </div>
@@ -934,8 +953,22 @@ async function setupFileDrop() {
                   @paginate="onPaginate"
                   @sort="onSort"
                 />
-                <div v-else-if="activeTab.isExecuting" class="h-full flex items-center justify-center text-muted-foreground text-sm">
-                  <Loader2 class="h-5 w-5 animate-spin mr-2" /> {{ t('common.loading') }}
+                <div v-else-if="activeTab.isExecuting" class="h-full flex flex-col items-center justify-center gap-3 text-muted-foreground text-sm">
+                  <div class="flex items-center">
+                    <Loader2 class="h-5 w-5 animate-spin mr-2" />
+                    {{ t(queryExecutionLabelKey(activeTab)) }}
+                  </div>
+                  <Button
+                    variant="destructive"
+                    size="sm"
+                    class="h-7 gap-1.5"
+                    :disabled="!canCancelQueryExecution(activeTab)"
+                    @click="cancelActiveExecution"
+                  >
+                    <Loader2 v-if="activeTab.isCancelling" class="h-3.5 w-3.5 animate-spin" />
+                    <Square v-else class="h-3.5 w-3.5 fill-current" />
+                    {{ t('toolbar.stopQuery') }}
+                  </Button>
                 </div>
               </div>
             </template>

src/i18n/locales/en.ts

@@ -7,6 +7,7 @@ export default {
     newQuery: "New Query",
     execute: "Execute",
     executeShortcut: "Execute selection/query (Cmd+Enter)",
+    stopQuery: "Stop query",
     formatSql: "Format SQL",
     formatSqlFailed: "Failed to format SQL",
   },
@@ -123,6 +124,7 @@ export default {
   common: {
     language: "Language",
     loading: "Loading...",
+    stopping: "Stopping...",
     close: "Close",
   },
   ai: {

src/i18n/locales/zh-CN.ts

@@ -7,6 +7,7 @@ export default {
     newQuery: "新建查询",
     execute: "执行",
     executeShortcut: "执行选中/全部 (Cmd+Enter)",
+    stopQuery: "停止执行",
     formatSql: "格式化 SQL",
     formatSqlFailed: "SQL 格式化失败",
   },
@@ -123,6 +124,7 @@ export default {
   common: {
     language: "语言",
     loading: "加载中...",
+    stopping: "正在停止...",
     close: "关闭",
   },
   ai: {

src/lib/queryExecutionState.ts

@@ -0,0 +1,13 @@
+export interface QueryExecutionStateLike {
+  isExecuting: boolean;
+  isCancelling?: boolean;
+  executionId?: string;
+}
+
+export function canCancelQueryExecution(state: QueryExecutionStateLike): boolean {
+  return state.isExecuting && !!state.executionId && !state.isCancelling;
+}
+
+export function queryExecutionLabelKey(state: Pick<QueryExecutionStateLike, "isCancelling">): "common.loading" | "common.stopping" {
+  return state.isCancelling ? "common.stopping" : "common.loading";
+}

src/lib/tauri.ts

@@ -95,8 +95,12 @@ export async function getColumns(
   return invoke("get_columns", { connectionId, database, schema, table });
 }
 
-export async function executeQuery(connectionId: string, database: string, sql: string): Promise<QueryResult> {
-  return invoke("execute_query", { connectionId, database, sql });
+export async function executeQuery(connectionId: string, database: string, sql: string, executionId?: string): Promise<QueryResult> {
+  return invoke("execute_query", { connectionId, database, sql, executionId });
+}
+
+export async function cancelQuery(executionId: string): Promise<boolean> {
+  return invoke("cancel_query", { executionId });
 }
 
 export async function executeBatch(connectionId: string, database: string, statements: string[]): Promise<QueryResult> {

src/stores/queryStore.ts

@@ -2,6 +2,7 @@ import { defineStore } from "pinia";
 import { ref } from "vue";
 import type { QueryTab } from "@/types/database";
 import { orderPinnedFirst } from "@/lib/pinnedItems";
+import { canCancelQueryExecution } from "@/lib/queryExecutionState";
 import { closeAllTabsState, closeOtherTabsState } from "@/lib/tabCloseActions";
 import * as api from "@/lib/tauri";
 
@@ -31,6 +32,7 @@ export const useQueryStore = defineStore("query", () => {
       database,
       sql: "",
       isExecuting: false,
+      isCancelling: false,
       mode,
     };
     tabs.value.push(tab);
@@ -41,6 +43,7 @@ export const useQueryStore = defineStore("query", () => {
   function closeTab(id: string) {
     const idx = tabs.value.findIndex((t) => t.id === id);
     if (idx < 0) return;
+    if (tabs.value[idx].isExecuting) void cancelTabExecution(id);
     tabs.value.splice(idx, 1);
     if (activeTabId.value === id) {
       activeTabId.value = tabs.value[Math.min(idx, tabs.value.length - 1)]?.id ?? null;
@@ -48,12 +51,18 @@ export const useQueryStore = defineStore("query", () => {
   }
 
   function closeOtherTabs(id: string) {
+    tabs.value
+      .filter((tab) => tab.id !== id && tab.isExecuting)
+      .forEach((tab) => void cancelTabExecution(tab.id));
     const next = closeOtherTabsState(tabs.value, activeTabId.value, id);
     tabs.value = next.tabs;
     activeTabId.value = next.activeTabId;
   }
 
   function closeAllTabs() {
+    tabs.value
+      .filter((tab) => tab.isExecuting)
+      .forEach((tab) => void cancelTabExecution(tab.id));
     const next = closeAllTabsState(tabs.value, activeTabId.value);
     tabs.value = next.tabs;
     activeTabId.value = next.activeTabId;
@@ -97,7 +106,12 @@ export const useQueryStore = defineStore("query", () => {
 
   function setExecuting(id: string, isExecuting: boolean) {
     const tab = tabs.value.find((t) => t.id === id);
-    if (tab) tab.isExecuting = isExecuting;
+    if (!tab) return;
+    tab.isExecuting = isExecuting;
+    if (!isExecuting) {
+      tab.isCancelling = false;
+      tab.executionId = undefined;
+    }
   }
 
   function toErrorResult(e: any): NonNullable<QueryTab["result"]> {
@@ -114,6 +128,8 @@ export const useQueryStore = defineStore("query", () => {
     if (!tab) return;
     tab.result = toErrorResult(e);
     tab.isExecuting = false;
+    tab.isCancelling = false;
+    tab.executionId = undefined;
   }
 
   async function executeCurrentTab() {
@@ -132,18 +148,53 @@ export const useQueryStore = defineStore("query", () => {
     const tab = tabs.value.find((t) => t.id === id);
     if (!tab || !sql.trim()) return;
 
+    const executionId = crypto.randomUUID();
     tab.isExecuting = true;
+    tab.isCancelling = false;
+    tab.executionId = executionId;
     tab.lastExecutedSql = sql;
     try {
-      tab.result = await api.executeQuery(tab.connectionId, tab.database, sql);
+      tab.result = await api.executeQuery(tab.connectionId, tab.database, sql, executionId);
     } catch (e: any) {
-      tab.result = toErrorResult(e);
+      const current = tabs.value.find((t) => t.id === id);
+      if (current?.executionId === executionId) {
+        current.result = toErrorResult(e);
+      }
     } finally {
-      tab.isExecuting = false;
+      const current = tabs.value.find((t) => t.id === id);
+      if (current?.executionId === executionId) {
+        current.isExecuting = false;
+        current.isCancelling = false;
+        current.executionId = undefined;
+      }
     }
     trimResultCache();
   }
 
+  async function cancelTabExecution(id: string) {
+    const tab = tabs.value.find((t) => t.id === id);
+    if (!tab || !canCancelQueryExecution(tab)) return false;
+
+    const executionId = tab.executionId;
+    if (!executionId) return false;
+    tab.isCancelling = true;
+    try {
+      const canceled = await api.cancelQuery(executionId);
+      if (!canceled) {
+        const current = tabs.value.find((t) => t.id === id);
+        if (current && current.executionId === executionId) current.isCancelling = false;
+      }
+      return canceled;
+    } catch (e: any) {
+      const current = tabs.value.find((t) => t.id === id);
+      if (current && current.executionId === executionId) {
+        current.isCancelling = false;
+        current.result = toErrorResult(e);
+      }
+      return false;
+    }
+  }
+
   function trimResultCache() {
     const inactive = tabs.value.filter((t) => t.id !== activeTabId.value && t.result);
     if (inactive.length > MAX_CACHED_RESULTS) {
@@ -169,5 +220,6 @@ export const useQueryStore = defineStore("query", () => {
     executeCurrentTab,
     executeCurrentSql,
     executeTabSql,
+    cancelTabExecution,
   };
 });

src/types/database.ts

@@ -102,6 +102,8 @@ export interface QueryTab {
   pinned?: boolean;
   result?: QueryResult;
   isExecuting: boolean;
+  isCancelling?: boolean;
+  executionId?: string;
   mode: "data" | "query" | "redis" | "mongo";
   tableMeta?: {
     schema?: string;