From 19b5b70fb7c76cde50e1e8f2b64df1efbdf79d8f Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Fri, 12 Jun 2026 01:09:41 +0800
Subject: [PATCH] feat: confirm before closing unsaved SQL query tabs

---
 apps/desktop/src/App.vue                      | 61 ++++++++++++++++++-
 .../src/components/layout/AppTabBar.vue       | 42 ++++++++++++-
 apps/desktop/src/i18n/locales/en.ts           |  3 +
 apps/desktop/src/i18n/locales/zh-CN.ts        |  3 +
 apps/desktop/src/stores/queryStore.ts         | 52 +++++++++++++++-
 apps/desktop/src/types/database.ts            |  1 +
 6 files changed, 157 insertions(+), 5 deletions(-)

diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue
index 33dc49896..dc3099066 100644
--- a/apps/desktop/src/App.vue
+++ b/apps/desktop/src/App.vue
@@ -121,6 +121,8 @@ const newQueryContextSource = ref<"tab" | "sidebar">("tab");
 const showSaveSqlDialog = ref(false);
 const saveSqlName = ref("");
 const saveSqlFolderId = ref("");
+const pendingSaveAndCloseTabId = ref<string | null>(null);
+const pendingPrevActiveTabId = ref<string | null>(null);
 const ROOT_SAVED_SQL_FOLDER = "__root__";
 
 const activeTab = computed(() => queryStore.tabs.find((t) => t.id === queryStore.activeTabId));
@@ -390,6 +392,36 @@ function defaultSavedSqlName(title: string) {
   return normalized.endsWith(".sql") ? normalized : `${normalized}.sql`;
 }
 
+async function handleSaveTab(tabId: string) {
+  const tab = queryStore.tabs.find((t) => t.id === tabId);
+  if (!tab || !tab.sql.trim()) return;
+  const existing = tab.savedSqlId ? savedSqlStore.getFile(tab.savedSqlId) : undefined;
+  if (existing) {
+    const updated = await savedSqlStore.saveFile({
+      id: existing.id,
+      connectionId: tab.connectionId,
+      folderId: existing.folderId,
+      name: existing.name,
+      database: tab.database,
+      schema: tab.schema,
+      sql: tab.sql,
+    });
+    queryStore.linkSavedSql(tab.id, updated.id, updated.name);
+    queryStore.markTabClean(tab);
+    toast(t("savedSql.saved"), 2000);
+    queryStore.closeTab(tabId, { force: true });
+    return;
+  }
+  // No existing saved SQL — open save dialog, then close after save
+  const prevActive = queryStore.activeTabId;
+  queryStore.activeTabId = tabId;
+  saveSqlName.value = defaultSavedSqlName(tab.title);
+  saveSqlFolderId.value = ROOT_SAVED_SQL_FOLDER;
+  pendingSaveAndCloseTabId.value = tabId;
+  pendingPrevActiveTabId.value = prevActive;
+  showSaveSqlDialog.value = true;
+}
+
 async function openSaveSqlDialog() {
   const tab = activeTab.value;
   if (!tab || !tab.sql.trim()) return;
@@ -459,7 +491,15 @@ async function confirmSaveSqlToLibrary() {
       sql: tab.sql,
     });
     queryStore.linkSavedSql(tab.id, saved.id, saved.name);
+    queryStore.markTabClean(tab);
     showSaveSqlDialog.value = false;
