From 7511753f71fcc0f6aabca768aa67e12c41e8f8d7 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Thu, 18 Jun 2026 16:15:58 +0800
Subject: [PATCH] fix(desktop): handle object source tab save close

---
 apps/desktop/src/App.vue                      | 15 +++++++++---
 .../src/components/sidebar/TreeItem.vue       |  1 +
 packages/app-tests/queryStore.test.ts         | 24 +++++++++++++++++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue
index 4f5ab45fd..c2f6b9ece 100644
--- a/apps/desktop/src/App.vue
+++ b/apps/desktop/src/App.vue
@@ -40,7 +40,7 @@ import { uuid } from "@/lib/utils";
 import { isTauriRuntime } from "@/lib/tauriRuntime";
 import { openQueryResultArchiveFile } from "@/lib/queryResultArchiveFile";
 import { sqlFileTitleFromPath } from "@/lib/sqlFileOpen";
-import type { ConnectionConfig } from "@/types/database";
+import type { ConnectionConfig, QueryTab } from "@/types/database";
 import { parseConnectionDeepLink, type ConnectionDeepLinkDraft } from "@/lib/connectionDeepLink";
 import {
   isBrowserReloadShortcut,
@@ -404,6 +404,11 @@ function defaultSavedSqlName(title: string) {
 async function handleSaveTab(tabId: string) {
   const tab = queryStore.tabs.find((t) => t.id === tabId);
   if (!tab || !tab.sql.trim()) return;
+  if (tab.objectSource) {
+    const saved = await saveActiveObjectSource(tab);
+    if (saved) queryStore.closeTab(tabId, { force: true });
+    return;
+  }
   const existing = tab.savedSqlId ? savedSqlStore.getFile(tab.savedSqlId) : undefined;
   if (existing) {
     const updated = await savedSqlStore.saveFile({
@@ -450,6 +455,7 @@ async function openSaveSqlDialog() {
       sql: tab.sql,
     });
     queryStore.linkSavedSql(tab.id, updated.id, updated.name);
+    queryStore.markTabClean(tab);
     toast(t("savedSql.saved"), 2000);
     return;
   }
@@ -459,10 +465,10 @@ async function openSaveSqlDialog() {
   showSaveSqlDialog.value = true;
 }
 
-async function saveActiveObjectSource(tab: NonNullable<typeof activeTab.value>) {
+async function saveActiveObjectSource(tab: QueryTab): Promise<boolean> {
   const connection = connectionStore.getConfig(tab.connectionId);
   const source = tab.objectSource;
-  if (!connection || !source) return;
+  if (!connection || !source) return false;
 
   try {
     const statements = await buildExecutableObjectSourceStatements({
@@ -479,9 +485,12 @@ async function saveActiveObjectSource(tab: NonNullable<typeof activeTab.value>)
         await api.executeScript(tab.connectionId, tab.database, sql, source.schema || tab.schema);
       }
     }
+    queryStore.markTabClean(tab);
     toast(t("objects.sourceSaved"), 2000);
+    return true;
   } catch (e: any) {
     toast(t("objects.sourceSaveFailed", { message: e?.message || String(e) }), 5000);
+    return false;
   }
 }
 
diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue
index 14d1fcdbc..45964cc87 100644
--- a/apps/desktop/src/components/sidebar/TreeItem.vue
+++ b/apps/desktop/src/components/sidebar/TreeItem.vue
@@ -1488,6 +1488,7 @@ function viewObjectSource() {
           objectType,
         });
       }
+      queryStore.markTabClean(queryStore.tabs.find((tab) => tab.id === tabId));
     })
     .catch((e: any) => {
       toast(e?.message || String(e), 5000);
diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts
index 3b3975041..f8c62d285 100644
--- a/packages/app-tests/queryStore.test.ts
+++ b/packages/app-tests/queryStore.test.ts
@@ -82,6 +82,30 @@ test("renames query tab titles", () => {
   assert.equal(tab?.customTitle, true);
 });
 
+test("marked-clean object source tabs close without unsaved confirmation", () => {
+  setActivePinia(createPinia());
+  const store = useQueryStore();
+  const tabId = store.createTab("conn-1", "db", "Source - refresh_orders");
+  store.updateSql(tabId, "CREATE PROCEDURE refresh_orders() SELECT 1;");
+  store.setObjectSource(tabId, {
+    schema: "public",
+    name: "refresh_orders",
+    objectType: "PROCEDURE",
+  });
+
+  const tab = store.tabs.find((item) => item.id === tabId);
+  assert.ok(tab);
+  assert.equal(store.isTabDirty(tab), true);
+
+  store.markTabClean(tab);
+  assert.equal(store.isTabDirty(tab), false);
+
+  store.closeTab(tabId);
+
+  assert.equal(store.showCloseConfirm, false);
+  assert.equal(store.tabs.some((item) => item.id === tabId), false);
+});
+
 test("editing query sql preserves the displayed result editability state", () => {
   setActivePinia(createPinia());
   const store = useQueryStore();