From 332daba30295854e3505d0d937e32d658b9de023 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 24 Aug 2026 18:36:54 +0000 Subject: [PATCH] fix(query): protect unsaved tabs on connection close Closes #4641 --- ...ionStore.disconnectDataTabMetadata.spec.ts | 40 +++++ .../connectionStore.oneTimeCleanup.spec.ts | 4 +- .../queryStore.structure-close.spec.ts | 22 +++ apps/desktop/src/stores/connectionStore.ts | 2 +- apps/desktop/src/stores/queryStore.ts | 56 ++++++- packages/app-tests/queryStore.test.ts | 157 ++++++++++++++++++ 6 files changed, 272 insertions(+), 9 deletions(-) diff --git a/apps/desktop/src/stores/__tests__/connectionStore.disconnectDataTabMetadata.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.disconnectDataTabMetadata.spec.ts index 1fec5edf62..aa31b73195 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.disconnectDataTabMetadata.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.disconnectDataTabMetadata.spec.ts @@ -33,6 +33,46 @@ describe("connectionStore disconnect data-tab metadata freshness", () => { setActivePinia(createPinia()); }); + it("disconnect keeps an unsaved query pending until the user resolves scoped close", async () => { + const checkConnectionHealth = vi.fn().mockResolvedValue(undefined); + const disconnectDb = vi.fn().mockResolvedValue(undefined); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth, + disconnectDb, + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + loadSchemaCache: vi.fn().mockResolvedValue(null), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSchemaCache: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useQueryStore } = await import("@/stores/queryStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + const store = useConnectionStore(); + useSettingsStore().editorSettings.disconnectTabHandlingMode = "close-tabs"; + store.connections = [pgConnection("conn-a"), pgConnection("conn-b")]; + store.connectedIds.add("conn-a"); + + const queryStore = useQueryStore(); + const queryId = queryStore.createTab("conn-a", "app", "draft query"); + queryStore.updateSql(queryId, "select 1;"); + const outsideId = queryStore.createTab("conn-b", "app", "outside"); + + await store.disconnect("conn-a"); + + expect(disconnectDb).toHaveBeenCalledWith("conn-a", undefined); + expect(store.connectedIds.has("conn-a")).toBe(false); + expect(queryStore.showCloseConfirm).toBe(true); + expect(queryStore.pendingCloseTabId).toBe(queryId); + expect(queryStore.tabs.map((tab) => tab.id)).toEqual([queryId, outsideId]); + + queryStore.forceCloseAllPendingTabs(); + expect(queryStore.tabs.map((tab) => tab.id)).toEqual([outsideId]); + }, 10_000); + it("clears data-tab metadata freshness for the disconnected connection only", async () => { const checkConnectionHealth = vi.fn().mockResolvedValue(undefined); const disconnectDb = vi.fn().mockResolvedValue(undefined); diff --git a/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts index 4c567d71c9..c8e3e1b23a 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts @@ -132,12 +132,14 @@ describe("connectionStore one_time runtime cleanup", () => { const store = useConnectionStore(); const queryStore = useQueryStore(); store.connections = [previewConnection()]; - queryStore.createTab("preview-1", "", "sales.parquet", "query"); + const queryId = queryStore.createTab("preview-1", "", "sales.parquet", "query"); + queryStore.updateSql(queryId, "select 1;"); expect(queryStore.tabs.some((tab) => tab.connectionId === "preview-1")).toBe(true); await store.removeConnection("preview-1"); expect(queryStore.tabs.some((tab) => tab.connectionId === "preview-1")).toBe(false); + expect(queryStore.showCloseConfirm).toBe(false); }); it("deleteConnectionGroups closes the tabs of removed one_time connections", async () => { diff --git a/apps/desktop/src/stores/__tests__/queryStore.structure-close.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.structure-close.spec.ts index 79fb81ca30..d405b04c79 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.structure-close.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.structure-close.spec.ts @@ -77,6 +77,28 @@ describe("queryStore structure editor close protection", () => { expect(store.closeConfirmContext).toBe("app"); }); + it("protects scoped dirty structure tabs when SQL confirmation is disabled", async () => { + const { useSettingsStore } = await import("@/stores/settingsStore"); + const { useQueryStore } = await import("@/stores/queryStore"); + const settingsStore = useSettingsStore(); + settingsStore.editorSettings.confirmUnsavedSqlClose = false; + const store = useQueryStore(); + const structureId = store.openTableStructure("mysql-1", "app", undefined, "users"); + const structureTab = store.tabs.find((item) => item.id === structureId)!; + structureTab.structureDraft = { dirty: true } as typeof structureTab.structureDraft; + const queryId = store.createTab("mysql-1", "app", "Query", "query"); + store.updateSql(queryId, "select 1"); + const outsideId = store.createTab("mysql-2", "app", "Outside", "query"); + + store.closeConnectionTabs("mysql-1"); + + expect(store.showCloseConfirm).toBe(true); + expect(store.pendingCloseTabId).toBe(structureId); + expect(store.closeConfirmDirtyTabIds).toEqual([structureId]); + store.cancelClosePendingTab(); + expect(store.tabs.map((tab) => tab.id)).toEqual([structureId, queryId, outsideId]); + }); + it("discards a structure draft before force closing", async () => { const { useQueryStore } = await import("@/stores/queryStore"); const store = useQueryStore(); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index f5679bf6f4..c40938bf3d 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -824,7 +824,7 @@ export const useConnectionStore = defineStore("connection", () => { const { useQueryStore } = await import("@/stores/queryStore"); const queryStore = useQueryStore(); for (const connectionId of connectionIds) { - queryStore.closeConnectionTabs(connectionId); + queryStore.closeConnectionTabs(connectionId, { force: true }); } } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 79d82dde07..08bd4fdbbc 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2429,7 +2429,7 @@ export const useQueryStore = defineStore("query", () => { if (pendingBatchCloseTabIds.value) { return pendingBatchCloseTabIds.value .map((id) => tabs.value.find((tab) => tab.id === id)) - .filter((tab): tab is QueryTab => !!tab && isTabDirty(tab)) + .filter((tab): tab is QueryTab => !!tab && shouldConfirmTabClose(tab)) .map((tab) => tab.id); } const pendingTab = pendingCloseTabId.value ? tabs.value.find((tab) => tab.id === pendingCloseTabId.value) : undefined; @@ -2579,9 +2579,31 @@ export const useQueryStore = defineStore("query", () => { onComplete?.(); return; } - pendingBatchCloseTabIds.value = uniqueIds; - pendingBatchCloseFinalActiveTabId.value = finalActiveTabId; - pendingBatchCloseComplete = onComplete ?? null; + + const existingIds = pendingBatchCloseTabIds.value; + if (existingIds) { + // Sidebar bulk disconnects settle independently, so later scopes must join the open dialog instead of replacing it. + const combinedIds = [...new Set([...existingIds, ...uniqueIds])]; + const preferredFinalActiveTabId = finalActiveTabId !== undefined ? finalActiveTabId : pendingBatchCloseFinalActiveTabId.value; + pendingBatchCloseTabIds.value = combinedIds; + pendingBatchCloseFinalActiveTabId.value = preferredFinalActiveTabId && combinedIds.includes(preferredFinalActiveTabId) ? activeTabAfterClosing(combinedIds, preferredFinalActiveTabId) : preferredFinalActiveTabId; + if (onComplete) { + const previousComplete = pendingBatchCloseComplete; + pendingBatchCloseComplete = previousComplete + ? () => { + try { + previousComplete(); + } finally { + onComplete(); + } + } + : onComplete; + } + } else { + pendingBatchCloseTabIds.value = uniqueIds; + pendingBatchCloseFinalActiveTabId.value = finalActiveTabId; + pendingBatchCloseComplete = onComplete ?? null; + } continuePendingBatchClose(); } @@ -2757,6 +2779,17 @@ export const useQueryStore = defineStore("query", () => { return tabs.value.find((tab) => !closingIds.has(tab.id))?.id ?? null; } + function activeTabAfterClosing(ids: string[], preferredActiveTabId = activeTabId.value) { + const closingIds = new Set(ids); + if (preferredActiveTabId && !closingIds.has(preferredActiveTabId) && tabs.value.some((tab) => tab.id === preferredActiveTabId)) { + return preferredActiveTabId; + } + const preferredIndex = preferredActiveTabId ? tabs.value.findIndex((tab) => tab.id === preferredActiveTabId) : -1; + const remainingTabs = tabs.value.filter((tab) => !closingIds.has(tab.id)); + if (preferredIndex < 0) return remainingTabs[0]?.id ?? null; + return remainingTabs[Math.min(preferredIndex, remainingTabs.length - 1)]?.id ?? null; + } + function closeOtherRegularTabs(id: string) { const tab = tabs.value.find((item) => item.id === id); if (!tab || tab.pinned) return; @@ -2905,12 +2938,21 @@ export const useQueryStore = defineStore("query", () => { } } - function closeConnectionTabs(connectionId: string) { - closeTabsWhere((tab) => tab.connectionId === connectionId); + function closeScopedTabsWhere(predicate: (tab: QueryTab) => boolean, options: { force?: boolean } = {}) { + const ids = tabs.value.filter((tab) => predicate(tab)).map((tab) => tab.id); + if (options.force) { + closeTabsWhere(predicate); + return; + } + beginBatchClose(ids, activeTabAfterClosing(ids)); + } + + function closeConnectionTabs(connectionId: string, options?: { force?: boolean }) { + closeScopedTabsWhere((tab) => tab.connectionId === connectionId, options); } function closeDatabaseTabs(connectionId: string, database: string) { - closeTabsWhere((tab) => tab.connectionId === connectionId && tab.database === database); + closeScopedTabsWhere((tab) => tab.connectionId === connectionId && tab.database === database); } function tabMatchesDroppedTableObject(tab: QueryTab, target: DroppedTableObjectTarget): boolean { diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index dcb2b6dbd8..5997e6c490 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -695,6 +695,163 @@ test("close all tabs pauses on unsaved query tabs", () => { assert.equal(store.activeTabId, null); }); +test("connection-scoped close pauses before dropping unsaved query tabs", () => { + setActivePinia(createPinia()); + const store = useQueryStore(); + const queryId = store.createTab("conn-1", "db", "draft query"); + store.updateSql(queryId, "select 1;"); + const dataId = store.createTab("conn-1", "db", "users", "data"); + const outsideId = store.createTab("conn-2", "db", "outside"); + store.activeTabId = dataId; + + store.closeConnectionTabs("conn-1"); + + assert.equal(store.showCloseConfirm, true); + assert.equal(store.pendingCloseTabId, queryId); + assert.equal(store.closeConfirmContext, "batch"); + assert.deepEqual(store.closeConfirmDirtyTabIds, [queryId]); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [queryId, dataId, outsideId], + ); + + store.cancelClosePendingTab(); + + assert.equal(store.showCloseConfirm, false); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [queryId, dataId, outsideId], + ); +}); + +test("database-scoped close completes the pending batch after all dirty tabs are saved", () => { + setActivePinia(createPinia()); + const store = useQueryStore(); + const firstId = store.createTab("conn-1", "db", "first query"); + store.updateSql(firstId, "select 1;"); + const secondId = store.createTab("conn-1", "db", "second query"); + store.updateSql(secondId, "select 2;"); + const otherDatabaseId = store.createTab("conn-1", "analytics", "outside"); + + store.closeDatabaseTabs("conn-1", "db"); + + assert.equal(store.showCloseConfirm, true); + assert.deepEqual(store.closeConfirmDirtyTabIds, [firstId, secondId]); + store.markTabClean(store.tabs.find((tab) => tab.id === firstId)); + store.markTabClean(store.tabs.find((tab) => tab.id === secondId)); + + assert.equal(store.completePendingCloseAfterSaveAll(), "tabs"); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [otherDatabaseId], + ); +}); + +test("connection-scoped close resumes after a saved tab and restores its prompt after save failure", () => { + setActivePinia(createPinia()); + const store = useQueryStore(); + const firstId = store.createTab("conn-1", "db", "first query"); + store.updateSql(firstId, "select 1;"); + const secondId = store.createTab("conn-1", "db", "second query"); + store.updateSql(secondId, "select 2;"); + const outsideId = store.createTab("conn-2", "db", "outside"); + + store.closeConnectionTabs("conn-1"); + + assert.equal(store.saveAndClosePendingTab(), firstId); + store.markTabClean(store.tabs.find((tab) => tab.id === firstId)); + store.closeTab(firstId, { force: true }); + assert.equal(store.showCloseConfirm, true); + assert.equal(store.pendingCloseTabId, secondId); + + assert.equal(store.saveAndClosePendingTab(), secondId); + assert.equal(store.resumeCloseConfirm(), true); + assert.equal(store.pendingCloseTabId, secondId); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [secondId, outsideId], + ); + + assert.equal(store.saveAndClosePendingTab(), secondId); + store.markTabClean(store.tabs.find((tab) => tab.id === secondId)); + store.closeTab(secondId, { force: true }); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [outsideId], + ); + assert.equal(store.activeTabId, outsideId); +}); + +test("connection-scoped close honors disabled unsaved SQL confirmation", () => { + const restoreStorage = installMemoryStorage(); + try { + setActivePinia(createPinia()); + useSettingsStore().updateEditorSettings({ confirmUnsavedSqlClose: false }); + const store = useQueryStore(); + const queryId = store.createTab("conn-1", "db", "draft query"); + store.updateSql(queryId, "select 1;"); + const outsideId = store.createTab("conn-2", "db", "outside"); + + store.closeConnectionTabs("conn-1"); + + assert.equal(store.showCloseConfirm, false); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [outsideId], + ); + } finally { + restoreStorage(); + } +}); + +test("concurrent connection-scoped closes coalesce every pending scope", () => { + setActivePinia(createPinia()); + const store = useQueryStore(); + const firstConnectionId = store.createTab("conn-1", "db", "first draft"); + store.updateSql(firstConnectionId, "select 1;"); + const secondConnectionId = store.createTab("conn-2", "db", "second draft"); + store.updateSql(secondConnectionId, "select 2;"); + const keepId = store.createTab("conn-3", "db", "keep"); + + store.closeConnectionTabs("conn-1"); + store.closeConnectionTabs("conn-2"); + + assert.deepEqual(store.closeConfirmDirtyTabIds, [firstConnectionId, secondConnectionId]); + store.forceCloseAllPendingTabs(); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [keepId], + ); + assert.equal(store.activeTabId, keepId); +}); + +test("scoped close coalesces with pending completion state", () => { + setActivePinia(createPinia()); + const store = useQueryStore(); + const firstConnectionId = store.createTab("conn-1", "db", "first draft"); + store.updateSql(firstConnectionId, "select 1;"); + const keepId = store.createTab("conn-3", "db", "keep"); + const secondConnectionId = store.createTab("conn-2", "db", "second draft"); + store.updateSql(secondConnectionId, "select 2;"); + let completions = 0; + + store.closeRightTabs(keepId, () => { + completions += 1; + }); + store.closeConnectionTabs("conn-1"); + + assert.equal(store.showCloseConfirm, true); + assert.deepEqual(store.closeConfirmDirtyTabIds, [secondConnectionId, firstConnectionId]); + store.forceCloseAllPendingTabs(); + + assert.equal(completions, 1); + assert.deepEqual( + store.tabs.map((tab) => tab.id), + [keepId], + ); + assert.equal(store.activeTabId, keepId); +}); + test("disabled unsaved SQL close confirmation closes dirty tabs directly", () => { const restoreStorage = installMemoryStorage(); try {