+    if (pendingSaveAndCloseTabId.value) {
+      const closeId = pendingSaveAndCloseTabId.value;
+      pendingSaveAndCloseTabId.value = null;
+      if (pendingPrevActiveTabId.value) queryStore.activeTabId = pendingPrevActiveTabId.value;
+      pendingPrevActiveTabId.value = null;
+      queryStore.closeTab(closeId, { force: true });
+    }
     toast(t("savedSql.saved"), 2000);
   } catch (e: any) {
     toast(t("savedSql.saveFailed", { message: e?.message || String(e) }), 5000);
@@ -1056,7 +1096,7 @@ onUnmounted(() => {
 
           <div :class="isClassicLayout ? 'flex-1 min-w-0 overflow-hidden' : 'flex-1 min-w-0 overflow-hidden rounded-md border border-border/80 bg-background'">
             <div class="h-full flex flex-col min-w-0">
-              <AppTabBar :show-driver-store="showDriverStore" :agent-driver-update-count="toolbarAgentDriverUpdateCount" @toggle-driver-store="showDriverStore = true" @close-driver-store="showDriverStore = false" />
+              <AppTabBar :show-driver-store="showDriverStore" :agent-driver-update-count="toolbarAgentDriverUpdateCount" @toggle-driver-store="showDriverStore = true" @close-driver-store="showDriverStore = false" @save-tab="handleSaveTab" />
               <DriverStorePage v-if="showDriverStore" class="flex-1 min-h-0" :update-notifications-enabled="updateNotificationsEnabled" @update-count-change="updateAgentDriverUpdateCount" />
               <div v-else-if="activeTab" class="flex flex-col flex-1 min-h-0">
                 <EditorToolbar
@@ -1219,7 +1259,15 @@ onUnmounted(() => {
         </Transition>
       </div>
 
-      <Dialog v-model:open="showSaveSqlDialog">
+      <Dialog
+        :open="showSaveSqlDialog"
+        @update:open="
+          (open: boolean) => {
+            showSaveSqlDialog = open;
+            if (!open) pendingSaveAndCloseTabId = null;
+          }
+        "
+      >
         <DialogContent class="sm:max-w-[420px]">
           <DialogHeader>
             <DialogTitle>{{ t("savedSql.saveToLibrary") }}</DialogTitle>
@@ -1245,7 +1293,14 @@ onUnmounted(() => {
             </div>
           </div>
           <DialogFooter>
-            <Button variant="outline" @click="showSaveSqlDialog = false">{{ t("dangerDialog.cancel") }}</Button>
+            <Button
+              variant="outline"
+              @click="
+                showSaveSqlDialog = false;
+                pendingSaveAndCloseTabId = null;
+              "
+              >{{ t("dangerDialog.cancel") }}</Button
+            >
             <Button :disabled="!saveSqlName.trim()" @click="confirmSaveSqlToLibrary">{{ t("savedSql.save") }}</Button>
           </DialogFooter>
         </DialogContent>
diff --git a/apps/desktop/src/components/layout/AppTabBar.vue b/apps/desktop/src/components/layout/AppTabBar.vue
index f0d69f69e..69e1ffcba 100644
--- a/apps/desktop/src/components/layout/AppTabBar.vue
+++ b/apps/desktop/src/components/layout/AppTabBar.vue
@@ -2,10 +2,12 @@
 import { computed, ref, watch, nextTick } from "vue";
 import type { CSSProperties } from "vue";
 import { useI18n } from "vue-i18n";
-import { X, Pin, ChevronDown, Table2, Code2, TableProperties, PencilRuler, KeyRound, Pencil, Package, Check, Lock, Copy } from "@lucide/vue";
+import { X, Pin, ChevronDown, Table2, Code2, TableProperties, PencilRuler, KeyRound, Pencil, Package, Check, Lock, Copy, AlertTriangle } from "@lucide/vue";
 import CustomContextMenu, { type ContextMenuItem } from "@/components/ui/CustomContextMenu.vue";
 import LightDropdown from "@/components/ui/LightDropdown.vue";
 import { Tooltip, TooltipTrigger, TooltipContent } from "@/components/ui/tooltip";
+import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Button } from "@/components/ui/button";
 import { useQueryStore } from "@/stores/queryStore";
 import { useSettingsStore } from "@/stores/settingsStore";
 import { useTabScroll } from "@/composables/useTabScroll";
@@ -22,6 +24,7 @@ const props = defineProps<{
 const emit = defineEmits<{
   "toggle-driver-store": [];
   "close-driver-store": [];
+  "save-tab": [tabId: string];
 }>();
 
 const { t } = useI18n();
@@ -113,6 +116,19 @@ function getTabMenuItems(tab: QueryTab): ContextMenuItem[] {
   ];
 }
 
+function handleSaveAndClose() {
+  const id = queryStore.saveAndClosePendingTab();
+  if (id) emit("save-tab", id);
+}
+
+function handleDiscardAndClose() {
+  queryStore.forceClosePendingTab();
+}
+
+function handleCancelClose() {
+  queryStore.cancelClosePendingTab();
+}
+
 const tabsContainerRef = ref<HTMLElement | null>(null);
 const { hasTabOverflow, scrollThumbLeftPercent, scrollThumbWidthPercent, isScrollbarDragging, updateScrollButtons, onTabsWheel, startScrollbarDrag } = useTabScroll(tabsContainerRef);
 const tabScrollBehavior = ref<ScrollBehavior>("smooth");
@@ -387,6 +403,30 @@ function activateTab(tabId: string) {
       />
     </div>
   </div>
+
+  <Dialog
+    :open="queryStore.showCloseConfirm"
+    @update:open="
+      (open) => {
+        if (!open) queryStore.cancelClosePendingTab();
+      }
+    "
+  >
+    <DialogContent class="sm:max-w-md">
+      <DialogHeader>
+        <DialogTitle class="flex items-center gap-2">
+          <AlertTriangle class="h-5 w-5 text-amber-500" />
+          {{ t("editor.unsavedChangesTitle") }}
+        </DialogTitle>
+      </DialogHeader>
+      <p class="text-sm text-muted-foreground">{{ t("editor.unsavedChangesMessage") }}</p>
+      <DialogFooter>
+        <Button variant="outline" @click="handleCancelClose">{{ t("common.cancel") }}</Button>
+        <Button variant="secondary" @click="handleDiscardAndClose">{{ t("editor.discardChanges") }}</Button>
+        <Button @click="handleSaveAndClose">{{ t("savedSql.save") }}</Button>
+      </DialogFooter>
+    </DialogContent>
+  </Dialog>
 </template>
 
 <style scoped>
diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts
index c429e18e9..369540db8 100644
--- a/apps/desktop/src/i18n/locales/en.ts
+++ b/apps/desktop/src/i18n/locales/en.ts
@@ -347,6 +347,9 @@
     showResultsPane: "Show results",
     hideResultsPane: "Hide results",
     noDatabase: "No database selected",
+    unsavedChangesTitle: "Unsaved changes",
+    unsavedChangesMessage: "This tab has unsaved changes that will be lost if closed. Save before closing?",
+    discardChanges: "Discard",
     selectConnection: "Select connection",
     searchConnection: "Search connections...",
     selectDatabase: "Select database",
diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts
index a1dfeb0aa..7efc32e35 100644
--- a/apps/desktop/src/i18n/locales/zh-CN.ts
+++ b/apps/desktop/src/i18n/locales/zh-CN.ts
@@ -348,6 +348,9 @@
     showResultsPane: "显示结果",
     hideResultsPane: "收起结果",
     noDatabase: "未选择数据库",
+    unsavedChangesTitle: "未保存的更改",
+    unsavedChangesMessage: "此标签页有未保存的更改，关闭后将丢失。是否保存后再关闭？",
+    discardChanges: "不保存",
     selectConnection: "选择连接",
     searchConnection: "搜索连接...",
     selectDatabase: "选择数据库",
diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts
index 5184a0e9c..7c799856b 100644
--- a/apps/desktop/src/stores/queryStore.ts
+++ b/apps/desktop/src/stores/queryStore.ts
@@ -113,6 +113,8 @@ export const useQueryStore = defineStore("query", () => {
   const restored = loadSavedTabs();
   const tabs = ref<QueryTab[]>(restored.tabs);
   const activeTabId = ref<string | null>(restored.activeTabId);
+  const showCloseConfirm = ref(false);
+  const pendingCloseTabId = ref<string | null>(null);
   for (const tab of restored.tabs) {
     if (tab.mode === "data") void deleteTabResultSnapshot(tabResultCacheKey(tab.id));
   }
@@ -263,6 +265,7 @@ export const useQueryStore = defineStore("query", () => {
       isExplaining: false,
       mode,
     };
+    if (mode === "query") tab.originalSql = "";
     tabs.value.push(tab);
     activeTabId.value = id;
     return id;
@@ -353,7 +356,26 @@ export const useQueryStore = defineStore("query", () => {
     return id;
   }
 
-  function closeTab(id: string) {
+  function isTabDirty(tab: QueryTab): boolean {
+    if (tab.mode !== "query") return false;
+    if (!tab.sql.trim()) return false;
+    const original = tab.originalSql;
+    if (original === undefined) return !!tab.savedSqlId;
+    return tab.sql !== original;
+  }
+
+  function markTabClean(tab: QueryTab | undefined) {
+    if (tab) tab.originalSql = tab.sql;
+  }
+
+  function closeTab(id: string, { force = false }: { force?: boolean } = {}) {
+    const tab = tabs.value.find((t) => t.id === id);
+    if (!tab) return;
+    if (!force && isTabDirty(tab)) {
+      pendingCloseTabId.value = id;
+      showCloseConfirm.value = true;
+      return;
+    }
     const idx = tabs.value.findIndex((t) => t.id === id);
     if (idx < 0) return;
     clearDataGridPendingSnapshotsForTab(id);
@@ -368,6 +390,26 @@ export const useQueryStore = defineStore("query", () => {
     }
   }
 
+  function forceClosePendingTab() {
+    const id = pendingCloseTabId.value;
+    pendingCloseTabId.value = null;
+    showCloseConfirm.value = false;
+    if (id) closeTab(id, { force: true });
+  }
+
+  function cancelClosePendingTab() {
+    pendingCloseTabId.value = null;
+    showCloseConfirm.value = false;
+  }
+
+  function saveAndClosePendingTab() {
+    const id = pendingCloseTabId.value;
+    pendingCloseTabId.value = null;
+    showCloseConfirm.value = false;
+    if (id) return id;
+    return null;
+  }
+
   function closeOtherTabs(id: string) {
     tabs.value
       .filter((tab) => tab.id !== id)
@@ -568,6 +610,7 @@ export const useQueryStore = defineStore("query", () => {
       schema: file.schema,
       sql: file.sql,
       savedSqlId: file.id,
+      originalSql: file.sql,
       isExecuting: false,
       isCancelling: false,
       isExplaining: false,
@@ -1620,8 +1663,15 @@ export const useQueryStore = defineStore("query", () => {
   return {
     tabs,
     activeTabId,
+    showCloseConfirm,
+    pendingCloseTabId,
     createTab,
     closeTab,
+    forceClosePendingTab,
+    cancelClosePendingTab,
+    saveAndClosePendingTab,
+    isTabDirty,
+    markTabClean,
     closeOtherTabs,
     closeAllTabs,
     duplicateTab,
diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts
index e54766a5a..efdf1562e 100644
--- a/apps/desktop/src/types/database.ts
+++ b/apps/desktop/src/types/database.ts
@@ -395,6 +395,7 @@ export interface QueryTab {
   schema?: string;
   sql: string;
   savedSqlId?: string;
+  originalSql?: string;
   lastExecutedSql?: string;
   resultBaseSql?: string;
   resultSortedSql?: string;