From 295b67dd364950d8413ef7963eaa9dea78045367 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Thu, 16 Jul 2026 00:19:24 +0800 Subject: [PATCH] feat(core): optimize query resource lifecycle --- .../src/components/layout/ContentArea.vue | 11 +- .../src/components/sidebar/ConnectionTree.vue | 18 +- .../src/components/sidebar/TreeItem.vue | 54 ++- apps/desktop/src/i18n/locales/en.ts | 2 + .../fixtures/queryResultLifecycle.ts | 57 +++ .../sidebarDataOpenCoordinator.spec.ts | 140 ++++-- .../__tests__/tabs/queryResultSize.spec.ts | 43 ++ .../tabs/resourceLifecycleBaseline.spec.ts | 57 +++ .../lib/__tests__/tabs/tabResultCache.spec.ts | 93 +++- .../resourceLifecycleDiagnostics.ts | 25 ++ .../lib/sidebar/sidebarDataOpenCoordinator.ts | 158 +++++-- apps/desktop/src/lib/tabs/queryResultSize.ts | 51 +++ apps/desktop/src/lib/tabs/tabResultCache.ts | 420 +++++++++++++++--- apps/desktop/src/stores/connectionStore.ts | 2 + apps/desktop/src/stores/queryStore.ts | 210 +++++---- apps/desktop/src/types/database.ts | 2 + crates/dbx-core/src/connection.rs | 110 +++-- crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/query_cancel.rs | 219 ++++++++- crates/dbx-core/src/storage.rs | 279 +++++++++++- crates/dbx-core/src/task_supervisor.rs | 111 +++++ crates/dbx-web/src/main.rs | 14 +- crates/dbx-web/src/routes/connection.rs | 1 + .../dbx-web/src/routes/tab_runtime_cache.rs | 50 ++- packages/app-tests/queryStore.test.ts | 8 +- .../app-tests/sidebarContextMenuHost.test.ts | 2 +- src-tauri/src/commands/connection.rs | 1 + src-tauri/src/commands/tab_runtime_cache.rs | 46 +- src-tauri/src/lib.rs | 7 +- 29 files changed, 1857 insertions(+), 335 deletions(-) create mode 100644 apps/desktop/src/lib/__tests__/fixtures/queryResultLifecycle.ts create mode 100644 apps/desktop/src/lib/__tests__/tabs/queryResultSize.spec.ts create mode 100644 apps/desktop/src/lib/__tests__/tabs/resourceLifecycleBaseline.spec.ts create mode 100644 apps/desktop/src/lib/diagnostics/resourceLifecycleDiagnostics.ts create mode 100644 apps/desktop/src/lib/tabs/queryResultSize.ts create mode 100644 crates/dbx-core/src/task_supervisor.rs diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 5210273d36..275e465ea8 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -311,7 +311,10 @@ const activeQueryError = computed(() => { if (!result?.columns.includes("Error")) return ""; return String(result.rows[0]?.[0] ?? ""); }); -const hasQueryOutput = computed(() => !!props.activeTab.result || !!props.activeTab.explainPlan || !!props.activeTab.explainError || !!props.activeTab.explainTableResult || !!props.activeTab.explainTableError || props.activeTab.isExecuting === true || props.activeTab.isExplaining === true); +const hasQueryOutput = computed( + () => + !!props.activeTab.result || props.activeTab.resultEvicted === true || !!props.activeTab.explainPlan || !!props.activeTab.explainError || !!props.activeTab.explainTableResult || !!props.activeTab.explainTableError || props.activeTab.isExecuting === true || props.activeTab.isExplaining === true, +); const visibleResultItems = computed(() => tabularResultItems(props.activeTab.results ?? (props.activeTab.result ? [props.activeTab.result] : undefined))); const tabularResults = computed(() => tabularResultItems(props.activeTab.results)); const allResultExportSheets = computed(() => @@ -1196,6 +1199,12 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand :cancelling="activeTab.isCancelling" @cancel="emit('cancel')" /> +
+
{{ t("editor.cachedResultUnavailable") }}
+ +
{{ t("editor.pressToExecute", { mod: shortcutModifier }) }}
{{ t("editor.pressToSaveSql", { mod: shortcutModifier }) }}
diff --git a/apps/desktop/src/components/sidebar/ConnectionTree.vue b/apps/desktop/src/components/sidebar/ConnectionTree.vue index a85870ffd6..df37609ba5 100644 --- a/apps/desktop/src/components/sidebar/ConnectionTree.vue +++ b/apps/desktop/src/components/sidebar/ConnectionTree.vue @@ -961,12 +961,20 @@ function openSidebarProcedure(node: TreeNode) { sidebarProcedureOpen.value = true; } -function openSidebarData(node: TreeNode, requireSelection: boolean, runner: (node: TreeNode, request: SidebarDataOpenRequest) => Promise) { +function openSidebarData(node: TreeNode, requireSelection: boolean, openMode: "default" | "new-tab", runner: (node: TreeNode, request: SidebarDataOpenRequest) => Promise) { const target = createSidebarActionTarget(node); - runSidebarDataOpenImmediately((request) => { - if (requireSelection && store.selectedTreeNodeId !== target.id) return; - return runner(target, request); - }); + runSidebarDataOpenImmediately( + { + connectionKey: target.connectionId || target.id, + // Explicit new-tab opens are intentional independent work; ordinary + // navigation keeps latest-request-wins behavior. + supersede: openMode !== "new-tab", + }, + (request) => { + if (requireSelection && store.selectedTreeNodeId !== target.id) return; + return runner(target, request); + }, + ); } function openSidebarVisibleDatabases(node: TreeNode) { diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 445657ad52..b0c7af0057 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -80,6 +80,7 @@ import { buildTableDeleteTemplate, buildTableInsertTemplate, buildTableSelectTem import { connectionFilePath, defaultSqliteBackupFileName, isMemorySqlitePath, sqliteBackupSourcePath } from "@/lib/connection/connectionFile"; import { driverStoreFocusForInstallError } from "@/lib/connection/agentDriverInstallHint"; import { revealPathInFileManager } from "@/lib/backend/tauri"; +import { appendDebugLog, isDebugLoggingEnabled } from "@/lib/backend/debugLog"; import { clearActiveTableReferencePayload, createTableReferencePayload, createTableReferenceDropEvent, setActiveTableReferencePayload, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop"; import { usesSyntheticRowIdKey } from "@/lib/table/tableEditing"; import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; @@ -354,7 +355,7 @@ const emit = defineEmits<{ "open-ddl": [node: TreeNode]; "open-object-source": [node: TreeNode, initialEditing: boolean]; "open-procedure": [node: TreeNode]; - "open-data": [node: TreeNode, requireSelection: boolean, runner: (node: TreeNode, request: SidebarDataOpenRequest) => Promise]; + "open-data": [node: TreeNode, requireSelection: boolean, openMode: DataTabOpenMode, runner: (node: TreeNode, request: SidebarDataOpenRequest) => Promise]; "open-visible-databases": [node: TreeNode]; "open-visible-schemas": [node: TreeNode]; "open-danger-dialog": [request: SidebarDangerDialogRequest]; @@ -1353,15 +1354,15 @@ async function openDamengJobAdmin() { } function scheduleOpenData(node: TreeNode) { - emit("open-data", node, true, openData); + emit("open-data", node, true, "default", openData); } function openDataImmediately(node: TreeNode = props.node) { - emit("open-data", node, false, openData); + emit("open-data", node, false, "default", openData); } function openDataInNewTabImmediately(node: TreeNode = props.node) { - emit("open-data", node, false, (target, request) => openData(target, request, "new-tab")); + emit("open-data", node, false, "new-tab", (target, request) => openData(target, request, "new-tab")); } async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMode: DataTabOpenMode = "default") { @@ -1371,9 +1372,12 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo const startedAt = performance.now(); let lastPhaseAt = startedAt; const elapsed = () => `${Math.round(performance.now() - startedAt)}ms`; + const openDataLog = (level: "info" | "warn" | "error" | "debug", event: string, details: Record) => { + appendDebugLog(level, `[DBX][openData:${event}]`, details); + }; const logPhase = (phase: string, extra: Record = {}) => { const now = performance.now(); - console.info("[DBX][openData:phase]", { + openDataLog("info", "phase", { traceId, phase, deltaMs: Math.round(now - lastPhaseAt), @@ -1382,14 +1386,11 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo }); lastPhaseAt = now; }; - console.info("[DBX][openData:start]", { + openDataLog("info", "start", { traceId, type: node.type, - connectionId: node.connectionId, - database: node.database, - schema: node.schema, - table: node.label, dbType: config?.db_type, + openMode, }); const tableSchema = connectionObjectTreeNodeSchema(config, node.database, node.schema); const tableType = node.type === "view" ? "VIEW" : node.type === "materialized_view" ? "MATERIALIZED_VIEW" : (node.tableType ?? "TABLE"); @@ -1445,7 +1446,7 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo } return queryStore.createTab(node.connectionId, node.database, node.label, "data", tableSchema); })(); - console.info("[DBX][openData:tab-created]", { traceId, tabId, elapsed: elapsed() }); + openDataLog("info", "tab-created", { traceId, tabId, elapsed: elapsed() }); logPhase("tab-created", { tabId }); // Cancel any previous execution on this tab before starting a new one @@ -1519,24 +1520,21 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo }; try { - console.info("[DBX][openData:ensure-connected:start]", { traceId, elapsed: elapsed() }); + openDataLog("info", "ensure-connected:start", { traceId, elapsed: elapsed() }); await connectionStore.ensureConnected(node.connectionId); if (!isActive()) { logPhase("superseded-after-ensure-connected", { tabId }); return; } - console.info("[DBX][openData:ensure-connected:done]", { traceId, elapsed: elapsed() }); + openDataLog("info", "ensure-connected:done", { traceId, elapsed: elapsed() }); logPhase("ensure-connected", { tabId }); if (!config) throw new Error("Connection config not found"); const limit = tableOpenPageLimit(); const refreshTableMetaInBackground = async () => { const metadataStartedAt = performance.now(); - console.info("[DBX][openData:metadata:start]", { + openDataLog("info", "metadata:start", { traceId, - database: node.database, - schema: querySchema, - table: node.label, elapsed: elapsed(), }); try { @@ -1549,10 +1547,10 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo databaseType: metadataDatabaseType, driverProfile: config.driver_profile || config.db_type, catalog: node.catalog, - traceLogger: (event) => console.debug("[DBX][openData:metadata:trace]", { sourceTraceId: traceId, ...event }), + traceLogger: isDebugLoggingEnabled() ? (event) => openDataLog("debug", "metadata:trace", { sourceTraceId: traceId, ...event }) : undefined, }); if (!isCurrentDataTab()) { - console.info("[DBX][openData:metadata:stale]", { + openDataLog("info", "metadata:stale", { traceId, tabId, columnCount: loadedMetadata.metadata.columns.length, @@ -1562,7 +1560,7 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo } const nextTableMeta = tableMetadataToDataTabMeta(loadedMetadata.metadata, tableSchema); queryStore.setTableMeta(tabId, nextTableMeta); - console.info("[DBX][openData:metadata:done]", { + openDataLog("info", "metadata:done", { traceId, tabId, columnCount: nextTableMeta.columns.length, @@ -1573,12 +1571,12 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo metadataMs: Math.round(performance.now() - metadataStartedAt), }); } catch (error) { - console.warn("[DBX][openData:metadata:error]", { traceId, tabId, elapsed: elapsed(), error }); + openDataLog("warn", "metadata:error", { traceId, tabId, elapsed: elapsed(), error }); } }; const shouldRefreshTableMeta = !cachedTableMeta; if (cachedTableMeta) { - console.info("[DBX][openData:metadata:cache-hit]", { + openDataLog("info", "metadata:cache-hit", { traceId, tabId, columnCount: cachedTableMeta.columns.length, @@ -1613,24 +1611,24 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo limit, includeRowId, }); - console.info("[DBX][openData:sql-built]", { + openDataLog("info", "sql-built", { traceId, - primaryKeys, + primaryKeyCount: primaryKeys.length, includeRowId, - sql, + sqlLength: sql.length, elapsed: elapsed(), }); logPhase("sql-built", { tabId, columnCount: columns.length, primaryKeyCount: primaryKeys.length }); queryStore.updateSql(tabId, sql); logPhase("sql-updated", { tabId }); - console.info("[DBX][openData:execute:start]", { traceId, tabId, elapsed: elapsed() }); + openDataLog("info", "execute:start", { traceId, tabId, elapsed: elapsed() }); await queryStore.executeTabSql(tabId, sql, { sourceTraceId: traceId, skipEnsureConnected: true, pagination: { limit, offset: 0 }, }); - console.info("[DBX][openData:execute:done]", { traceId, tabId, elapsed: elapsed() }); + openDataLog("info", "execute:done", { traceId, tabId, elapsed: elapsed() }); logPhase("execute-tab-sql", { tabId }); if (shouldRefreshTableMeta && isCurrentDataTab()) { void refreshTableMetaInBackground(); @@ -1641,7 +1639,7 @@ async function openData(node: TreeNode, request?: SidebarDataOpenRequest, openMo logPhase("superseded-after-error", { tabId }); return; } - console.error("[DBX][openData:error]", { traceId, elapsed: elapsed(), error: e }); + openDataLog("error", "error", { traceId, elapsed: elapsed(), error: e }); logPhase("error", { tabId }); queryStore.setErrorResult(tabId, e); } diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index b383a92fa0..b651c27d83 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -1038,6 +1038,8 @@ export default { noColumnsFound: "No columns found", queryError: "Query Error", dataUnavailable: "Table data needs to be reloaded.", + cachedResultUnavailable: "The cached result is missing or incompatible.", + reexecuteQuery: "Run query again", dataUnavailableHintPrefix: "Press ", dataUnavailableHintSuffix: " or click Refresh below to reload.", refresh: "Refresh", diff --git a/apps/desktop/src/lib/__tests__/fixtures/queryResultLifecycle.ts b/apps/desktop/src/lib/__tests__/fixtures/queryResultLifecycle.ts new file mode 100644 index 0000000000..e46903cb72 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/fixtures/queryResultLifecycle.ts @@ -0,0 +1,57 @@ +import type { QueryTab } from "@/types/database"; +import type { TabResultSnapshot } from "@/lib/tabs/tabResultCache"; + +export function queryResultLifecycleSnapshot(): TabResultSnapshot { + const first = { + columns: ["id", "name"], + rows: [[1, "Ada"]], + affected_rows: 0, + execution_time_ms: 4, + sourceStatement: "SELECT id, name FROM users", + }; + const second = { + columns: ["count"], + rows: [[1]], + affected_rows: 0, + execution_time_ms: 2, + sourceStatement: "SELECT COUNT(*) FROM users", + }; + const resultRuns: NonNullable = [ + { + id: "run-1", + sequence: 1, + result: first, + results: [first, second], + activeResultIndex: 0, + resultLocalSortOriginalRows: [[1, "Ada"]], + resultBaseSql: "SELECT id, name FROM users", + resultPageSql: "SELECT id, name FROM users LIMIT 100 OFFSET 200", + resultPageLimit: 100, + resultPageOffset: 200, + tableMeta: { + schema: "public", + tableName: "users", + tableType: "TABLE", + columns: [{ name: "id", data_type: "integer", nullable: false }], + primaryKeys: ["id"], + }, + }, + ]; + return { + result: first, + results: [first, second], + activeResultIndex: 0, + resultLocalSortOriginalRows: [[1, "Ada"]], + resultRuns, + activeResultRunId: "run-1", + queryAnalysis: { tableName: "users", schema: "public", allowInsert: true, allowInsertDelete: true }, + querySourceColumns: ["id", "name"], + tableMeta: resultRuns[0].tableMeta, + resultPageSql: "SELECT id, name FROM users LIMIT 100 OFFSET 200", + resultPageLimit: 100, + resultPageOffset: 200, + resultCountSql: "SELECT COUNT(*) FROM users", + resultTotalRowCount: 301, + cachedAt: 1, + }; +} diff --git a/apps/desktop/src/lib/__tests__/sidebar/sidebarDataOpenCoordinator.spec.ts b/apps/desktop/src/lib/__tests__/sidebar/sidebarDataOpenCoordinator.spec.ts index b3dae50285..8b678ab380 100644 --- a/apps/desktop/src/lib/__tests__/sidebar/sidebarDataOpenCoordinator.spec.ts +++ b/apps/desktop/src/lib/__tests__/sidebar/sidebarDataOpenCoordinator.spec.ts @@ -1,95 +1,141 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { cancelPendingSidebarDataOpen, runSidebarDataOpenImmediately } from "@/lib/sidebar/sidebarDataOpenCoordinator"; +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + describe("sidebarDataOpenCoordinator", () => { afterEach(() => { cancelPendingSidebarDataOpen(); }); - it("cancels the active request when a newer open starts", async () => { + it("cancels the active request when a newer ordinary open starts", async () => { const cancel = vi.fn(); + const running = deferred(); let firstIsCurrent = true; - runSidebarDataOpenImmediately(async (request) => { + runSidebarDataOpenImmediately({ connectionKey: "conn" }, async (request) => { request.registerCancel(cancel); - await Promise.resolve(); + await running.promise; firstIsCurrent = request.isCurrent(); }); - runSidebarDataOpenImmediately(() => undefined); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, () => undefined); await Promise.resolve(); expect(cancel).toHaveBeenCalledOnce(); + expect(firstIsCurrent).toBe(true); + + running.resolve(); + await Promise.resolve(); + await Promise.resolve(); expect(firstIsCurrent).toBe(false); }); + it("replaces a superseded queued request before it starts", async () => { + const running = deferred(); + const second = vi.fn(); + const third = vi.fn(); + + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, () => running.promise); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, second); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, third); + running.resolve(); + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + + expect(second).not.toHaveBeenCalled(); + expect(third).toHaveBeenCalledOnce(); + }); + + it("bounds active work per connection", async () => { + const first = deferred(); + const starts: string[] = []; + + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, async () => { + starts.push("first"); + await first.promise; + }); + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, () => { + starts.push("second"); + }); + + expect(starts).toEqual(["first"]); + first.resolve(); + await vi.waitFor(() => expect(starts).toEqual(["first", "second"])); + }); + + it("allows different connections to make progress independently", () => { + const blocked = deferred(); + const otherConnection = vi.fn(); + + runSidebarDataOpenImmediately({ connectionKey: "conn-a", supersede: false }, () => blocked.promise); + runSidebarDataOpenImmediately({ connectionKey: "conn-b", supersede: false }, otherConnection); + + expect(otherConnection).toHaveBeenCalledOnce(); + blocked.resolve(); + }); + + it("keeps explicit new-tab opens independent while still queueing them", async () => { + const first = deferred(); + const firstCancel = vi.fn(); + const second = vi.fn(); + + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, async (request) => { + request.registerCancel(firstCancel); + await first.promise; + }); + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, second); + + expect(firstCancel).not.toHaveBeenCalled(); + expect(second).not.toHaveBeenCalled(); + first.resolve(); + await vi.waitFor(() => expect(second).toHaveBeenCalledOnce()); + }); + it("runs registered cancellation immediately for an already stale request", async () => { const cancel = vi.fn(); let registerLateCancel: (() => void) | undefined; + const running = deferred(); - runSidebarDataOpenImmediately((request) => { + runSidebarDataOpenImmediately({ connectionKey: "conn" }, async (request) => { registerLateCancel = () => request.registerCancel(cancel); + await running.promise; }); - runSidebarDataOpenImmediately(() => undefined); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, () => undefined); registerLateCancel?.(); await Promise.resolve(); expect(cancel).toHaveBeenCalledOnce(); + running.resolve(); }); it("recovers after a runner throws synchronously", async () => { const nextRunner = vi.fn(); - runSidebarDataOpenImmediately(() => { + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, () => { throw new Error("boom"); }); await Promise.resolve(); - runSidebarDataOpenImmediately(nextRunner); await Promise.resolve(); + runSidebarDataOpenImmediately({ connectionKey: "conn", supersede: false }, nextRunner); expect(nextRunner).toHaveBeenCalledOnce(); }); - it("keeps only the newest rapid navigation request current", async () => { - const states: boolean[] = []; - const captures: Array<() => void> = []; + it("aborts a superseded request signal", () => { + let signal: AbortSignal | undefined; - for (let index = 0; index < 3; index += 1) { - runSidebarDataOpenImmediately((request) => { - captures.push(() => states.push(request.isCurrent())); - }); - } - captures.forEach((capture) => capture()); - - expect(states).toEqual([false, false, true]); - }); - - it("does not cancel work that was not registered as sidebar-owned", async () => { - const sidebarCancel = vi.fn(); - const unrelatedEditorCancel = vi.fn(); - - runSidebarDataOpenImmediately((request) => request.registerCancel(sidebarCancel)); - runSidebarDataOpenImmediately(() => undefined); - await Promise.resolve(); - - expect(sidebarCancel).toHaveBeenCalledOnce(); - expect(unrelatedEditorCancel).not.toHaveBeenCalled(); - }); - - it("prevents a superseded connection completion from becoming current", async () => { - let finishConnection: (() => void) | undefined; - let completionWasCurrent = true; - - runSidebarDataOpenImmediately(async (request) => { - await new Promise((resolve) => { - finishConnection = resolve; - }); - completionWasCurrent = request.isCurrent(); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, (request) => { + signal = request.signal; }); - runSidebarDataOpenImmediately(() => undefined); - finishConnection?.(); - await Promise.resolve(); - await Promise.resolve(); + runSidebarDataOpenImmediately({ connectionKey: "conn" }, () => undefined); - expect(completionWasCurrent).toBe(false); + expect(signal?.aborted).toBe(true); }); }); diff --git a/apps/desktop/src/lib/__tests__/tabs/queryResultSize.spec.ts b/apps/desktop/src/lib/__tests__/tabs/queryResultSize.spec.ts new file mode 100644 index 0000000000..ad0d9c302e --- /dev/null +++ b/apps/desktop/src/lib/__tests__/tabs/queryResultSize.spec.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from "vitest"; +import { estimateQueryResultBytes, estimateQueryResultsBytes, selectInactiveResultEvictions } from "@/lib/tabs/queryResultSize"; + +function result(rows: unknown[][]) { + return { columns: ["value"], rows, affected_rows: 0, execution_time_ms: 1 }; +} + +describe("query result size estimation", () => { + it("grows with string and binary payloads", () => { + const small = estimateQueryResultBytes(result([["a"]])); + const large = estimateQueryResultBytes(result([["a".repeat(10_000)], [new Uint8Array(20_000)]])); + expect(large).toBeGreaterThan(small + 30_000); + }); + + it("counts multiple results and tolerates cyclic document values", () => { + const document: Record = { name: "node" }; + document.self = document; + const first = result([[document]]); + const second = result(Array.from({ length: 100 }, (_, index) => [index])); + expect(estimateQueryResultsBytes(undefined, [first, second])).toBeGreaterThan(estimateQueryResultBytes(first)); + }); +}); + +describe("inactive result byte budgets", () => { + it("evicts oldest small results to the count limit", () => { + const entries = Array.from({ length: 8 }, (_, index) => ({ id: `${index}`, estimatedBytes: 10, accessedAt: index })); + expect(selectInactiveResultEvictions(entries, 1_000, 5)).toEqual(["0", "1", "2"]); + }); + + it("evicts least-recently-used large results to the byte limit", () => { + const entries = [ + { id: "old", estimatedBytes: 80, accessedAt: 1 }, + { id: "middle", estimatedBytes: 70, accessedAt: 2 }, + { id: "new", estimatedBytes: 60, accessedAt: 3 }, + ]; + expect(selectInactiveResultEvictions(entries, 100, 10)).toEqual(["old", "middle"]); + }); + + it("does not consider the active result even when it exceeds the budget", () => { + const inactive = [{ id: "inactive", estimatedBytes: 10, accessedAt: 1 }]; + expect(selectInactiveResultEvictions(inactive, 100, 5)).toEqual([]); + }); +}); diff --git a/apps/desktop/src/lib/__tests__/tabs/resourceLifecycleBaseline.spec.ts b/apps/desktop/src/lib/__tests__/tabs/resourceLifecycleBaseline.spec.ts new file mode 100644 index 0000000000..178281ea47 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/tabs/resourceLifecycleBaseline.spec.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from "vitest"; +import { encodeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; +import { estimateQueryResultBytes } from "@/lib/tabs/queryResultSize"; +import { cancelPendingSidebarDataOpen, runSidebarDataOpenImmediately, sidebarDataOpenDiagnostics } from "@/lib/sidebar/sidebarDataOpenCoordinator"; + +describe("resource lifecycle performance harness", () => { + it("records deterministic serialization and memory-estimation measurements", () => { + const result = { + columns: ["id", "label", "payload"], + rows: Array.from({ length: 25_000 }, (_, index) => [index, `row-${index}`, "x".repeat(64)]), + affected_rows: 0, + execution_time_ms: 1, + }; + + const estimateStartedAt = performance.now(); + const estimatedBytes = estimateQueryResultBytes(result); + const estimateDurationMs = performance.now() - estimateStartedAt; + const serializeStartedAt = performance.now(); + const encodedBytes = encodeTabResultSnapshot({ result, cachedAt: 1 }).byteLength; + const serializationDurationMs = performance.now() - serializeStartedAt; + + console.info("[resource-lifecycle-baseline]", JSON.stringify({ rows: result.rows.length, estimatedBytes, encodedBytes, estimateDurationMs, serializationDurationMs })); + expect(estimatedBytes).toBeGreaterThan(encodedBytes); + expect(estimateDurationMs).toBeLessThan(250); + expect(serializationDurationMs).toBeLessThan(1_500); + }); +}); + +describe("rapid table-open performance harness", () => { + it("coalesces a 100-request burst to one active and one latest queued request", async () => { + let release!: () => void; + const blocked = new Promise((resolve) => { + release = resolve; + }); + const committed: number[] = []; + const startedAt = performance.now(); + runSidebarDataOpenImmediately({ connectionKey: "connection" }, async (request) => { + await blocked; + if (request.isCurrent()) committed.push(0); + }); + for (let index = 1; index < 100; index += 1) { + runSidebarDataOpenImmediately({ connectionKey: "connection" }, (request) => { + if (request.isCurrent()) committed.push(index); + }); + } + const enqueueDurationMs = performance.now() - startedAt; + const burstDiagnostics = sidebarDataOpenDiagnostics(); + release(); + await new Promise((resolve) => setTimeout(resolve, 0)); + const totalDurationMs = performance.now() - startedAt; + console.info("[rapid-table-open-baseline]", JSON.stringify({ enqueueDurationMs, totalDurationMs, ...burstDiagnostics })); + + expect(burstDiagnostics).toEqual({ activeTasks: 1, queuedTasks: 1, trackedRequests: 2 }); + expect(committed).toEqual([99]); + cancelPendingSidebarDataOpen(); + }); +}); diff --git a/apps/desktop/src/lib/__tests__/tabs/tabResultCache.spec.ts b/apps/desktop/src/lib/__tests__/tabs/tabResultCache.spec.ts index badad44d45..0dbe1475cc 100644 --- a/apps/desktop/src/lib/__tests__/tabs/tabResultCache.spec.ts +++ b/apps/desktop/src/lib/__tests__/tabs/tabResultCache.spec.ts @@ -1,7 +1,78 @@ -import { describe, expect, it } from "vitest"; -import { decodeTabResultSnapshot, encodeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; +import { describe, expect, it, vi } from "vitest"; +import { decodeTabResultSnapshot, encodeTabResultSnapshot, readResultCacheBackends, promoteFallbackResultCacheRead, resultCacheBackendOrder, resultCacheRuntimeConfig, selectResultCachePruneKeys, writeResultCacheBackends, type ResultCacheBackend } from "@/lib/tabs/tabResultCache"; +import { queryResultLifecycleSnapshot } from "@/lib/__tests__/fixtures/queryResultLifecycle"; + +function backend(name: ResultCacheBackend["name"], overrides: Partial = {}): ResultCacheBackend { + return { + name, + available: () => true, + read: async () => undefined, + write: async () => true, + delete: async () => undefined, + listMetadata: async () => [], + prune: async () => ({ deletedEntries: 0, deletedBytes: 0, orphanDeletions: 0, remainingEntries: 0, remainingBytes: 0 }), + deleteOwner: async () => undefined, + ...overrides, + }; +} describe("tab result cache statement execution metadata", () => { + it("selects one authoritative backend before fallback", () => { + expect(resultCacheBackendOrder(true)).toEqual(["runtime", "indexed-db"]); + expect(resultCacheBackendOrder(false)).toEqual(["indexed-db", "runtime"]); + }); + + it("uses explicit browser runtime configuration", () => { + expect(resultCacheRuntimeConfig(false, { VITE_DBX_RESULT_CACHE_BACKEND: "http", VITE_DBX_RESULT_CACHE_FALLBACK: "true" } as ImportMetaEnv)).toEqual({ + primary: "runtime", + fallbackEnabled: true, + }); + }); + + it("writes only to the first successful backend", async () => { + const primaryWrite = vi.fn(async () => true); + const fallbackWrite = vi.fn(async () => true); + await writeResultCacheBackends([backend("runtime", { write: primaryWrite }), backend("indexed-db", { write: fallbackWrite })], "key", new Uint8Array([1]), { rowCount: 1, columnCount: 1 }); + expect(primaryWrite).toHaveBeenCalledOnce(); + expect(fallbackWrite).not.toHaveBeenCalled(); + }); + + it("reads legacy fallback data for promotion", async () => { + const bytes = new Uint8Array([1, 2, 3]); + const primary = backend("runtime"); + const fallback = backend("indexed-db", { read: async () => bytes }); + expect(await readResultCacheBackends([primary, fallback], "key")).toEqual({ bytes, backend: fallback }); + }); + + it("promotes legacy IndexedDB or native bytes to the authoritative backend", async () => { + const bytes = new Uint8Array([1, 2, 3]); + const primaryWrite = vi.fn(async () => true); + const primary = backend("runtime", { write: primaryWrite }); + const legacy = backend("indexed-db"); + + expect(await promoteFallbackResultCacheRead([primary, legacy], "key", { bytes, backend: legacy }, { rowCount: 2, columnCount: 1 }, "connection-1")).toBe(true); + expect(primaryWrite).toHaveBeenCalledWith("key", bytes, { rowCount: 2, columnCount: 1 }, "connection-1"); + }); + + it("skips an obsolete write before touching a backend", async () => { + const write = vi.fn(async () => true); + expect(await writeResultCacheBackends([backend("runtime", { write })], "key", new Uint8Array([1]), { rowCount: 1, columnCount: 1 }, () => false)).toBe(false); + expect(write).not.toHaveBeenCalled(); + }); + + it("prunes crash leftovers and then unreferenced LRU entries", () => { + const selection = selectResultCachePruneKeys( + [ + { key: "live", rowCount: 1, columnCount: 1, byteSize: 8, createdAt: 1, lastAccessedAt: 1 }, + { key: "crash", rowCount: 1, columnCount: 1, byteSize: 7, createdAt: 1, lastAccessedAt: 1 }, + { key: "recent", rowCount: 1, columnCount: 1, byteSize: 6, createdAt: 9_999, lastAccessedAt: 9_999 }, + ], + { liveKeys: ["live"], maxBytes: 10, orphanGraceMs: 100, maxAgeMs: 1_000 }, + 10_000, + ); + expect(selection.keys).toEqual(["crash", "recent"]); + expect(selection.result).toMatchObject({ orphanDeletions: 1, remainingBytes: 8 }); + }); it("preserves statement identity and the editor fingerprint", () => { const encoded = encodeTabResultSnapshot({ results: [ @@ -35,4 +106,22 @@ describe("tab result cache statement execution metadata", () => { { statementIndex: 1, error: true }, ]); }); + + it("restores multi-result, pagination, local-sort, and editable metadata fixtures", () => { + const restored = decodeTabResultSnapshot(encodeTabResultSnapshot(queryResultLifecycleSnapshot())); + + expect(restored?.results).toHaveLength(2); + expect(restored?.resultRuns?.[0]?.resultLocalSortOriginalRows).toEqual([[1, "Ada"]]); + expect(restored?.resultPageOffset).toBe(200); + expect(restored?.resultTotalRowCount).toBe(301); + expect(restored?.tableMeta?.primaryKeys).toEqual(["id"]); + expect(restored?.querySourceColumns).toEqual(["id", "name"]); + }); + + it("treats corrupt and unsupported snapshots as missing", () => { + expect(decodeTabResultSnapshot(new Uint8Array([0xff, 0x00]))).toBeUndefined(); + const encoded = encodeTabResultSnapshot(queryResultLifecycleSnapshot()); + encoded[0] = 0xff; + expect(decodeTabResultSnapshot(encoded)).toBeUndefined(); + }); }); diff --git a/apps/desktop/src/lib/diagnostics/resourceLifecycleDiagnostics.ts b/apps/desktop/src/lib/diagnostics/resourceLifecycleDiagnostics.ts new file mode 100644 index 0000000000..60d7052e96 --- /dev/null +++ b/apps/desktop/src/lib/diagnostics/resourceLifecycleDiagnostics.ts @@ -0,0 +1,25 @@ +import type { QueryTab } from "@/types/database"; +import { getResultCacheDiagnostics } from "@/lib/tabs/tabResultCache"; + +const cancellation = { + count: 0, + totalLatencyMs: 0, + maxLatencyMs: 0, +}; + +export function recordQueryCancellationLatency(durationMs: number) { + cancellation.count += 1; + cancellation.totalLatencyMs += durationMs; + cancellation.maxLatencyMs = Math.max(cancellation.maxLatencyMs, durationMs); +} + +export function resourceLifecycleDiagnostics(tabs: readonly QueryTab[]) { + const cache = getResultCacheDiagnostics(); + return { + activeTasks: tabs.filter((tab) => tab.isExecuting || tab.isExplaining || tab.isCancelling).length, + cancellationCount: cancellation.count, + averageCancellationLatencyMs: cancellation.count ? cancellation.totalLatencyMs / cancellation.count : 0, + maxCancellationLatencyMs: cancellation.maxLatencyMs, + ...cache, + }; +} diff --git a/apps/desktop/src/lib/sidebar/sidebarDataOpenCoordinator.ts b/apps/desktop/src/lib/sidebar/sidebarDataOpenCoordinator.ts index 53365ffe9e..ded24db65d 100644 --- a/apps/desktop/src/lib/sidebar/sidebarDataOpenCoordinator.ts +++ b/apps/desktop/src/lib/sidebar/sidebarDataOpenCoordinator.ts @@ -1,60 +1,166 @@ export interface SidebarDataOpenRequest { isCurrent: () => boolean; + signal: AbortSignal; registerCancel: (cancel: () => void | Promise) => void; } +export interface SidebarDataOpenOptions { + connectionKey: string; + supersede?: boolean; +} + type OpenDataRunner = (request: SidebarDataOpenRequest) => void | Promise; -let generation = 0; -let activeGeneration = 0; -let activeCancel: (() => void | Promise) | null = null; +interface QueuedOpen { + id: number; + connectionKey: string; + superseding: boolean; + controller: AbortController; + runner: OpenDataRunner; + cancel: (() => void | Promise) | null; + started: boolean; + finished: boolean; +} + +const DEFAULT_CONNECTION_KEY = "__sidebar__"; +const MAX_ACTIVE_PER_CONNECTION = 1; + +let nextRequestId = 0; +let latestSupersedingRequestId = 0; +const queuedByConnection = new Map(); +const activeCountByConnection = new Map(); +const requests = new Map(); function runCancellation(cancel: (() => void | Promise) | null) { if (!cancel) return; void Promise.resolve(cancel()).catch(() => undefined); } -function supersedeCurrentRequest(): number { - generation += 1; - runCancellation(activeCancel); - activeCancel = null; - activeGeneration = 0; - return generation; +function isCurrent(request: QueuedOpen): boolean { + if (request.finished || request.controller.signal.aborted) return false; + return !request.superseding || request.id === latestSupersedingRequestId; } -function executeRequest(requestGeneration: number, runner: OpenDataRunner) { - if (requestGeneration !== generation) return; - const request: SidebarDataOpenRequest = { - isCurrent: () => requestGeneration === generation, +function removeQueuedRequest(request: QueuedOpen) { + const queue = queuedByConnection.get(request.connectionKey); + if (!queue) return; + const index = queue.indexOf(request); + if (index >= 0) queue.splice(index, 1); + if (queue.length === 0) queuedByConnection.delete(request.connectionKey); +} + +function cancelRequest(request: QueuedOpen) { + if (request.finished) return; + request.controller.abort(); + runCancellation(request.cancel); + request.cancel = null; + if (!request.started) { + request.finished = true; + removeQueuedRequest(request); + requests.delete(request.id); + } +} + +function supersedeCurrentRequest() { + const current = requests.get(latestSupersedingRequestId); + if (current) cancelRequest(current); +} + +function finishRequest(request: QueuedOpen) { + if (request.finished) return; + request.finished = true; + request.cancel = null; + requests.delete(request.id); + const activeCount = activeCountByConnection.get(request.connectionKey) ?? 0; + if (activeCount <= 1) activeCountByConnection.delete(request.connectionKey); + else activeCountByConnection.set(request.connectionKey, activeCount - 1); + drainConnection(request.connectionKey); +} + +function executeRequest(request: QueuedOpen) { + if (!isCurrent(request)) { + cancelRequest(request); + return; + } + + request.started = true; + activeCountByConnection.set(request.connectionKey, (activeCountByConnection.get(request.connectionKey) ?? 0) + 1); + const openRequest: SidebarDataOpenRequest = { + isCurrent: () => isCurrent(request), + signal: request.controller.signal, registerCancel: (cancel) => { - if (requestGeneration !== generation) { + if (!isCurrent(request)) { runCancellation(cancel); return; } - activeGeneration = requestGeneration; - activeCancel = cancel; + request.cancel = cancel; }, }; + let result: void | Promise; try { - result = runner(request); + result = request.runner(openRequest); } catch { result = undefined; } void Promise.resolve(result) .catch(() => undefined) - .finally(() => { - if (activeGeneration !== requestGeneration) return; - activeGeneration = 0; - activeCancel = null; - }); + .finally(() => finishRequest(request)); } -export function runSidebarDataOpenImmediately(runner: OpenDataRunner) { - const requestGeneration = supersedeCurrentRequest(); - executeRequest(requestGeneration, runner); +function drainConnection(connectionKey: string) { + if ((activeCountByConnection.get(connectionKey) ?? 0) >= MAX_ACTIVE_PER_CONNECTION) return; + const queue = queuedByConnection.get(connectionKey); + const request = queue?.shift(); + if (!request) { + queuedByConnection.delete(connectionKey); + return; + } + if (queue?.length === 0) queuedByConnection.delete(connectionKey); + executeRequest(request); +} + +function enqueueRequest(options: SidebarDataOpenOptions, runner: OpenDataRunner) { + const superseding = options.supersede !== false; + if (superseding) supersedeCurrentRequest(); + + const request: QueuedOpen = { + id: ++nextRequestId, + connectionKey: options.connectionKey || DEFAULT_CONNECTION_KEY, + superseding, + controller: new AbortController(), + runner, + cancel: null, + started: false, + finished: false, + }; + if (superseding) latestSupersedingRequestId = request.id; + requests.set(request.id, request); + const queue = queuedByConnection.get(request.connectionKey) ?? []; + queue.push(request); + queuedByConnection.set(request.connectionKey, queue); + drainConnection(request.connectionKey); +} + +export function runSidebarDataOpenImmediately(runner: OpenDataRunner): void; +export function runSidebarDataOpenImmediately(options: SidebarDataOpenOptions, runner: OpenDataRunner): void; +export function runSidebarDataOpenImmediately(optionsOrRunner: SidebarDataOpenOptions | OpenDataRunner, maybeRunner?: OpenDataRunner) { + if (typeof optionsOrRunner === "function") { + enqueueRequest({ connectionKey: DEFAULT_CONNECTION_KEY }, optionsOrRunner); + return; + } + if (maybeRunner) enqueueRequest(optionsOrRunner, maybeRunner); } export function cancelPendingSidebarDataOpen() { - supersedeCurrentRequest(); + for (const request of [...requests.values()]) cancelRequest(request); + latestSupersedingRequestId = 0; +} + +export function sidebarDataOpenDiagnostics() { + return { + activeTasks: [...activeCountByConnection.values()].reduce((total, count) => total + count, 0), + queuedTasks: [...queuedByConnection.values()].reduce((total, queue) => total + queue.length, 0), + trackedRequests: requests.size, + }; } diff --git a/apps/desktop/src/lib/tabs/queryResultSize.ts b/apps/desktop/src/lib/tabs/queryResultSize.ts new file mode 100644 index 0000000000..4a77ffeb08 --- /dev/null +++ b/apps/desktop/src/lib/tabs/queryResultSize.ts @@ -0,0 +1,51 @@ +import type { QueryResult } from "@/types/database"; + +const ARRAY_OVERHEAD_BYTES = 24; +const OBJECT_OVERHEAD_BYTES = 32; +const VALUE_SLOT_BYTES = 8; + +function estimateValueBytes(value: unknown, seen: WeakSet): number { + if (value == null) return VALUE_SLOT_BYTES; + if (typeof value === "string") return VALUE_SLOT_BYTES + value.length * 2; + if (typeof value === "number" || typeof value === "bigint") return 16; + if (typeof value === "boolean") return 8; + if (value instanceof Uint8Array) return OBJECT_OVERHEAD_BYTES + value.byteLength; + if (value instanceof ArrayBuffer) return OBJECT_OVERHEAD_BYTES + value.byteLength; + if (value instanceof Date) return OBJECT_OVERHEAD_BYTES; + if (typeof value !== "object") return VALUE_SLOT_BYTES; + if (seen.has(value)) return VALUE_SLOT_BYTES; + seen.add(value); + if (Array.isArray(value)) { + return ARRAY_OVERHEAD_BYTES + value.reduce((total, item) => total + estimateValueBytes(item, seen), 0); + } + return OBJECT_OVERHEAD_BYTES + Object.entries(value).reduce((total, [key, entry]) => total + key.length * 2 + estimateValueBytes(entry, seen), 0); +} + +export function estimateQueryResultBytes(result: QueryResult | undefined): number { + if (!result) return 0; + return estimateValueBytes(result, new WeakSet()); +} + +export function estimateQueryResultsBytes(result: QueryResult | undefined, results: QueryResult[] | undefined): number { + if (results?.length) return results.reduce((total, item) => total + estimateQueryResultBytes(item), ARRAY_OVERHEAD_BYTES); + return estimateQueryResultBytes(result); +} + +export interface InactiveResultResidency { + id: string; + estimatedBytes: number; + accessedAt: number; +} + +export function selectInactiveResultEvictions(entries: InactiveResultResidency[], maxBytes: number, maxCount: number): string[] { + const ordered = [...entries].sort((left, right) => left.accessedAt - right.accessedAt); + let bytes = ordered.reduce((total, entry) => total + entry.estimatedBytes, 0); + const evicted: string[] = []; + while (ordered.length - evicted.length > maxCount || bytes > maxBytes) { + const entry = ordered[evicted.length]; + if (!entry) break; + evicted.push(entry.id); + bytes -= entry.estimatedBytes; + } + return evicted; +} diff --git a/apps/desktop/src/lib/tabs/tabResultCache.ts b/apps/desktop/src/lib/tabs/tabResultCache.ts index cafbd82505..180edaeb33 100644 --- a/apps/desktop/src/lib/tabs/tabResultCache.ts +++ b/apps/desktop/src/lib/tabs/tabResultCache.ts @@ -5,8 +5,12 @@ import { isTauriRuntime } from "@/lib/backend/tauriRuntime"; import { apiUrl } from "@/lib/common/webPath"; const DB_NAME = "dbx-tab-runtime-cache"; -const DB_VERSION = 1; +const DB_VERSION = 2; const RESULT_STORE = "resultSnapshots"; +const RESULT_METADATA_STORE = "resultSnapshotMetadata"; +const DEFAULT_PERSISTENT_CACHE_BYTES = 512 * 1024 * 1024; +const DEFAULT_ORPHAN_GRACE_MS = 24 * 60 * 60 * 1000; +const DEFAULT_MAX_AGE_MS = 30 * 24 * 60 * 60 * 1000; const PAYLOAD_MAGIC = "DBX_TAB_RESULT_CACHE"; const PAYLOAD_VERSION = 1; const PAYLOAD_CODEC = "msgpack-columnar"; @@ -80,6 +84,58 @@ interface TabResultCacheEnvelope { payload: TabResultSnapshotPayload; } +export type ResultCacheBackendName = "indexed-db" | "runtime"; + +export interface ResultCacheMetadata { + key: string; + rowCount: number; + columnCount: number; + byteSize: number; + createdAt: number; + lastAccessedAt: number; + ownerId?: string; +} + +export interface ResultCachePruneOptions { + liveKeys: string[]; + maxBytes: number; + orphanGraceMs: number; + maxAgeMs?: number; +} + +export interface ResultCachePruneResult { + deletedEntries: number; + deletedBytes: number; + orphanDeletions: number; + remainingEntries: number; + remainingBytes: number; +} + +export interface ResultCacheBackend { + name: ResultCacheBackendName; + available: () => boolean; + read: (key: string) => Promise; + write: (key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }, ownerId?: string) => Promise; + delete: (key: string) => Promise; + listMetadata: () => Promise; + prune: (options: ResultCachePruneOptions) => Promise; + deleteOwner: (ownerId: string) => Promise; +} + +const resultCacheDiagnostics = { + cacheBytes: 0, + evictions: 0, + serializationCount: 0, + serializationDurationMs: 0, + serializedBytes: 0, + orphanDeletions: 0, + corruptSnapshots: 0, +}; + +export function getResultCacheDiagnostics() { + return { ...resultCacheDiagnostics }; +} + function indexedDb(): IDBFactory | undefined { return typeof globalThis.indexedDB === "undefined" ? undefined : globalThis.indexedDB; } @@ -118,6 +174,7 @@ function openCacheDb(): Promise { request.onupgradeneeded = () => { const db = request.result; if (!db.objectStoreNames.contains(RESULT_STORE)) db.createObjectStore(RESULT_STORE); + if (!db.objectStoreNames.contains(RESULT_METADATA_STORE)) db.createObjectStore(RESULT_METADATA_STORE); }; request.onsuccess = () => resolve(request.result); request.onerror = () => { @@ -129,6 +186,129 @@ function openCacheDb(): Promise { return dbPromise; } +async function writeIndexedDbCache(key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }, ownerId?: string): Promise { + const db = await openCacheDb(); + if (!db) return false; + const previous = (await requestToPromise(db.transaction(RESULT_METADATA_STORE, "readonly").objectStore(RESULT_METADATA_STORE).get(key))) as ResultCacheMetadata | undefined; + const transaction = db.transaction([RESULT_STORE, RESULT_METADATA_STORE], "readwrite"); + const metadataStore = transaction.objectStore(RESULT_METADATA_STORE); + const now = Date.now(); + await Promise.all([ + requestToPromise(transaction.objectStore(RESULT_STORE).put(bytes, key)), + requestToPromise( + metadataStore.put( + { + key, + rowCount: stats.rowCount, + columnCount: stats.columnCount, + byteSize: bytes.byteLength, + createdAt: previous?.createdAt ?? now, + lastAccessedAt: now, + ownerId, + } satisfies ResultCacheMetadata, + key, + ), + ), + ]); + return true; +} + +async function readIndexedDbCache(key: string): Promise { + const db = await openCacheDb(); + if (!db) return undefined; + const value = await requestToPromise(db.transaction(RESULT_STORE, "readonly").objectStore(RESULT_STORE).get(key)); + if (isBinaryPayload(value)) { + const previous = (await requestToPromise(db.transaction(RESULT_METADATA_STORE, "readonly").objectStore(RESULT_METADATA_STORE).get(key))) as ResultCacheMetadata | undefined; + const store = db.transaction(RESULT_METADATA_STORE, "readwrite").objectStore(RESULT_METADATA_STORE); + const now = Date.now(); + await requestToPromise(store.put(previous ? { ...previous, lastAccessedAt: now } : { key, rowCount: 0, columnCount: 0, byteSize: value.byteLength, createdAt: now, lastAccessedAt: now }, key)); + } + return isBinaryPayload(value) ? new Uint8Array(value) : undefined; +} + +async function deleteIndexedDbCache(key: string): Promise { + const db = await openCacheDb(); + if (!db) return; + const transaction = db.transaction([RESULT_STORE, RESULT_METADATA_STORE], "readwrite"); + await Promise.all([requestToPromise(transaction.objectStore(RESULT_STORE).delete(key)), requestToPromise(transaction.objectStore(RESULT_METADATA_STORE).delete(key))]); +} + +async function listIndexedDbCacheMetadata(): Promise { + const db = await openCacheDb(); + if (!db) return []; + const transaction = db.transaction([RESULT_STORE, RESULT_METADATA_STORE], "readonly"); + const metadataRequest = transaction.objectStore(RESULT_METADATA_STORE).getAll(); + const keysRequest = transaction.objectStore(RESULT_STORE).getAllKeys(); + const [metadata, keys] = await Promise.all([requestToPromise(metadataRequest) as Promise, requestToPromise(keysRequest)]); + const byKey = new Map(metadata.map((entry) => [entry.key, entry])); + const missingKeys = keys.map(String).filter((key) => !byKey.has(key)); + const migratedMetadata: ResultCacheMetadata[] = []; + for (const key of missingKeys) { + const value = await requestToPromise(db.transaction(RESULT_STORE, "readonly").objectStore(RESULT_STORE).get(key)); + const now = Date.now(); + const entry: ResultCacheMetadata = { + key, + rowCount: 0, + columnCount: 0, + byteSize: isBinaryPayload(value) ? value.byteLength : 0, + createdAt: now, + lastAccessedAt: now, + }; + byKey.set(key, entry); + migratedMetadata.push(entry); + } + if (migratedMetadata.length) { + const migration = db.transaction(RESULT_METADATA_STORE, "readwrite").objectStore(RESULT_METADATA_STORE); + await Promise.all(migratedMetadata.map((entry) => requestToPromise(migration.put(entry, entry.key)))); + } + return [...byKey.values()]; +} + +export function selectResultCachePruneKeys(metadata: ResultCacheMetadata[], options: ResultCachePruneOptions, now = Date.now()): { keys: string[]; result: ResultCachePruneResult } { + const live = new Set(options.liveKeys); + const ordered = [...metadata].sort((left, right) => left.lastAccessedAt - right.lastAccessedAt || left.key.localeCompare(right.key)); + const deleted = new Set(); + let remainingBytes = ordered.reduce((total, entry) => total + entry.byteSize, 0); + let orphanDeletions = 0; + for (const entry of ordered) { + if (live.has(entry.key)) continue; + const orphanExpired = now - entry.createdAt >= Math.max(0, options.orphanGraceMs); + const ageExpired = options.maxAgeMs !== undefined && now - entry.lastAccessedAt >= Math.max(0, options.maxAgeMs); + if (!orphanExpired && !ageExpired) continue; + deleted.add(entry.key); + remainingBytes -= entry.byteSize; + if (orphanExpired) orphanDeletions += 1; + } + for (const entry of ordered) { + if (remainingBytes <= Math.max(0, options.maxBytes)) break; + if (live.has(entry.key) || deleted.has(entry.key)) continue; + deleted.add(entry.key); + remainingBytes -= entry.byteSize; + } + const deletedBytes = ordered.filter((entry) => deleted.has(entry.key)).reduce((total, entry) => total + entry.byteSize, 0); + return { + keys: [...deleted], + result: { + deletedEntries: deleted.size, + deletedBytes, + orphanDeletions, + remainingEntries: ordered.length - deleted.size, + remainingBytes, + }, + }; +} + +async function pruneIndexedDbCache(options: ResultCachePruneOptions): Promise { + const selection = selectResultCachePruneKeys(await listIndexedDbCacheMetadata(), options); + await Promise.all(selection.keys.map(deleteIndexedDbCache)); + return selection.result; +} + +async function deleteIndexedDbCacheOwner(ownerId: string): Promise { + const metadata = await listIndexedDbCacheMetadata(); + await Promise.all(metadata.filter((entry) => entry.ownerId === ownerId).map((entry) => deleteIndexedDbCache(entry.key))); +} + function clonePlain(value: T): T { const raw = toRaw(value); if (typeof structuredClone === "function") return structuredClone(raw); @@ -272,7 +452,7 @@ function canUseRemoteRuntimeCache(): boolean { return typeof btoa !== "undefined" && typeof atob !== "undefined" && (isTauriRuntime() || typeof fetch !== "undefined"); } -async function writeRemoteRuntimeCache(key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }): Promise { +async function writeRemoteRuntimeCache(key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }, ownerId?: string): Promise { if (!canUseRemoteRuntimeCache()) return false; try { if (isTauriRuntime()) { @@ -282,6 +462,7 @@ async function writeRemoteRuntimeCache(key: string, bytes: Uint8Array, stats: { payloadBase64: bytesToBase64(bytes), rowCount: stats.rowCount, columnCount: stats.columnCount, + ownerId, }); return true; } @@ -293,6 +474,7 @@ async function writeRemoteRuntimeCache(key: string, bytes: Uint8Array, stats: { payloadBase64: bytesToBase64(bytes), rowCount: stats.rowCount, columnCount: stats.columnCount, + ownerId, }), }); return response.ok; @@ -302,6 +484,53 @@ async function writeRemoteRuntimeCache(key: string, bytes: Uint8Array, stats: { } } +function normalizeRuntimeMetadata(entry: Record): ResultCacheMetadata { + return { + key: String(entry.key ?? ""), + rowCount: Number(entry.rowCount ?? 0), + columnCount: Number(entry.columnCount ?? 0), + byteSize: Number(entry.byteSize ?? 0), + createdAt: Number(entry.createdAt ?? 0), + lastAccessedAt: Number(entry.lastAccessedAt ?? 0), + ownerId: typeof entry.ownerId === "string" ? entry.ownerId : undefined, + }; +} + +async function listRemoteRuntimeCacheMetadata(): Promise { + if (!canUseRemoteRuntimeCache()) return []; + if (isTauriRuntime()) { + const { invoke } = await import("@tauri-apps/api/core"); + const entries = await invoke[]>("list_tab_runtime_cache_metadata"); + return entries.map(normalizeRuntimeMetadata); + } + const response = await fetch(apiUrl("/api/tab-runtime-cache/metadata")); + if (!response.ok) return []; + return ((await response.json()) as Record[]).map(normalizeRuntimeMetadata); +} + +async function pruneRemoteRuntimeCache(options: ResultCachePruneOptions): Promise { + if (isTauriRuntime()) { + const { invoke } = await import("@tauri-apps/api/core"); + return invoke("prune_tab_runtime_cache", { request: options }); + } + const response = await fetch(apiUrl("/api/tab-runtime-cache/prune"), { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(options), + }); + if (!response.ok) throw new Error(`Runtime cache prune failed: ${response.status}`); + return response.json() as Promise; +} + +async function deleteRemoteRuntimeCacheOwner(ownerId: string): Promise { + if (isTauriRuntime()) { + const { invoke } = await import("@tauri-apps/api/core"); + await invoke("delete_tab_runtime_cache_owner", { ownerId }); + return; + } + await fetch(apiUrl(`/api/tab-runtime-cache/owner?owner_id=${encodeURIComponent(ownerId)}`), { method: "DELETE" }); +} + async function readRemoteRuntimeCache(key: string): Promise { if (!canUseRemoteRuntimeCache()) return undefined; try { @@ -334,12 +563,82 @@ async function deleteRemoteRuntimeCache(key: string): Promise { } } -function scheduleRemoteRuntimeCacheWrite(key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }, version: number) { - window.setTimeout(async () => { - if (!isCurrentCacheKeyVersion(key, version)) return; - await writeRemoteRuntimeCache(key, bytes, stats); - clearCacheKeyVersionIfCurrent(key, version); - }, 0); +const indexedDbBackend: ResultCacheBackend = { + name: "indexed-db", + available: () => indexedDb() !== undefined, + read: readIndexedDbCache, + write: writeIndexedDbCache, + delete: deleteIndexedDbCache, + listMetadata: listIndexedDbCacheMetadata, + prune: pruneIndexedDbCache, + deleteOwner: deleteIndexedDbCacheOwner, +}; + +const runtimeBackend: ResultCacheBackend = { + name: "runtime", + available: canUseRemoteRuntimeCache, + read: readRemoteRuntimeCache, + write: writeRemoteRuntimeCache, + delete: deleteRemoteRuntimeCache, + listMetadata: listRemoteRuntimeCacheMetadata, + prune: pruneRemoteRuntimeCache, + deleteOwner: deleteRemoteRuntimeCacheOwner, +}; + +export interface ResultCacheRuntimeConfig { + primary: ResultCacheBackendName; + fallbackEnabled: boolean; +} + +export function resultCacheRuntimeConfig(tauriRuntime = isTauriRuntime(), environment = import.meta.env): ResultCacheRuntimeConfig { + const configured = environment.VITE_DBX_RESULT_CACHE_BACKEND; + const primary = tauriRuntime || configured === "http" || configured === "runtime" ? "runtime" : "indexed-db"; + return { primary, fallbackEnabled: environment.VITE_DBX_RESULT_CACHE_FALLBACK !== "false" }; +} + +export function resultCacheBackendOrder(tauriRuntime = isTauriRuntime(), config = resultCacheRuntimeConfig(tauriRuntime)): ResultCacheBackendName[] { + return config.primary === "runtime" ? ["runtime", "indexed-db"] : ["indexed-db", "runtime"]; +} + +function availableResultCacheBackends(includeLegacyFallback = true): ResultCacheBackend[] { + const byName: Record = { + "indexed-db": indexedDbBackend, + runtime: runtimeBackend, + }; + const config = resultCacheRuntimeConfig(); + const order = resultCacheBackendOrder(isTauriRuntime(), config); + const selected = includeLegacyFallback || config.fallbackEnabled ? order : order.slice(0, 1); + return selected.map((name) => byName[name]).filter((backend) => backend.available()); +} + +export async function writeResultCacheBackends(backends: ResultCacheBackend[], key: string, bytes: Uint8Array, stats: { rowCount: number; columnCount: number }, isCurrent: () => boolean = () => true, ownerId?: string): Promise { + for (const backend of backends) { + if (!isCurrent()) return false; + try { + if (await backend.write(key, bytes, stats, ownerId)) return true; + } catch (error) { + console.warn("[DBX][tab-result-cache:write:error]", { key, backend: backend.name, error }); + } + } + return false; +} + +export async function readResultCacheBackends(backends: ResultCacheBackend[], key: string): Promise<{ bytes: Uint8Array; backend: ResultCacheBackend } | undefined> { + for (const backend of backends) { + try { + const bytes = await backend.read(key); + if (bytes) return { bytes, backend }; + } catch (error) { + console.warn("[DBX][tab-result-cache:read:error]", { key, backend: backend.name, error }); + } + } + return undefined; +} + +export async function promoteFallbackResultCacheRead(backends: ResultCacheBackend[], key: string, cached: { bytes: Uint8Array; backend: ResultCacheBackend }, stats: { rowCount: number; columnCount: number }, ownerId?: string): Promise { + const primary = backends[0]; + if (!primary || cached.backend === primary) return false; + return primary.write(key, cached.bytes, stats, ownerId); } function isRecord(value: unknown): value is Record { @@ -375,7 +674,13 @@ export function encodeTabResultSnapshot(snapshot: TabResultSnapshot): Uint8Array } export function decodeTabResultSnapshot(bytes: Uint8Array | ArrayBuffer): TabResultSnapshot | undefined { - const decoded = decode(bytes instanceof Uint8Array ? bytes : new Uint8Array(bytes)); + let decoded: unknown; + try { + decoded = decode(bytes instanceof Uint8Array ? bytes : new Uint8Array(bytes)); + } catch { + resultCacheDiagnostics.corruptSnapshots += 1; + return undefined; + } if (!isRecord(decoded)) return undefined; if (decoded.magic !== PAYLOAD_MAGIC || decoded.version !== PAYLOAD_VERSION || decoded.codec !== PAYLOAD_CODEC) { return undefined; @@ -413,57 +718,48 @@ export function buildTabResultSnapshot(tab: QueryTab): TabResultSnapshot | undef }; } -export async function writeTabResultSnapshot(key: string, snapshot: TabResultSnapshot | undefined): Promise { +function afterActiveResultPaint(): Promise { + if (typeof window === "undefined" || typeof requestAnimationFrame === "undefined") return Promise.resolve(); + return new Promise((resolve) => { + let settled = false; + const finish = () => { + if (settled) return; + settled = true; + resolve(); + }; + requestAnimationFrame(finish); + // Background WebViews and test DOMs may throttle animation frames indefinitely. + setTimeout(finish, 100); + }); +} + +export async function writeTabResultSnapshot(key: string, snapshot: TabResultSnapshot | undefined, ownerId?: string): Promise { if (!snapshot) return false; const version = bumpCacheKeyVersion(key); - const encoded = encodeTabResultSnapshot(snapshot); - const stats = resultStats(snapshot); - let wroteLocal = false; - try { - const db = await openCacheDb(); - if (db) { - const tx = db.transaction(RESULT_STORE, "readwrite"); - await requestToPromise(tx.objectStore(RESULT_STORE).put(encoded, key)); - wroteLocal = true; - } - } catch (error) { - console.warn("[DBX][tab-result-cache:write:error]", { key, error }); - } - if (wroteLocal) { - if (typeof window === "undefined") { - void writeRemoteRuntimeCache(key, encoded, stats).finally(() => clearCacheKeyVersionIfCurrent(key, version)); - } else { - scheduleRemoteRuntimeCacheWrite(key, encoded, stats, version); - } - return true; - } + await afterActiveResultPaint(); if (!isCurrentCacheKeyVersion(key, version)) return false; + const startedAt = typeof performance === "undefined" ? Date.now() : performance.now(); + const encoded = encodeTabResultSnapshot(snapshot); + const finishedAt = typeof performance === "undefined" ? Date.now() : performance.now(); + resultCacheDiagnostics.serializationCount += 1; + resultCacheDiagnostics.serializationDurationMs += finishedAt - startedAt; + resultCacheDiagnostics.serializedBytes += encoded.byteLength; + const stats = resultStats(snapshot); try { - return await writeRemoteRuntimeCache(key, encoded, stats); + return await writeResultCacheBackends(availableResultCacheBackends(false), key, encoded, stats, () => isCurrentCacheKeyVersion(key, version), ownerId); } finally { clearCacheKeyVersionIfCurrent(key, version); } } export async function readTabResultSnapshot(key: string): Promise { - try { - const db = await openCacheDb(); - if (db) { - const value = await requestToPromise(db.transaction(RESULT_STORE, "readonly").objectStore(RESULT_STORE).get(key)); - if (isBinaryPayload(value)) return decodeTabResultSnapshot(value); - if (value) return value as TabResultSnapshot; - } - const remoteBytes = await readRemoteRuntimeCache(key); - const remoteSnapshot = remoteBytes ? decodeTabResultSnapshot(remoteBytes) : undefined; - if (remoteBytes && remoteSnapshot && db) { - void requestToPromise(db.transaction(RESULT_STORE, "readwrite").objectStore(RESULT_STORE).put(remoteBytes, key)); - } - return remoteSnapshot; - } catch (error) { - console.warn("[DBX][tab-result-cache:read:error]", { key, error }); - const remoteBytes = await readRemoteRuntimeCache(key); - return remoteBytes ? decodeTabResultSnapshot(remoteBytes) : undefined; - } + const backends = availableResultCacheBackends(true); + const cached = await readResultCacheBackends(backends, key); + if (!cached) return undefined; + const snapshot = decodeTabResultSnapshot(cached.bytes); + if (!snapshot) return undefined; + void promoteFallbackResultCacheRead(backends, key, cached, resultStats(snapshot)); + return snapshot; } export async function deleteTabResultSnapshot(key: string): Promise { @@ -474,13 +770,29 @@ export async function deleteTabResultSnapshot(key: string): Promise { else setTimeout(cleanup, 5000); }; try { - const db = await openCacheDb(); - if (db) await requestToPromise(db.transaction(RESULT_STORE, "readwrite").objectStore(RESULT_STORE).delete(key)); - await deleteRemoteRuntimeCache(key); + await Promise.all(availableResultCacheBackends().map((backend) => backend.delete(key))); clearVersionLater(); } catch (error) { console.warn("[DBX][tab-result-cache:delete:error]", { key, error }); - await deleteRemoteRuntimeCache(key); clearVersionLater(); } } + +export async function pruneTabResultSnapshots(liveKeys: Iterable, options: Partial> = {}): Promise { + const primary = availableResultCacheBackends(false)[0]; + if (!primary) return undefined; + const result = await primary.prune({ + liveKeys: [...new Set(liveKeys)], + maxBytes: options.maxBytes ?? DEFAULT_PERSISTENT_CACHE_BYTES, + orphanGraceMs: options.orphanGraceMs ?? DEFAULT_ORPHAN_GRACE_MS, + maxAgeMs: options.maxAgeMs ?? DEFAULT_MAX_AGE_MS, + }); + resultCacheDiagnostics.cacheBytes = result.remainingBytes; + resultCacheDiagnostics.evictions += result.deletedEntries; + resultCacheDiagnostics.orphanDeletions += result.orphanDeletions; + return result; +} + +export async function deleteTabResultSnapshotsForOwner(ownerId: string): Promise { + await Promise.all(availableResultCacheBackends(true).map((backend) => backend.deleteOwner(ownerId))); +} diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 114a62ccde..6f2278133f 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -51,6 +51,7 @@ import { findDatabaseTreeNode } from "@/lib/sidebar/treeRefreshTarget"; import { shouldMarkDisconnected } from "@/lib/connection/connectionHealth"; import { connectionAttemptOriginalErrorMessage, connectionAttemptTimeoutMessage, connectionAttemptTimeoutMs } from "@/lib/connection/connectionAttemptTimeout"; import { requiresSqlServerLegacyCompatibilityComponent, SQLSERVER_LEGACY_COMPATIBILITY_DRIVER_KEY } from "@/lib/connection/sqlServerLegacyCompatibility"; +import { deleteTabResultSnapshotsForOwner } from "@/lib/tabs/tabResultCache"; import { connectionUsesVisibleSchemaFilter, filterDatabaseNamesForConnection, filterSchemaNamesForConnection, filterVisibleDatabaseNames, normalizeVisibleDatabaseSelection } from "@/lib/database/visibleDatabases"; import { buildObjectGroupPlaceholderNodes, @@ -1756,6 +1757,7 @@ export const useConnectionStore = defineStore("connection", () => { for (const id of removedIds) { invalidateCompletionCache(id); clearLoadedChildrenCache(id); + void deleteTabResultSnapshotsForOwner(id); } } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 145fc5d7d7..a6db3852ff 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -44,7 +44,8 @@ import { normalizeResultPageSize } from "@/lib/dataGrid/paginationPageSize"; import { executableStatementRanges, splitSqlStatementRanges } from "@/lib/sql/sqlStatementRanges"; import { externalSqlFileDisplayTitles, normalizeExternalSqlPath } from "@/lib/sql/sqlFileOpen"; import { clearDataGridPendingSnapshotsForTab } from "@/composables/useDataGridEditor"; -import { buildTabResultSnapshot, deleteTabResultSnapshot, readTabResultSnapshot, tabResultCacheKey, writeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; +import { buildTabResultSnapshot, deleteTabResultSnapshot, pruneTabResultSnapshots, readTabResultSnapshot, tabResultCacheKey, writeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; +import { estimateQueryResultsBytes, selectInactiveResultEvictions } from "@/lib/tabs/queryResultSize"; import { queryResultBaseSql, queryResultExecutionSql } from "@/lib/tabs/tabPresentation"; import { isMysqlExecutionErrorResult } from "@/lib/query/queryResultError"; import { decodeQueryResultArchive, encodeQueryResultArchive, type DecodedQueryResultArchive } from "@/lib/query/queryResultArchive"; @@ -52,6 +53,8 @@ import * as api from "@/lib/backend/api"; import { useConnectionStore } from "@/stores/connectionStore"; import { useSettingsStore } from "@/stores/settingsStore"; import { useSavedSqlStore } from "@/stores/savedSqlStore"; +import { recordQueryCancellationLatency, resourceLifecycleDiagnostics } from "@/lib/diagnostics/resourceLifecycleDiagnostics"; +import { appendDebugLog } from "@/lib/backend/debugLog"; import { createSavedSqlEditorPosition, initSavedSqlEditorPositions, restoreSavedSqlEditorPosition, saveSavedSqlEditorPosition } from "@/lib/app/savedSqlEditorPosition"; import { ensureSqlExtension } from "@/lib/savedSql/savedSqlFileName"; import { safeLocalStorageGet, safeLocalStorageRemove } from "@/lib/backend/safeStorage"; @@ -416,6 +419,11 @@ export const useQueryStore = defineStore("query", () => { return tableStructureRefreshVersions.value[tableStructureKey(connectionId, database, schema, tableName)] ?? 0; } const MAX_CACHED_RESULTS = 5; + const MAX_CACHED_RESULT_BYTES = 128 * 1024 * 1024; + + function queryExecutionLog(level: "debug" | "info" | "warn" | "error", event: string, details: Record) { + appendDebugLog(level, `[DBX][executeTabSql:${event}]`, details); + } async function closeResultSession(tab: QueryTab | undefined, preserveSessionId?: string, throwOnError = false) { const sessionId = tab?.resultSessionId ?? tab?.result?.session_id; @@ -481,6 +489,7 @@ export const useQueryStore = defineStore("query", () => { function touchResult(tab: QueryTab | undefined, accessedAt = Date.now()) { if (tab?.result || tab?.results) { tab.resultAccessedAt = accessedAt; + tab.resultEstimatedBytes = estimateQueryResultsBytes(tab.result, tab.results); tab.resultCacheState = "memory"; tab.resultEvicted = undefined; } @@ -496,6 +505,7 @@ export const useQueryStore = defineStore("query", () => { tab.resultSortMode = undefined; tab.resultSessionId = undefined; tab.resultAccessedAt = undefined; + tab.resultEstimatedBytes = undefined; tab.queryAnalysis = undefined; tab.querySourceColumns = undefined; tab.queryEditabilityReason = undefined; @@ -521,6 +531,7 @@ export const useQueryStore = defineStore("query", () => { run.resultLocalSortOriginalRows = undefined; run.resultLocalSortOriginalMongoDocuments = undefined; run.resultSessionId = undefined; + run.resultEstimatedBytes = undefined; run.queryAnalysis = undefined; run.querySourceColumns = undefined; run.queryEditabilityReason = undefined; @@ -556,6 +567,7 @@ export const useQueryStore = defineStore("query", () => { tab.resultAccessedAt = run.resultAccessedAt; tab.resultCacheKey = run.resultCacheKey; tab.resultCacheState = run.resultCacheState; + tab.resultEstimatedBytes = run.resultEstimatedBytes ?? estimateQueryResultsBytes(run.result, run.results); tab.resultEvicted = run.resultEvicted; tab.queryAnalysis = run.queryAnalysis; tab.querySourceColumns = run.querySourceColumns; @@ -632,24 +644,28 @@ export const useQueryStore = defineStore("query", () => { const key = run.resultCacheKey ?? resultRunCacheKey(tab.id, run.id); run.resultCacheKey = key; run.resultCacheState = "memory"; - return writeTabResultSnapshot(key, { - result: run.result, - results: run.results, - activeResultIndex: run.activeResultIndex, - resultEditorFingerprint: run.resultEditorFingerprint, - resultRuns: [run], - activeResultRunId: run.id, - queryAnalysis: run.queryAnalysis, - querySourceColumns: run.querySourceColumns, - queryEditabilityReason: run.queryEditabilityReason, - tableMeta: run.tableMeta, - resultPageSql: run.resultPageSql, - resultPageLimit: run.resultPageLimit, - resultPageOffset: run.resultPageOffset, - resultCountSql: run.resultCountSql, - resultTotalRowCount: run.resultTotalRowCount, - cachedAt: Date.now(), - }); + return writeTabResultSnapshot( + key, + { + result: run.result, + results: run.results, + activeResultIndex: run.activeResultIndex, + resultEditorFingerprint: run.resultEditorFingerprint, + resultRuns: [run], + activeResultRunId: run.id, + queryAnalysis: run.queryAnalysis, + querySourceColumns: run.querySourceColumns, + queryEditabilityReason: run.queryEditabilityReason, + tableMeta: run.tableMeta, + resultPageSql: run.resultPageSql, + resultPageLimit: run.resultPageLimit, + resultPageOffset: run.resultPageOffset, + resultCountSql: run.resultCountSql, + resultTotalRowCount: run.resultTotalRowCount, + cachedAt: Date.now(), + }, + tab.connectionId, + ); } function evictInactiveResultRunPayloads(tab: QueryTab) { @@ -698,6 +714,7 @@ export const useQueryStore = defineStore("query", () => { resultTotalRowCountLoading: tab.resultTotalRowCountLoading, resultSessionId: tab.resultSessionId, resultAccessedAt: tab.resultAccessedAt, + resultEstimatedBytes: tab.resultEstimatedBytes, resultCacheKey: tab.resultCacheKey, resultCacheState: tab.resultCacheState, resultEvicted: tab.resultEvicted, @@ -750,6 +767,7 @@ export const useQueryStore = defineStore("query", () => { resultTotalRowCountLoading: tab.resultTotalRowCountLoading, resultSessionId: tab.resultSessionId, resultAccessedAt: tab.resultAccessedAt, + resultEstimatedBytes: tab.resultEstimatedBytes, resultCacheKey: tab.resultCacheKey, resultCacheState: tab.resultCacheState, resultEvicted: tab.resultEvicted, @@ -823,7 +841,7 @@ export const useQueryStore = defineStore("query", () => { async function evictCachedResult(tab: QueryTab) { await closeResultSession(tab); const cacheKey = tabResultCacheKey(tab.id); - const cached = await writeTabResultSnapshot(cacheKey, buildTabResultSnapshot(tab)); + const cached = await writeTabResultSnapshot(cacheKey, buildTabResultSnapshot(tab), tab.connectionId); tab.resultCacheKey = cached ? cacheKey : undefined; tab.resultCacheState = cached ? "disk" : "missing"; clearResultPayload(tab, { evicted: true }); @@ -838,6 +856,16 @@ export const useQueryStore = defineStore("query", () => { } } + function scheduleResultCacheMaintenance() { + const maintain = () => { + const liveKeys = tabs.value.flatMap((tab) => [tab.resultCacheKey, ...(tab.resultRuns?.map((run) => run.resultCacheKey) ?? [])]).filter((key): key is string => !!key); + void pruneTabResultSnapshots(liveKeys).catch((error) => console.warn("[DBX][result-cache:maintenance:error]", error)); + }; + if (typeof requestIdleCallback !== "undefined") requestIdleCallback(maintain, { timeout: 5000 }); + else if (typeof window !== "undefined") window.setTimeout(maintain, 0); + else setTimeout(maintain, 0); + } + async function initOpenTabs() { if (isOpenTabsLoaded.value) return; const saved = await api.loadOpenTabsState().catch(() => null); @@ -851,6 +879,7 @@ export const useQueryStore = defineStore("query", () => { await saveTabs(tabs.value, activeTabId.value).catch(() => undefined); } isOpenTabsLoaded.value = true; + scheduleResultCacheMaintenance(); return; } @@ -863,6 +892,7 @@ export const useQueryStore = defineStore("query", () => { // would resurrect old tabs if the user later changes the setting. clearLegacySavedTabs(); isOpenTabsLoaded.value = true; + scheduleResultCacheMaintenance(); return; } try { @@ -875,6 +905,7 @@ export const useQueryStore = defineStore("query", () => { } } isOpenTabsLoaded.value = true; + scheduleResultCacheMaintenance(); } const _persistSnapshot = computed(() => @@ -1452,6 +1483,7 @@ export const useQueryStore = defineStore("query", () => { void closeResultSession(tabs.value[idx]); void closeClientConnectionSession(tabs.value[idx]); clearResultRunSnapshots(tabs.value[idx]); + void deleteTabResultSnapshot(tabResultCacheKey(id)); clearResultPayload(tabs.value[idx]); tabs.value.splice(idx, 1); if (tab.externalSqlPath) refreshExternalSqlFileTitles(); @@ -1703,6 +1735,7 @@ export const useQueryStore = defineStore("query", () => { void closeResultSession(tab); void closeClientConnectionSession(tab); clearResultRunSnapshots(tab); + void deleteTabResultSnapshot(tabResultCacheKey(tab.id)); clearResultPayload(tab); }); @@ -2265,7 +2298,7 @@ export const useQueryStore = defineStore("query", () => { tableName: metadataTableName, }; const knownTableType = tab.tableMeta?.tableName.toLowerCase() === metadataTableName.toLowerCase() && normalizeOptionalSchema(tab.tableMeta.schema) === normalizeOptionalSchema(metadataSchema) ? tab.tableMeta.tableType : undefined; - console.info("[DBX][executeTabSql:metadata:table:start]", { + queryExecutionLog("info", "metadata:table:start", { traceId, schema: metadataSchema, table: metadataTableName, @@ -2281,11 +2314,11 @@ export const useQueryStore = defineStore("query", () => { databaseType: dbType, driverProfile: conn?.driver_profile || conn?.db_type, catalog: metadataCatalog, - traceLogger: (event) => console.debug("[DBX][executeTabSql:metadata:table-trace]", { sourceTraceId: traceId, ...event }), + traceLogger: (event) => queryExecutionLog("debug", "metadata:table-trace", { sourceTraceId: traceId, ...event }), }); const columns = loadedMetadata.metadata.columns; const primaryKeys = loadedMetadata.metadata.primaryKeys; - console.info("[DBX][executeTabSql:metadata:table:done]", { + queryExecutionLog("info", "metadata:table:done", { traceId, columnCount: columns.length, primaryKeyCount: primaryKeys.length, @@ -2356,7 +2389,7 @@ export const useQueryStore = defineStore("query", () => { sourceExpressions: databaseType === "oracle" && missingPrimaryKeys.includes(DBX_ROWID_COLUMN) ? { [DBX_ROWID_COLUMN]: "ROWIDTOCHAR(ROWID)" } : undefined, }); if (!rewritten) return unchanged; - console.info("[DBX][executeTabSql:hidden-primary-keys]", { + queryExecutionLog("info", "hidden-primary-keys", { traceId, table: loaded.tableMeta.tableName, keyCount: rewritten.projections.length, @@ -2366,7 +2399,7 @@ export const useQueryStore = defineStore("query", () => { } catch (error) { // Metadata enrichment is optional. Query execution must retain its prior // behavior when metadata is unavailable or the SQL cannot be rewritten. - console.warn("[DBX][executeTabSql:hidden-primary-keys:skip]", { traceId, error, elapsed: elapsed() }); + queryExecutionLog("warn", "hidden-primary-keys:skip", { traceId, error, elapsed: elapsed() }); return unchanged; } } @@ -2382,9 +2415,9 @@ export const useQueryStore = defineStore("query", () => { }; } - console.info("[DBX][executeTabSql:metadata:editability:start]", { traceId, elapsed: elapsed?.() }); + queryExecutionLog("info", "metadata:editability:start", { traceId, elapsed: elapsed?.() }); const editability = await api.analyzeEditableQueryEditability(sql); - console.info("[DBX][executeTabSql:metadata:editability:done]", { + queryExecutionLog("info", "metadata:editability:done", { traceId, editable: editability.editable, reason: editability.editable ? undefined : editability.reason, @@ -2533,7 +2566,7 @@ export const useQueryStore = defineStore("query", () => { void (async () => { const tab = tabs.value.find((t) => t.id === tabId); if (!tab || tab.result !== result) return; - console.info("[DBX][executeTabSql:metadata:start]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "metadata:start", { traceId, elapsed: elapsed() }); const patch = await buildQueryMetadataPatch(tab, sql, traceId, elapsed, hiddenPrimaryKeys); if (patch?.queryAnalysis && hiddenPrimaryKeys.length > 0) { patch.queryAnalysis = { ...patch.queryAnalysis, allowInsert: false }; @@ -2542,9 +2575,9 @@ export const useQueryStore = defineStore("query", () => { if (patch && current?.result === result) { applyQueryMetadataPatch(current, patch); syncActiveResultRunFromDisplayed(current); - console.info("[DBX][executeTabSql:metadata:done]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "metadata:done", { traceId, elapsed: elapsed() }); } else { - console.warn("[DBX][executeTabSql:metadata:stale]", { traceId, elapsed: elapsed() }); + queryExecutionLog("warn", "metadata:stale", { traceId, elapsed: elapsed() }); } })(); } @@ -2604,7 +2637,7 @@ export const useQueryStore = defineStore("query", () => { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); return; } - console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() }); + queryExecutionLog("info", "count:start", { traceId: options.traceId, elapsed: options.elapsed() }); const countResult = await api.executeQuery(options.connectionId, options.database, countTarget.sql, countTarget.schema, countExecutionId, { clientSessionId, timeoutSecs: options.timeoutSecs, @@ -2615,14 +2648,14 @@ export const useQueryStore = defineStore("query", () => { return; } setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, total); - console.info("[DBX][executeTabSql:count:done]", { + queryExecutionLog("info", "count:done", { traceId: options.traceId, total, elapsed: options.elapsed(), }); } catch (error) { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); - console.warn("[DBX][executeTabSql:count:error]", { + queryExecutionLog("warn", "count:error", { traceId: options.traceId, elapsed: options.elapsed(), error, @@ -2686,15 +2719,12 @@ export const useQueryStore = defineStore("query", () => { if (!options?.preserveResultDuringExecution || !tab.result) { clearResultPayload(tab); } - console.info("[DBX][executeTabSql:start]", { + queryExecutionLog("info", "start", { traceId, tabId: id, mode: tab.mode, - connectionId: tab.connectionId, - database: tab.database, - schema: tab.schema, sourceTraceId: options?.sourceTraceId, - sql, + sqlLength: sql.length, }); const queryBaseSql = options?.resultBaseSql ?? sql; let sqlToExecute = sql; @@ -2715,13 +2745,13 @@ export const useQueryStore = defineStore("query", () => { const mongoNeedsConnection = mongoCommands.some(({ command }) => command.kind !== "use"); if (options?.skipEnsureConnected) { - console.info("[DBX][executeTabSql:ensure-connected:skip]", { traceId, elapsed: elapsed(), reason: "caller" }); + queryExecutionLog("info", "ensure-connected:skip", { traceId, elapsed: elapsed(), reason: "caller" }); } else if (conn?.db_type === "mongodb" && mongoCommands.length > 0 && !mongoNeedsConnection) { - console.info("[DBX][executeTabSql:ensure-connected:skip]", { traceId, elapsed: elapsed(), reason: "mongo-use-only" }); + queryExecutionLog("info", "ensure-connected:skip", { traceId, elapsed: elapsed(), reason: "mongo-use-only" }); } else { - console.info("[DBX][executeTabSql:ensure-connected:start]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "ensure-connected:start", { traceId, elapsed: elapsed() }); await connStore.ensureConnected(tab.connectionId); - console.info("[DBX][executeTabSql:ensure-connected:done]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "ensure-connected:done", { traceId, elapsed: elapsed() }); } conn = connStore.getConfig(tab.connectionId); if (parsedMongoCommands === undefined && conn?.db_type === "mongodb") { @@ -2732,9 +2762,9 @@ export const useQueryStore = defineStore("query", () => { const useAgentCursor = usesAgentCursorForQuery(conn?.db_type); const queryTimeoutSecs = queryTimeoutSecsForConnection(conn); const settingsStore = useSettingsStore(); - console.info("[DBX][executeTabSql:previous-session-close:start]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "previous-session-close:start", { traceId, elapsed: elapsed() }); await previousResultSessionClose; - console.info("[DBX][executeTabSql:previous-session-close:done]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "previous-session-close:done", { traceId, elapsed: elapsed() }); // Redis command execution — split multi-line input into individual commands if (conn?.db_type === "redis") { @@ -2745,7 +2775,7 @@ export const useQueryStore = defineStore("query", () => { .map((line) => line.trim()) .filter((line) => line.length > 0); if (commands.length === 0) return; - console.info("[DBX][executeTabSql:redis:start]", { traceId, db: currentDb, commandCount: commands.length, sql }); + queryExecutionLog("info", "redis:start", { traceId, db: currentDb, commandCount: commands.length, sqlLength: sql.length }); const allResults: QueryResult[] = []; const commandRanges = executableStatementRanges(sql, "redis"); @@ -2769,7 +2799,7 @@ export const useQueryStore = defineStore("query", () => { allResults.push(annotateQueryResultSource({ columns: ["Error"], rows: [[e?.message ?? String(e)]], affected_rows: 0, execution_time_ms: 0 }, command, undefined, undefined, sourceRange)); } } - console.info("[DBX][executeTabSql:redis:done]", { traceId, commandCount: commands.length, elapsed: elapsed() }); + queryExecutionLog("info", "redis:done", { traceId, commandCount: commands.length, elapsed: elapsed() }); const current = tabs.value.find((t) => t.id === id); if (current?.executionId === executionId) { @@ -2809,7 +2839,7 @@ export const useQueryStore = defineStore("query", () => { } if (mongoCommands.length > 0) { - console.info("[DBX][executeTabSql:mongo:start]", { traceId, database: tab.database, commandCount: mongoCommands.length, sql }); + queryExecutionLog("info", "mongo:start", { traceId, commandCount: mongoCommands.length, sqlLength: sql.length }); const allResults: QueryResult[] = []; // Track the effective db as we walk the batch so later commands observe @@ -2832,12 +2862,12 @@ export const useQueryStore = defineStore("query", () => { try { switch (mongoCommand.kind) { case "find": { - console.info("[DBX][executeTabSql:mongo-find:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); + queryExecutionLog("info", "mongo-find:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.skip, mongoCommand.limit, mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, executionId); const queryResult = markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total))); allResults.push(queryResult); mongoEditTarget = mongoCommands.length === 1 && queryResult.columns.includes("_id") ? { collection: mongoCommand.collection, idColumn: "_id" } : undefined; - console.info("[DBX][executeTabSql:mongo-find:done]", { + queryExecutionLog("info", "mongo-find:done", { traceId, collection: mongoCommand.collection, database: currentDatabase, @@ -2848,12 +2878,12 @@ export const useQueryStore = defineStore("query", () => { break; } case "findOne": { - console.info("[DBX][executeTabSql:mongo-find-one:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); + queryExecutionLog("info", "mongo-find-one:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); const result = await api.mongoFindOne(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.projection, mongoCommand.options, executionId); const queryResult = markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total))); allResults.push(queryResult); mongoEditTarget = mongoCommands.length === 1 && queryResult.columns.includes("_id") ? { collection: mongoCommand.collection, idColumn: "_id" } : undefined; - console.info("[DBX][executeTabSql:mongo-find-one:done]", { + queryExecutionLog("info", "mongo-find-one:done", { traceId, collection: mongoCommand.collection, database: currentDatabase, @@ -2863,11 +2893,11 @@ export const useQueryStore = defineStore("query", () => { break; } case "version": { - console.info("[DBX][executeTabSql:mongo-version:start]", { traceId, database: currentDatabase }); + queryExecutionLog("info", "mongo-version:start", { traceId, database: currentDatabase }); const version = await api.mongoServerVersion(tab.connectionId, currentDatabase, executionId); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoVersionToQueryResult(version, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-version:done]", { + queryExecutionLog("info", "mongo-version:done", { traceId, database: currentDatabase, version, @@ -2876,11 +2906,11 @@ export const useQueryStore = defineStore("query", () => { break; } case "countDocuments": { - console.info("[DBX][executeTabSql:mongo-count:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); + queryExecutionLog("info", "mongo-count:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); const total = await api.mongoCountDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.mode, executionId); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCountToQueryResult(total, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-count:done]", { + queryExecutionLog("info", "mongo-count:done", { traceId, collection: mongoCommand.collection, database: currentDatabase, @@ -2894,12 +2924,12 @@ export const useQueryStore = defineStore("query", () => { const safety = evaluateMongoAggregateSafety(mongoCommand, options.mongoSafety); if (!safety.allowed) throw new Error(safety.reason); } - console.info("[DBX][executeTabSql:mongo-aggregate:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); + queryExecutionLog("info", "mongo-aggregate:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); const aggregateMaxRows = normalizeResultPageSize(pageLimit ?? options?.pagination?.limit ?? settingsStore.editorSettings.pageSize); const result = await api.mongoAggregateDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.pipeline, aggregateMaxRows, executionId); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-aggregate:done]", { + queryExecutionLog("info", "mongo-aggregate:done", { traceId, collection: mongoCommand.collection, database: currentDatabase, @@ -2910,11 +2940,11 @@ export const useQueryStore = defineStore("query", () => { break; } case "getIndexes": { - console.info("[DBX][executeTabSql:mongo-indexes:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); + queryExecutionLog("info", "mongo-indexes:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); const indexes = await api.listIndexes(tab.connectionId, currentDatabase, "", mongoCommand.collection); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoIndexesToQueryResult(indexes, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-indexes:done]", { + queryExecutionLog("info", "mongo-indexes:done", { traceId, collection: mongoCommand.collection, database: currentDatabase, @@ -2924,7 +2954,7 @@ export const useQueryStore = defineStore("query", () => { break; } case "collectionStats": { - console.info("[DBX][executeTabSql:mongo-collection-stats:start]", { + queryExecutionLog("info", "mongo-collection-stats:start", { traceId, collection: mongoCommand.collection, metric: mongoCommand.metric, @@ -2933,7 +2963,7 @@ export const useQueryStore = defineStore("query", () => { const stats = await api.mongoCollectionStats(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.scale, executionId); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCollectionStatsToQueryResult(mongoCommand.metric, stats as unknown as Record, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-collection-stats:done]", { + queryExecutionLog("info", "mongo-collection-stats:done", { traceId, collection: mongoCommand.collection, metric: mongoCommand.metric, @@ -2949,7 +2979,7 @@ export const useQueryStore = defineStore("query", () => { const safety = evaluateMongoWriteSafety(mongoCommand, options.mongoSafety); if (!safety.allowed) throw new Error(safety.reason); } - console.info("[DBX][executeTabSql:mongo-find-and-modify:start]", { + queryExecutionLog("info", "mongo-find-and-modify:start", { traceId, kind: mongoCommand.kind, collection: mongoCommand.collection, @@ -2963,7 +2993,7 @@ export const useQueryStore = defineStore("query", () => { : await api.mongoFindOneAndDelete(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.options); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-find-and-modify:done]", { + queryExecutionLog("info", "mongo-find-and-modify:done", { traceId, kind: mongoCommand.kind, collection: mongoCommand.collection, @@ -2984,7 +3014,7 @@ export const useQueryStore = defineStore("query", () => { const safety = evaluateMongoWriteSafety(mongoCommand, options.mongoSafety); if (!safety.allowed) throw new Error(safety.reason); } - console.info("[DBX][executeTabSql:mongo-write:start]", { + queryExecutionLog("info", "mongo-write:start", { traceId, database: currentDatabase, kind: mongoCommand.kind, @@ -3010,7 +3040,7 @@ export const useQueryStore = defineStore("query", () => { const result = await api.mongoDeleteDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.many); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoWriteToQueryResult(result.affected_rows, performance.now() - commandStartedAt)))); } - console.info("[DBX][executeTabSql:mongo-write:done]", { + queryExecutionLog("info", "mongo-write:done", { traceId, database: currentDatabase, kind: mongoCommand.kind, @@ -3023,7 +3053,7 @@ export const useQueryStore = defineStore("query", () => { currentDatabase = mongoCommand.database; allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoUseToQueryResult(currentDatabase, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; - console.info("[DBX][executeTabSql:mongo-use:done]", { + queryExecutionLog("info", "mongo-use:done", { traceId, database: currentDatabase, elapsed: elapsed(), @@ -3039,7 +3069,7 @@ export const useQueryStore = defineStore("query", () => { } } - console.info("[DBX][executeTabSql:mongo:done]", { + queryExecutionLog("info", "mongo:done", { traceId, database: currentDatabase, commandCount: mongoCommands.length, @@ -3127,14 +3157,14 @@ export const useQueryStore = defineStore("query", () => { let executionPromise: Promise; if (tab.autoCommit === false) { if (!tab.txnSessionId) { - console.info("[DBX][executeTabSql:begin-manual-txn:start]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "begin-manual-txn:start", { traceId, elapsed: elapsed() }); tab.txnSessionId = await api.beginManualTransaction(tab.connectionId, executionDatabase, executionSchema); - console.info("[DBX][executeTabSql:begin-manual-txn:done]", { traceId, txnSessionId: tab.txnSessionId, elapsed: elapsed() }); + queryExecutionLog("info", "begin-manual-txn:done", { traceId, txnSessionId: tab.txnSessionId, elapsed: elapsed() }); } - console.info("[DBX][executeTabSql:execute-in-txn:invoke]", { traceId, txnSessionId: tab.txnSessionId, elapsed: elapsed() }); + queryExecutionLog("info", "execute-in-txn:invoke", { traceId, txnSessionId: tab.txnSessionId, elapsed: elapsed() }); executionPromise = api.executeInManualTransaction(tab.txnSessionId, sqlToExecute, executionDatabase, executionSchema, pageLimit); } else { - console.info("[DBX][executeTabSql:execute-multi:start]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "execute-multi:start", { traceId, elapsed: elapsed() }); // Query and data tabs use a tab-scoped pool so repeated executions keep // connection-local state and avoid MySQL pool resets on every refresh. const clientSessionId = tab.mode === "query" || tab.mode === "data" ? tabClientSessionId(tab) : undefined; @@ -3152,7 +3182,7 @@ export const useQueryStore = defineStore("query", () => { timeoutSecs: queryTimeoutSecs, continueOnError: settingsStore.editorSettings.continueOnErrorOnBatch, }; - console.info("[DBX][executeTabSql:execute-multi:invoke]", { + queryExecutionLog("info", "execute-multi:invoke", { traceId, elapsed: elapsed(), executionSchema, @@ -3169,7 +3199,7 @@ export const useQueryStore = defineStore("query", () => { } else if (hiddenPrimaryKeys.length > 0) { queryMetadataSql = queryBaseSql; } - console.info("[DBX][executeTabSql:execute-multi:done]", { + queryExecutionLog("info", "execute-multi:done", { traceId, resultCount: results.length, rowCounts: results.map((result) => result.rows.length), @@ -3262,7 +3292,7 @@ export const useQueryStore = defineStore("query", () => { timeoutSecs: queryTimeoutSecs, }); } - console.info("[DBX][executeTabSql:result:assigned]", { + queryExecutionLog("info", "result:assigned", { traceId, activeResultIndex: current.activeResultIndex, rowCount: current.result?.rows.length ?? 0, @@ -3274,14 +3304,14 @@ export const useQueryStore = defineStore("query", () => { analyzeQueryMetadataInBackground(id, displayedQueryMetadataSql(current, queryMetadataSql), current.result, traceId, elapsed, hiddenPrimaryKeys); } } else { - console.warn("[DBX][executeTabSql:stale-result]", { + queryExecutionLog("warn", "stale-result", { traceId, currentExecutionId: current?.executionId, elapsed: elapsed(), }); } } catch (e: any) { - console.error("[DBX][executeTabSql:error]", { traceId, elapsed: elapsed(), error: e }); + queryExecutionLog("error", "error", { traceId, elapsed: elapsed(), error: e }); // Sync connection state if the error indicates a lost connection useConnectionStore().recordConnectionLostError(tab.connectionId, e); // Handle manual transaction auto-rollback (e.g. deadlock detected by server, @@ -3332,9 +3362,9 @@ export const useQueryStore = defineStore("query", () => { current.isCancelling = false; current.queryExecutionStartedAt = undefined; current.executionId = undefined; - console.info("[DBX][executeTabSql:finish]", { traceId, elapsed: elapsed() }); + queryExecutionLog("info", "finish", { traceId, elapsed: elapsed() }); } else { - console.warn("[DBX][executeTabSql:finish-stale]", { + queryExecutionLog("warn", "finish-stale", { traceId, currentExecutionId: current?.executionId, elapsed: elapsed(), @@ -3561,6 +3591,7 @@ export const useQueryStore = defineStore("query", () => { const executionId = tab.executionId; if (!executionId) return false; tab.isCancelling = true; + const cancellationStartedAt = performance.now(); try { const canceled = await withCancelQueryTimeout(api.cancelQuery(executionId)); if (canceled) { @@ -3588,6 +3619,8 @@ export const useQueryStore = defineStore("query", () => { current.result = toErrorResult(e); } return false; + } finally { + recordQueryCancellationLatency(performance.now() - cancellationStartedAt); } } @@ -3648,9 +3681,20 @@ export const useQueryStore = defineStore("query", () => { } async function trimResultCache() { - const inactive = tabs.value.filter((t) => t.id !== activeTabId.value && (t.result || t.results)).sort((a, b) => (a.resultAccessedAt ?? 0) - (b.resultAccessedAt ?? 0)); - if (inactive.length > MAX_CACHED_RESULTS) { - const toEvict = inactive.slice(0, inactive.length - MAX_CACHED_RESULTS); + const inactive = tabs.value.filter((t) => t.id !== activeTabId.value && (t.result || t.results)); + const evictionIds = new Set( + selectInactiveResultEvictions( + inactive.map((tab) => ({ + id: tab.id, + estimatedBytes: tab.resultEstimatedBytes ?? estimateQueryResultsBytes(tab.result, tab.results), + accessedAt: tab.resultAccessedAt ?? 0, + })), + MAX_CACHED_RESULT_BYTES, + MAX_CACHED_RESULTS, + ), + ); + const toEvict = inactive.filter((tab) => evictionIds.has(tab.id)); + if (toEvict.length > 0) { await Promise.all(toEvict.map((t) => evictCachedResult(t))); } } @@ -3808,13 +3852,14 @@ export const useQueryStore = defineStore("query", () => { return openResultArchiveTab(archive); } - async function reloadEvictedTab(id: string) { + async function reloadEvictedTab(id: string, { reexecuteOnMissing = false }: { reexecuteOnMissing?: boolean } = {}) { const tab = tabs.value.find((t) => t.id === id); if (!tab || !tab.resultEvicted) return; if (tab.resultCacheKey) { const restored = restoreCachedResultPayload(tab, await readTabResultSnapshot(tab.resultCacheKey)); if (restored) return; tab.resultCacheState = "missing"; + if (!reexecuteOnMissing) return; } tab.resultEvicted = false; const sql = tab.lastExecutedSql ?? tab.sql; @@ -4120,6 +4165,7 @@ export const useQueryStore = defineStore("query", () => { importResultArchive, fetchTabResultForExport, buildQueryResultExportRequest, + getResourceLifecycleDiagnostics: () => resourceLifecycleDiagnostics(tabs.value), notifyConnectionMayBeLost, }; }); diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 265584e811..4927238723 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -557,6 +557,7 @@ export interface QueryResultRun { resultTotalRowCountLoading?: boolean; resultSessionId?: string; resultAccessedAt?: number; + resultEstimatedBytes?: number; resultCacheKey?: string; resultCacheState?: "memory" | "disk" | "missing"; resultEvicted?: boolean; @@ -780,6 +781,7 @@ export interface QueryTab { resultTotalRowCountLoading?: boolean; resultSessionId?: string; resultAccessedAt?: number; + resultEstimatedBytes?: number; resultCacheKey?: string; resultCacheState?: "memory" | "disk" | "missing"; pinned?: boolean; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index d837546664..59c559f6dc 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -4,6 +4,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::{Mutex, RwLock}; +#[cfg(feature = "duckdb-bundled")] use tokio::task::JoinHandle; use mysql_async::prelude::Queryable; @@ -29,6 +30,7 @@ use crate::path_utils::expand_tilde; use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv}; use crate::query_cancel::RunningQueries; use crate::storage::{normalize_duckdb_worker_max_processes, Storage, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT}; +use crate::task_supervisor::TaskSupervisor; pub const JDBC_PLUGIN_NOT_INSTALLED: &str = "JDBC plugin is not installed. Install the optional JDBC plugin to use this connection."; @@ -175,7 +177,7 @@ macro_rules! agent_connection_pool_database_type { pub struct AppState { pub connections: Arc>>, - keepalive_tasks: Arc>>>, + task_supervisor: TaskSupervisor, pool_activity: Arc>>, connection_attempts: RwLock>, pub configs: RwLock>, @@ -219,6 +221,7 @@ pub struct PoolActivityTouch { pool_key: String, connections: Arc>>, pool_activity: Arc>>, + task_supervisor: TaskSupervisor, } impl Drop for PoolActivityTouch { @@ -226,10 +229,7 @@ impl Drop for PoolActivityTouch { let pool_key = self.pool_key.clone(); let connections = self.connections.clone(); let pool_activity = self.pool_activity.clone(); - let Ok(handle) = tokio::runtime::Handle::try_current() else { - return; - }; - handle.spawn(async move { + self.task_supervisor.spawn_replace(format!("pool-activity:{pool_key}"), move |_| async move { if !connections.read().await.contains_key(&pool_key) { return; } @@ -569,7 +569,7 @@ impl AppState { ) -> Self { Self { connections: Arc::new(RwLock::new(HashMap::new())), - keepalive_tasks: Arc::new(RwLock::new(HashMap::new())), + task_supervisor: TaskSupervisor::new(), pool_activity: Arc::new(RwLock::new(HashMap::new())), connection_attempts: RwLock::new(HashMap::new()), configs: RwLock::new(HashMap::new()), @@ -926,13 +926,15 @@ impl AppState { let interval = Duration::from_secs(interval_secs.max(1)); let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1)); let connections = self.connections.clone(); - let keepalive_tasks = self.keepalive_tasks.clone(); let pool_activity = self.pool_activity.clone(); let cancel_contexts = self.postgres_cancel_contexts.clone(); let running_queries = self.running_queries.clone(); - let handle = tokio::spawn(async move { + self.task_supervisor.spawn_replace(format!("keepalive:{pool_key}"), move |shutdown| async move { loop { - tokio::time::sleep(interval).await; + tokio::select! { + _ = shutdown.cancelled() => break, + _ = tokio::time::sleep(interval) => {} + } if running_queries.is_pool_active(&key) { continue; @@ -944,7 +946,6 @@ impl AppState { Ok(Ok(())) => {} Ok(Err(err)) => { log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool"); - keepalive_tasks.write().await.remove(&key); pool_activity.write().await.remove(&key); cancel_contexts.write().await.remove(&key); let removed = connections.write().await.remove(&key); @@ -958,7 +959,6 @@ impl AppState { "Connection keepalive timed out for '{key}' after {}s; invalidating pool", timeout.as_secs() ); - keepalive_tasks.write().await.remove(&key); pool_activity.write().await.remove(&key); cancel_contexts.write().await.remove(&key); let removed = connections.write().await.remove(&key); @@ -971,26 +971,15 @@ impl AppState { } } }); - let previous = self.keepalive_tasks.write().await.insert(pool_key.to_string(), handle); - if let Some(previous) = previous { - previous.abort(); - } } async fn stop_keepalive_task(&self, pool_key: &str) { - let task = self.keepalive_tasks.write().await.remove(pool_key); - if let Some(task) = task { - task.abort(); - } + self.task_supervisor.stop(&format!("keepalive:{pool_key}")); } async fn stop_keepalive_tasks(&self, pool_keys: &[String]) { - let mut tasks = self.keepalive_tasks.write().await; - for pool_key in pool_keys { - if let Some(task) = tasks.remove(pool_key) { - task.abort(); - } - } + let keys: Vec = pool_keys.iter().map(|pool_key| format!("keepalive:{pool_key}")).collect(); + self.task_supervisor.stop_many(keys.iter().map(String::as_str)); } pub async fn touch_pool_activity(&self, pool_key: &str) { @@ -1007,9 +996,20 @@ impl AppState { pool_key: pool_key.to_string(), connections: self.connections.clone(), pool_activity: self.pool_activity.clone(), + task_supervisor: self.task_supervisor.clone(), } } + pub async fn shutdown_background_tasks(&self, deadline: Duration) { + self.running_queries.cancel_all(); + self.task_supervisor.shutdown(deadline).await; + } + + #[cfg(test)] + pub fn supervised_task_count(&self) -> usize { + self.task_supervisor.active_count() + } + pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result { self.get_or_create_pool_for_session(connection_id, database, None).await } @@ -2207,16 +2207,18 @@ impl AppState { #[cfg(feature = "duckdb-bundled")] pub fn spawn_duckdb_pool_cleanup(&self, pool_key: String, con: DuckDbHandle) { let connections = self.connections.clone(); - let keepalive_tasks = self.keepalive_tasks.clone(); + let supervisor = self.task_supervisor.clone(); let pool_activity = self.pool_activity.clone(); let postgres_cancel_contexts = self.postgres_cancel_contexts.clone(); - tokio::spawn(async move { + let task_key = format!("duckdb-cleanup:{pool_key}"); + supervisor.clone().spawn_replace(task_key, move |shutdown| async move { while Arc::strong_count(&con) > 2 { - tokio::time::sleep(Duration::from_millis(50)).await; - } - if let Some(handle) = keepalive_tasks.write().await.remove(&pool_key) { - handle.abort(); + tokio::select! { + _ = shutdown.cancelled() => return, + _ = tokio::time::sleep(Duration::from_millis(50)) => {} + } } + supervisor.stop(&format!("keepalive:{pool_key}")); pool_activity.write().await.remove(&pool_key); postgres_cancel_contexts.write().await.remove(&pool_key); let removed = { @@ -2241,20 +2243,28 @@ impl AppState { &self, pool_key: String, con: DuckDbHandle, - task: JoinHandle>, + mut task: JoinHandle>, ) { let connections = self.connections.clone(); - let keepalive_tasks = self.keepalive_tasks.clone(); + let supervisor = self.task_supervisor.clone(); let pool_activity = self.pool_activity.clone(); let postgres_cancel_contexts = self.postgres_cancel_contexts.clone(); - tokio::spawn(async move { - let _ = task.await; + let task_key = format!("duckdb-draining:{pool_key}"); + supervisor.clone().spawn_replace(task_key, move |shutdown| async move { + tokio::select! { + _ = shutdown.cancelled() => { + task.abort(); + return; + } + _ = &mut task => {} + } while Arc::strong_count(&con) > 2 { - tokio::time::sleep(Duration::from_millis(50)).await; - } - if let Some(handle) = keepalive_tasks.write().await.remove(&pool_key) { - handle.abort(); + tokio::select! { + _ = shutdown.cancelled() => return, + _ = tokio::time::sleep(Duration::from_millis(50)) => {} + } } + supervisor.stop(&format!("keepalive:{pool_key}")); pool_activity.write().await.remove(&pool_key); postgres_cancel_contexts.write().await.remove(&pool_key); let removed = { @@ -3288,26 +3298,6 @@ fn native_postgres_url_config(config: &ConnectionConfig) -> Option bool { - let left = expand_tilde(left); - let right = expand_tilde(right); - - if db::duckdb_driver::is_memory_database_path(&left) || db::duckdb_driver::is_memory_database_path(&right) { - return left.trim().eq_ignore_ascii_case(right.trim()); - } - - if let (Ok(left_path), Ok(right_path)) = (std::fs::canonicalize(&left), std::fs::canonicalize(&right)) { - return left_path == right_path; - } - - if cfg!(windows) { - left.eq_ignore_ascii_case(&right) - } else { - left == right - } -} - pub async fn probe_connection_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> Result<(), String> { if !uses_tcp_probe(config, host, port) { return Ok(()); @@ -4671,7 +4661,7 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(20)).await; assert!(state.connections.read().await.contains_key(pool_key)); - assert!(!state.keepalive_tasks.read().await.contains_key(pool_key)); + assert_eq!(state.supervised_task_count(), 0); let _ = std::fs::remove_dir_all(dir); } diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index f153fe0575..99bc3a58ac 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -60,6 +60,7 @@ pub mod storage; pub mod table_export; pub mod table_import; pub mod table_structure_sql; +pub mod task_supervisor; pub mod text_export; pub mod token_usage; pub mod transfer; diff --git a/crates/dbx-core/src/query_cancel.rs b/crates/dbx-core/src/query_cancel.rs index e985b41621..305d62556e 100644 --- a/crates/dbx-core/src/query_cancel.rs +++ b/crates/dbx-core/src/query_cancel.rs @@ -1,7 +1,15 @@ use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use tokio_util::sync::CancellationToken; +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RunningQueryDiagnostics { + pub active_execution_ids: Vec, + pub active_by_connection: HashMap, + pub interrupt_registrations: usize, +} + type InterruptFn = Box; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -25,6 +33,7 @@ pub struct RunningTaskMetadata { pub connection_id: Option, pub database: Option, pub client_session_id: Option, + pub owner_scope: Option, } impl RunningTaskMetadata { @@ -34,12 +43,25 @@ impl RunningTaskMetadata { client_session_id: Option, ) -> Self { let kind = task_kind_from_client_session_id(client_session_id.as_deref()); - Self { kind, connection_id: Some(connection_id.into()), database: Some(database.into()), client_session_id } + let owner_scope = client_session_id.as_deref().and_then(owner_scope_from_client_session_id); + Self { + kind, + connection_id: Some(connection_id.into()), + database: Some(database.into()), + client_session_id, + owner_scope, + } + } + + pub fn with_owner_scope(mut self, owner_scope: impl Into) -> Self { + self.owner_scope = Some(owner_scope.into()); + self } } #[derive(Clone)] struct RunningTask { + registration_id: u64, token: CancellationToken, metadata: RunningTaskMetadata, pool_key: Option, @@ -49,6 +71,7 @@ struct RunningTask { pub struct RunningQueries { inner: Arc>>, interrupts: Arc>>, + next_registration_id: Arc, } impl RunningQueries { @@ -58,12 +81,19 @@ impl RunningQueries { pub fn register_task(&self, execution_id: String, metadata: RunningTaskMetadata) -> RegisteredQuery { let token = CancellationToken::new(); - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(execution_id.clone(), RunningTask { token: token.clone(), metadata, pool_key: None }); + let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed) + 1; + let previous = self.inner.lock().unwrap_or_else(|e| e.into_inner()).insert( + execution_id.clone(), + RunningTask { registration_id, token: token.clone(), metadata, pool_key: None }, + ); + if let Some(previous) = previous { + previous.token.cancel(); + if let Some(interrupt) = self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(&execution_id) { + interrupt(); + } + } - RegisteredQuery { execution_id, token, running_queries: self.clone() } + RegisteredQuery { execution_id, registration_id, token, running_queries: self.clone() } } pub fn register_interrupt(&self, execution_id: &str, interrupt: impl Fn() + Send + 'static) { @@ -86,6 +116,52 @@ impl RunningQueries { } } + pub fn cancel_connection(&self, connection_id: &str) -> usize { + self.cancel_matching(|task| task.metadata.connection_id.as_deref() == Some(connection_id)) + } + + pub fn cancel_client_session(&self, client_session_id: &str) -> usize { + self.cancel_matching(|task| task.metadata.client_session_id.as_deref() == Some(client_session_id)) + } + + pub fn cancel_owner_scope(&self, owner_scope: &str) -> usize { + self.cancel_matching(|task| task.metadata.owner_scope.as_deref() == Some(owner_scope)) + } + + pub fn cancel_all(&self) -> usize { + self.cancel_matching(|_| true) + } + + pub fn diagnostics(&self) -> RunningQueryDiagnostics { + let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + let mut active_execution_ids = inner.keys().cloned().collect::>(); + active_execution_ids.sort(); + let mut active_by_connection = HashMap::new(); + for task in inner.values() { + if let Some(connection_id) = &task.metadata.connection_id { + *active_by_connection.entry(connection_id.clone()).or_insert(0) += 1; + } + } + drop(inner); + RunningQueryDiagnostics { + active_execution_ids, + active_by_connection, + interrupt_registrations: self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).len(), + } + } + + fn cancel_matching(&self, predicate: impl Fn(&RunningTask) -> bool) -> usize { + let execution_ids: Vec = self + .inner + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .filter(|(_, task)| predicate(task)) + .map(|(execution_id, _)| execution_id.clone()) + .collect(); + execution_ids.iter().filter(|execution_id| self.cancel(execution_id)).count() + } + pub fn set_pool_key(&self, execution_id: &str, pool_key: impl Into) { if let Some(task) = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get_mut(execution_id) { task.pool_key = Some(pool_key.into()); @@ -109,9 +185,27 @@ impl RunningQueries { self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.metadata.kind) } - fn remove(&self, execution_id: &str) { - self.inner.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); - self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); + #[cfg(test)] + pub fn registration_counts(&self) -> (usize, usize) { + ( + self.inner.lock().unwrap_or_else(|e| e.into_inner()).len(), + self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).len(), + ) + } + + fn remove(&self, execution_id: &str, registration_id: u64) { + let removed = { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.get(execution_id).is_some_and(|task| task.registration_id == registration_id) { + inner.remove(execution_id); + true + } else { + false + } + }; + if removed { + self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); + } } } @@ -131,8 +225,14 @@ fn task_kind_from_client_session_id(client_session_id: Option<&str>) -> RunningT } } +fn owner_scope_from_client_session_id(client_session_id: &str) -> Option { + let owner = client_session_id.trim().split(':').next()?.trim(); + (!owner.is_empty()).then(|| owner.to_string()) +} + pub struct RegisteredQuery { execution_id: String, + registration_id: u64, token: CancellationToken, running_queries: RunningQueries, } @@ -145,7 +245,7 @@ impl RegisteredQuery { impl Drop for RegisteredQuery { fn drop(&mut self) { - self.running_queries.remove(&self.execution_id); + self.running_queries.remove(&self.execution_id, self.registration_id); } } @@ -187,6 +287,61 @@ mod tests { assert!(!running.has("exec-1")); } + #[test] + fn stale_registration_drop_does_not_remove_replacement() { + let running = RunningQueries::default(); + let first = running.register("exec-1".to_string()); + let second = running.register("exec-1".to_string()); + + assert!(first.token().is_cancelled()); + drop(first); + assert!(running.has("exec-1")); + drop(second); + assert!(!running.has("exec-1")); + } + + #[test] + fn scoped_cancellation_only_cancels_matching_tasks() { + let running = RunningQueries::default(); + let first = running.register_task( + "exec-1".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())), + ); + let second = running.register_task( + "exec-2".to_string(), + RunningTaskMetadata::query("conn-2", "main", Some("tab-2".to_string())), + ); + + assert_eq!(running.cancel_connection("conn-1"), 1); + assert!(first.token().is_cancelled()); + assert!(!second.token().is_cancelled()); + assert_eq!(running.cancel_client_session("tab-2"), 1); + assert!(second.token().is_cancelled()); + } + + #[test] + fn query_metadata_derives_stable_owner_scope() { + let metadata = RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())); + assert_eq!(metadata.owner_scope.as_deref(), Some("tab-1")); + } + + #[test] + fn owner_scope_cancellation_matches_related_task_kinds() { + let running = RunningQueries::default(); + let query = running.register_task( + "exec-query".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())), + ); + let export = running.register_task( + "exec-export".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())), + ); + + assert_eq!(running.cancel_owner_scope("tab-1"), 2); + assert!(query.token().is_cancelled()); + assert!(export.token().is_cancelled()); + } + #[test] fn register_task_tracks_kind_and_pool_activity() { let running = RunningQueries::default(); @@ -202,4 +357,48 @@ mod tests { drop(registered); assert!(!running.is_pool_active("conn-1:main:session:tab-1_export")); } + + #[test] + fn unwind_removes_task_and_interrupt_registrations() { + let running = RunningQueries::default(); + let unwind_running = running.clone(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || { + let _registered = unwind_running.register("exec-unwind".to_string()); + unwind_running.register_interrupt("exec-unwind", || {}); + panic!("simulate driver unwind"); + })); + + assert!(result.is_err()); + assert_eq!(running.registration_counts(), (0, 0)); + } + + #[test] + fn completion_after_cancel_removes_all_registrations() { + let running = RunningQueries::default(); + let registered = running.register("exec-cancel".to_string()); + running.register_interrupt("exec-cancel", || {}); + + assert!(running.cancel("exec-cancel")); + drop(registered); + + assert_eq!(running.registration_counts(), (0, 0)); + } + + #[test] + fn diagnostics_expose_ids_and_scoped_counts_without_query_text() { + let running = RunningQueries::default(); + let _first = running.register_task( + "exec-2".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())), + ); + let _second = running.register_task( + "exec-1".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-2".to_string())), + ); + + let diagnostics = running.diagnostics(); + + assert_eq!(diagnostics.active_execution_ids, vec!["exec-1", "exec-2"]); + assert_eq!(diagnostics.active_by_connection.get("conn-1"), Some(&2)); + } } diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 1b575e84ea..e7e25fdb75 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -1,5 +1,6 @@ use std::collections::{HashMap, HashSet}; use std::path::Path; +use std::time::{SystemTime, UNIX_EPOCH}; use log::warn; use rusqlite::{params, params_from_iter, Connection, DatabaseName, OpenFlags, OptionalExtension, ToSql}; @@ -102,6 +103,7 @@ pub struct Storage { } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct TabRuntimeCacheEntry { pub key: String, pub payload: Vec, @@ -109,6 +111,32 @@ pub struct TabRuntimeCacheEntry { pub column_count: i64, pub byte_size: i64, pub updated_at: String, + pub created_at: i64, + pub last_accessed_at: i64, + pub owner_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TabRuntimeCacheMetadata { + pub key: String, + pub row_count: i64, + pub column_count: i64, + pub byte_size: i64, + pub updated_at: String, + pub created_at: i64, + pub last_accessed_at: i64, + pub owner_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TabRuntimeCachePruneResult { + pub deleted_entries: usize, + pub deleted_bytes: i64, + pub orphan_deletions: usize, + pub remaining_entries: usize, + pub remaining_bytes: i64, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -260,7 +288,10 @@ const SCHEMA_STATEMENTS: &[&str] = &[ row_count INTEGER NOT NULL DEFAULT 0, column_count INTEGER NOT NULL DEFAULT 0, byte_size INTEGER NOT NULL DEFAULT 0, - updated_at TEXT NOT NULL + updated_at TEXT NOT NULL, + created_at INTEGER NOT NULL DEFAULT 0, + last_accessed_at INTEGER NOT NULL DEFAULT 0, + owner_id TEXT )", "CREATE TABLE IF NOT EXISTS mq_token_records ( id TEXT PRIMARY KEY, @@ -319,6 +350,7 @@ impl Storage { } ensure_history_columns_sync(conn)?; ensure_saved_sql_columns_sync(conn)?; + ensure_tab_runtime_cache_columns_sync(conn)?; Ok(()) }) } @@ -432,6 +464,26 @@ fn ensure_saved_sql_columns_sync(conn: &Connection) -> Result<(), String> { Ok(()) } +fn ensure_tab_runtime_cache_columns_sync(conn: &Connection) -> Result<(), String> { + const COLUMNS: &[(&str, &str)] = &[ + ("created_at", "INTEGER NOT NULL DEFAULT 0"), + ("last_accessed_at", "INTEGER NOT NULL DEFAULT 0"), + ("owner_id", "TEXT"), + ]; + ensure_table_columns(conn, "tab_runtime_cache", COLUMNS)?; + let now = unix_timestamp_millis(); + // Legacy rows must receive a grace period instead of being treated as ancient crash leftovers. + conn.execute("UPDATE tab_runtime_cache SET created_at = ?1 WHERE created_at = 0", [now]) + .map_err(|e| e.to_string())?; + conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = created_at WHERE last_accessed_at = 0", []) + .map_err(|e| e.to_string())?; + Ok(()) +} + +fn unix_timestamp_millis() -> i64 { + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis().min(i64::MAX as u128) as i64 +} + fn ensure_table_columns(conn: &Connection, table_name: &str, columns: &[(&str, &str)]) -> Result<(), String> { let mut stmt = conn.prepare(&format!("SELECT name FROM pragma_table_info('{table_name}')")).map_err(|e| e.to_string())?; @@ -2082,18 +2134,21 @@ impl Storage { payload: Vec, row_count: i64, column_count: i64, + owner_id: Option, ) -> Result<(), String> { let key = key.to_string(); let byte_size = payload.len() as i64; + let now = unix_timestamp_millis(); self.with_conn(move |conn| { conn.execute( "INSERT INTO tab_runtime_cache \ - (cache_key, payload, row_count, column_count, byte_size, updated_at) \ - VALUES (?1, ?2, ?3, ?4, ?5, datetime('now')) \ + (cache_key, payload, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id) \ + VALUES (?1, ?2, ?3, ?4, ?5, datetime('now'), ?6, ?6, ?7) \ ON CONFLICT(cache_key) DO UPDATE SET \ payload = excluded.payload, row_count = excluded.row_count, column_count = excluded.column_count, \ - byte_size = excluded.byte_size, updated_at = excluded.updated_at", - params![key, payload, row_count, column_count, byte_size], + byte_size = excluded.byte_size, updated_at = excluded.updated_at, \ + last_accessed_at = excluded.last_accessed_at, owner_id = excluded.owner_id", + params![key, payload, row_count, column_count, byte_size, now, owner_id], ) .map(|_| ()) .map_err(|e| e.to_string()) @@ -2103,11 +2158,13 @@ impl Storage { pub async fn load_tab_runtime_cache(&self, key: &str) -> Result, String> { let key = key.to_string(); + let now = unix_timestamp_millis(); self.with_conn(move |conn| { - conn.query_row( - "SELECT cache_key, payload, row_count, column_count, byte_size, updated_at \ + let entry = conn + .query_row( + "SELECT cache_key, payload, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id \ FROM tab_runtime_cache WHERE cache_key = ?1", - [key], + [&key], |row| { Ok(TabRuntimeCacheEntry { key: row.get(0)?, @@ -2116,11 +2173,134 @@ impl Storage { column_count: row.get(3)?, byte_size: row.get(4)?, updated_at: row.get(5)?, + created_at: row.get(6)?, + last_accessed_at: now, + owner_id: row.get(8)?, }) }, ) .optional() - .map_err(|e| e.to_string()) + .map_err(|e| e.to_string())?; + if entry.is_some() { + conn.execute( + "UPDATE tab_runtime_cache SET last_accessed_at = ?2 WHERE cache_key = ?1", + params![key, now], + ) + .map_err(|e| e.to_string())?; + } + Ok(entry) + }) + .await + } + + pub async fn list_tab_runtime_cache_metadata(&self) -> Result, String> { + self.with_conn(move |conn| { + let mut statement = conn + .prepare( + "SELECT cache_key, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id \ + FROM tab_runtime_cache ORDER BY last_accessed_at ASC, cache_key ASC", + ) + .map_err(|e| e.to_string())?; + let metadata = statement + .query_map([], |row| { + Ok(TabRuntimeCacheMetadata { + key: row.get(0)?, + row_count: row.get(1)?, + column_count: row.get(2)?, + byte_size: row.get(3)?, + updated_at: row.get(4)?, + created_at: row.get(5)?, + last_accessed_at: row.get(6)?, + owner_id: row.get(7)?, + }) + }) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; + Ok(metadata) + }) + .await + } + + pub async fn prune_tab_runtime_cache( + &self, + live_keys: Vec, + max_bytes: i64, + orphan_grace_ms: i64, + max_age_ms: Option, + ) -> Result { + let now = unix_timestamp_millis(); + self.with_conn(move |conn| { + let live_keys: HashSet = live_keys.into_iter().collect(); + let mut statement = conn + .prepare( + "SELECT cache_key, byte_size, created_at, last_accessed_at \ + FROM tab_runtime_cache ORDER BY last_accessed_at ASC, cache_key ASC", + ) + .map_err(|e| e.to_string())?; + let entries = statement + .query_map([], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?, row.get::<_, i64>(2)?, row.get::<_, i64>(3)?)) + }) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; + drop(statement); + + let mut total_bytes = entries.iter().map(|(_, bytes, _, _)| *bytes).sum::(); + let mut deleted = HashSet::new(); + let mut orphan_deletions = 0usize; + for (key, bytes, created_at, last_accessed_at) in &entries { + if live_keys.contains(key) { + continue; + } + let orphan_expired = now.saturating_sub(*created_at) >= orphan_grace_ms.max(0); + let age_expired = + max_age_ms.is_some_and(|max_age| now.saturating_sub(*last_accessed_at) >= max_age.max(0)); + if orphan_expired || age_expired { + deleted.insert(key.clone()); + total_bytes = total_bytes.saturating_sub(*bytes); + if orphan_expired { + orphan_deletions += 1; + } + } + } + + for (key, bytes, _, _) in &entries { + if total_bytes <= max_bytes.max(0) { + break; + } + if live_keys.contains(key) || deleted.contains(key) { + continue; + } + deleted.insert(key.clone()); + total_bytes = total_bytes.saturating_sub(*bytes); + } + + let deleted_bytes = + entries.iter().filter(|(key, _, _, _)| deleted.contains(key)).map(|(_, bytes, _, _)| *bytes).sum(); + let transaction = conn.transaction().map_err(|e| e.to_string())?; + for key in &deleted { + transaction + .execute("DELETE FROM tab_runtime_cache WHERE cache_key = ?1", [key]) + .map_err(|e| e.to_string())?; + } + transaction.commit().map_err(|e| e.to_string())?; + Ok(TabRuntimeCachePruneResult { + deleted_entries: deleted.len(), + deleted_bytes, + orphan_deletions, + remaining_entries: entries.len().saturating_sub(deleted.len()), + remaining_bytes: total_bytes, + }) + }) + .await + } + + pub async fn delete_tab_runtime_cache_owner(&self, owner_id: &str) -> Result { + let owner_id = owner_id.to_string(); + self.with_conn(move |conn| { + conn.execute("DELETE FROM tab_runtime_cache WHERE owner_id = ?1", [owner_id]).map_err(|e| e.to_string()) }) .await } @@ -2535,6 +2715,7 @@ mod tests { ConnectionConfig, DatabaseConnectionInfo, DatabaseType, SshTunnelConfig, TransportLayerConfig, }; use crate::saved_sql::SavedSqlFile; + use rusqlite::Connection; use std::time::{SystemTime, UNIX_EPOCH}; fn temp_db_path(name: &str) -> std::path::PathBuf { @@ -3256,7 +3437,10 @@ mod tests { let path = temp_db_path("tab-runtime-cache"); let storage = Storage::open(&path).await.unwrap(); - storage.save_tab_runtime_cache("tab:1:result", vec![1, 2, 3, 4], 10, 3).await.unwrap(); + storage + .save_tab_runtime_cache("tab:1:result", vec![1, 2, 3, 4], 10, 3, Some("connection-1".to_string())) + .await + .unwrap(); let entry = storage.load_tab_runtime_cache("tab:1:result").await.unwrap().unwrap(); assert_eq!(entry.key, "tab:1:result"); @@ -3264,11 +3448,86 @@ mod tests { assert_eq!(entry.row_count, 10); assert_eq!(entry.column_count, 3); assert_eq!(entry.byte_size, 4); + assert_eq!(entry.owner_id.as_deref(), Some("connection-1")); + assert!(entry.created_at > 0); + assert!(entry.last_accessed_at >= entry.created_at); storage.delete_tab_runtime_cache("tab:1:result").await.unwrap(); assert_eq!(storage.load_tab_runtime_cache("tab:1:result").await.unwrap(), None); } + #[tokio::test] + async fn tab_runtime_cache_pruning_retains_live_entries_and_enforces_byte_budget() { + let path = temp_db_path("tab-runtime-cache-prune"); + let storage = Storage::open(&path).await.unwrap(); + for (key, size) in [("live", 6usize), ("old", 5), ("new", 4)] { + storage.save_tab_runtime_cache(key, vec![1; size], 1, 1, Some("connection-1".to_string())).await.unwrap(); + } + storage + .with_conn(|conn| { + conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = 1 WHERE cache_key = 'old'", []) + .map_err(|e| e.to_string())?; + conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = 2 WHERE cache_key = 'new'", []) + .map_err(|e| e.to_string())?; + Ok(()) + }) + .await + .unwrap(); + + let result = storage.prune_tab_runtime_cache(vec!["live".to_string()], 10, i64::MAX, None).await.unwrap(); + + assert_eq!(result.deleted_entries, 1); + assert!(storage.load_tab_runtime_cache("live").await.unwrap().is_some()); + assert!(storage.load_tab_runtime_cache("old").await.unwrap().is_none()); + assert!(storage.load_tab_runtime_cache("new").await.unwrap().is_some()); + } + + #[tokio::test] + async fn tab_runtime_cache_pruning_respects_orphan_grace_period() { + let path = temp_db_path("tab-runtime-cache-orphan-grace"); + let storage = Storage::open(&path).await.unwrap(); + storage.save_tab_runtime_cache("fresh", vec![1], 1, 1, None).await.unwrap(); + storage.save_tab_runtime_cache("crash-leftover", vec![2], 1, 1, None).await.unwrap(); + storage + .with_conn(|conn| { + conn.execute("UPDATE tab_runtime_cache SET created_at = 1 WHERE cache_key = 'crash-leftover'", []) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await + .unwrap(); + + let result = storage.prune_tab_runtime_cache(Vec::new(), 1024, 60_000, None).await.unwrap(); + + assert_eq!(result.orphan_deletions, 1); + assert!(storage.load_tab_runtime_cache("fresh").await.unwrap().is_some()); + assert!(storage.load_tab_runtime_cache("crash-leftover").await.unwrap().is_none()); + } + + #[tokio::test] + async fn tab_runtime_cache_migrates_legacy_schema_without_dropping_entries() { + let path = temp_db_path("tab-runtime-cache-legacy-schema"); + { + let connection = Connection::open(&path).unwrap(); + connection + .execute( + "CREATE TABLE tab_runtime_cache (cache_key TEXT PRIMARY KEY, payload BLOB NOT NULL, row_count INTEGER NOT NULL DEFAULT 0, column_count INTEGER NOT NULL DEFAULT 0, byte_size INTEGER NOT NULL DEFAULT 0, updated_at TEXT NOT NULL)", + [], + ) + .unwrap(); + connection + .execute("INSERT INTO tab_runtime_cache VALUES ('legacy', X'0102', 1, 1, 2, '2026-01-01')", []) + .unwrap(); + } + + let storage = Storage::open(&path).await.unwrap(); + let entry = storage.load_tab_runtime_cache("legacy").await.unwrap().unwrap(); + + assert_eq!(entry.payload, vec![1, 2]); + assert!(entry.created_at > 0); + assert!(entry.last_accessed_at >= entry.created_at); + } + #[tokio::test] async fn saved_sql_summary_omits_sql_text_and_loads_file_on_demand() { let path = temp_db_path("saved-sql-summary"); diff --git a/crates/dbx-core/src/task_supervisor.rs b/crates/dbx-core/src/task_supervisor.rs new file mode 100644 index 0000000000..2b1889f710 --- /dev/null +++ b/crates/dbx-core/src/task_supervisor.rs @@ -0,0 +1,111 @@ +use std::collections::HashMap; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +pub struct TaskSupervisor { + accepting: Arc, + cancellation: CancellationToken, + tasks: Arc>>>, +} + +impl TaskSupervisor { + pub fn new() -> Self { + Self { + accepting: Arc::new(AtomicBool::new(true)), + cancellation: CancellationToken::new(), + tasks: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub fn spawn_replace(&self, key: impl Into, task: F) -> bool + where + F: FnOnce(CancellationToken) -> Fut, + Fut: Future + Send + 'static, + { + if !self.accepting.load(Ordering::Acquire) { + return false; + } + let handle = tokio::spawn(task(self.cancellation.child_token())); + if let Some(previous) = self.tasks.lock().unwrap_or_else(|error| error.into_inner()).insert(key.into(), handle) + { + previous.abort(); + } + true + } + + pub fn stop(&self, key: &str) { + if let Some(task) = self.tasks.lock().unwrap_or_else(|error| error.into_inner()).remove(key) { + task.abort(); + } + } + + pub fn stop_many<'a>(&self, keys: impl IntoIterator) { + for key in keys { + self.stop(key); + } + } + + pub fn active_count(&self) -> usize { + self.tasks.lock().unwrap_or_else(|error| error.into_inner()).len() + } + + pub async fn shutdown(&self, deadline: Duration) { + self.accepting.store(false, Ordering::Release); + self.cancellation.cancel(); + let mut handles: Vec> = + self.tasks.lock().unwrap_or_else(|error| error.into_inner()).drain().map(|(_, handle)| handle).collect(); + let wait = async { + for handle in &mut handles { + let _ = handle.await; + } + }; + if tokio::time::timeout(deadline, wait).await.is_err() { + for handle in handles { + handle.abort(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::TaskSupervisor; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + #[tokio::test] + async fn replacement_aborts_previous_task() { + let supervisor = TaskSupervisor::new(); + let cancelled = Arc::new(AtomicBool::new(false)); + let flag = cancelled.clone(); + supervisor.spawn_replace("task", move |token| async move { + token.cancelled().await; + flag.store(true, Ordering::SeqCst); + }); + supervisor.spawn_replace("task", |_| async {}); + tokio::task::yield_now().await; + assert_eq!(supervisor.active_count(), 1); + assert!(!cancelled.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn shutdown_cancels_tasks_and_rejects_new_work() { + let supervisor = TaskSupervisor::new(); + let cancelled = Arc::new(AtomicBool::new(false)); + let flag = cancelled.clone(); + supervisor.spawn_replace("task", move |token| async move { + token.cancelled().await; + flag.store(true, Ordering::SeqCst); + }); + supervisor.shutdown(Duration::from_secs(1)).await; + assert!(cancelled.load(Ordering::SeqCst)); + assert_eq!(supervisor.active_count(), 0); + assert!(!supervisor.spawn_replace("late", |_| async {})); + } +} diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 8f5a9b3978..be78653838 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -314,6 +314,9 @@ async fn main() { .get(routes::tab_runtime_cache::load_tab_runtime_cache) .delete(routes::tab_runtime_cache::delete_tab_runtime_cache), ) + .route("/tab-runtime-cache/metadata", get(routes::tab_runtime_cache::list_tab_runtime_cache_metadata)) + .route("/tab-runtime-cache/prune", post(routes::tab_runtime_cache::prune_tab_runtime_cache)) + .route("/tab-runtime-cache/owner", delete(routes::tab_runtime_cache::delete_tab_runtime_cache_owner)) // Query .route("/query/execute", post(routes::query::execute_query)) .route("/query/execute-multi", post(routes::query::execute_multi)) @@ -632,5 +635,14 @@ async fn main() { } let listener = tokio::net::TcpListener::bind(addr).await.expect("Failed to bind address"); - axum::serve(listener, app).await.expect("Server error"); + let shutdown_state = web_state.app.clone(); + axum::serve(listener, app) + .with_graceful_shutdown(async { + if let Err(error) = tokio::signal::ctrl_c().await { + tracing::warn!("Failed to listen for shutdown signal: {error}"); + } + }) + .await + .expect("Server error"); + shutdown_state.shutdown_background_tasks(std::time::Duration::from_secs(3)).await; } diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index c6b05ea410..9866f152a2 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -179,6 +179,7 @@ pub async fn disconnect_db( if !should_disconnect { return Ok(Json(())); } + app.running_queries.cancel_connection(&body.connection_id); app.remove_connection_pools_detached(&body.connection_id).await; app.nacos_registry.drop_connection(&body.connection_id).await; #[cfg(feature = "mq-admin")] diff --git a/crates/dbx-web/src/routes/tab_runtime_cache.rs b/crates/dbx-web/src/routes/tab_runtime_cache.rs index 49ce8be8ae..5622827b85 100644 --- a/crates/dbx-web/src/routes/tab_runtime_cache.rs +++ b/crates/dbx-web/src/routes/tab_runtime_cache.rs @@ -15,6 +15,7 @@ pub struct SaveTabRuntimeCacheRequest { pub payload_base64: String, pub row_count: i64, pub column_count: i64, + pub owner_id: Option, } #[derive(Deserialize)] @@ -31,6 +32,9 @@ pub struct LoadTabRuntimeCacheResponse { pub column_count: i64, pub byte_size: i64, pub updated_at: String, + pub created_at: i64, + pub last_accessed_at: i64, + pub owner_id: Option, } pub async fn save_tab_runtime_cache( @@ -41,7 +45,7 @@ pub async fn save_tab_runtime_cache( state .app .storage - .save_tab_runtime_cache(&req.key, payload, req.row_count, req.column_count) + .save_tab_runtime_cache(&req.key, payload, req.row_count, req.column_count, req.owner_id) .await .map_err(AppError::internal)?; Ok(Json(())) @@ -59,9 +63,53 @@ pub async fn load_tab_runtime_cache( column_count: entry.column_count, byte_size: entry.byte_size, updated_at: entry.updated_at, + created_at: entry.created_at, + last_accessed_at: entry.last_accessed_at, + owner_id: entry.owner_id, }))) } +pub async fn list_tab_runtime_cache_metadata( + State(state): State>, +) -> Result>, AppError> { + Ok(Json(state.app.storage.list_tab_runtime_cache_metadata().await.map_err(AppError::internal)?)) +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct PruneTabRuntimeCacheRequest { + pub live_keys: Vec, + pub max_bytes: i64, + pub orphan_grace_ms: i64, + pub max_age_ms: Option, +} + +pub async fn prune_tab_runtime_cache( + State(state): State>, + Json(request): Json, +) -> Result, AppError> { + Ok(Json( + state + .app + .storage + .prune_tab_runtime_cache(request.live_keys, request.max_bytes, request.orphan_grace_ms, request.max_age_ms) + .await + .map_err(AppError::internal)?, + )) +} + +#[derive(Deserialize)] +pub struct TabRuntimeCacheOwnerQuery { + pub owner_id: String, +} + +pub async fn delete_tab_runtime_cache_owner( + State(state): State>, + Query(query): Query, +) -> Result, AppError> { + Ok(Json(state.app.storage.delete_tab_runtime_cache_owner(&query.owner_id).await.map_err(AppError::internal)?)) +} + pub async fn delete_tab_runtime_cache( State(state): State>, Query(query): Query, diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 8693c1a5eb..35327a659c 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -1180,7 +1180,7 @@ test("kept result runs evict inactive payloads without losing switch or archive const tab = store.tabs.find((item) => item.id === tabId); assert.ok(tab?.resultRuns?.[0]); assert.ok(tab.resultRuns[1]); - await waitFor(() => tab.resultRuns?.[0]?.result === undefined && tab.resultRuns?.[0]?.resultCacheState === "disk"); + await waitFor(() => tab.resultRuns?.[0]?.result === undefined && tab.resultRuns?.[0]?.resultCacheState === "disk", 3_000); assert.deepEqual(tab.result?.columns, ["run_2"]); assert.deepEqual(tab.resultRuns[1]?.result?.columns, ["run_2"]); @@ -1256,7 +1256,7 @@ test("removing the active result run restores a disk-backed adjacent run", async const tab = store.tabs.find((item) => item.id === tabId); assert.ok(tab?.resultRuns?.[0]); assert.ok(tab.resultRuns[1]); - await waitFor(() => tab.resultRuns?.[0]?.result === undefined && tab.resultRuns?.[0]?.resultCacheState === "disk"); + await waitFor(() => tab.resultRuns?.[0]?.result === undefined && tab.resultRuns?.[0]?.resultCacheState === "disk", 3_000); assert.equal(await store.removeResultRun(tabId, tab.resultRuns[1].id), true); @@ -2031,7 +2031,7 @@ test("evicting cached tab results releases multi-result payloads and sessions", await store.executeTabSql(tabId, `select ${i + 1}; select ${i + 1} as detail`); } - await waitFor(() => store.tabs.find((tab) => tab.id === tabIds[0])?.resultEvicted === true); + await waitFor(() => store.tabs.find((tab) => tab.id === tabIds[0])?.resultEvicted === true, 3_000); const evicted = store.tabs.find((tab) => tab.id === tabIds[0]); assert.equal(executeCount, 7); assert.equal(evicted?.result, undefined); @@ -2107,7 +2107,7 @@ test("result cache eviction keeps recently accessed inactive tabs", async () => tabIds.push(tabId); await store.executeTabSql(tabId, "select 7"); - await waitFor(() => store.tabs.find((tab) => tab.id === tabIds[1])?.resultEvicted === true); + await waitFor(() => store.tabs.find((tab) => tab.id === tabIds[1])?.resultEvicted === true, 3_000); const recentlyViewed = store.tabs.find((tab) => tab.id === tabIds[0]); const leastRecentlyUsed = store.tabs.find((tab) => tab.id === tabIds[1]); assert.ok(recentlyViewed?.result); diff --git a/packages/app-tests/sidebarContextMenuHost.test.ts b/packages/app-tests/sidebarContextMenuHost.test.ts index a828607b75..9b6bdd2739 100644 --- a/packages/app-tests/sidebarContextMenuHost.test.ts +++ b/packages/app-tests/sidebarContextMenuHost.test.ts @@ -31,7 +31,7 @@ test("tree host owns sidebar data-open generations", () => { const connectionTree = readFileSync("apps/desktop/src/components/sidebar/ConnectionTree.vue", "utf8"); assert.doesNotMatch(treeItem, /runSidebarDataOpenImmediately/); - assert.match(treeItem, /emit\("open-data", node, true, openData\)/); + assert.match(treeItem, /emit\("open-data", node, true, "default", openData\)/); assert.match(connectionTree, /function openSidebarData/); assert.match(connectionTree, /runSidebarDataOpenImmediately/); assert.match(connectionTree, /createSidebarActionTarget\(node\)/); diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index c7bd1e77df..8869b96ed2 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -1282,6 +1282,7 @@ pub async fn disconnect_db( if !should_disconnect { return Ok(()); } + state.running_queries.cancel_connection(&connection_id); state.remove_connection_pools_detached(&connection_id).await; drop_nacos_adapters_for_connection_ids(state.inner(), std::slice::from_ref(&connection_id)).await; drop_mq_adapters_for_connection_ids(state.inner(), std::slice::from_ref(&connection_id)).await; diff --git a/src-tauri/src/commands/tab_runtime_cache.rs b/src-tauri/src/commands/tab_runtime_cache.rs index 96fbb85cd0..f07f6d4836 100644 --- a/src-tauri/src/commands/tab_runtime_cache.rs +++ b/src-tauri/src/commands/tab_runtime_cache.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; use dbx_core::connection::AppState; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use tauri::State; #[derive(Serialize)] @@ -14,6 +14,9 @@ pub struct LoadTabRuntimeCacheResponse { pub column_count: i64, pub byte_size: i64, pub updated_at: String, + pub created_at: i64, + pub last_accessed_at: i64, + pub owner_id: Option, } #[tauri::command] @@ -23,9 +26,10 @@ pub async fn save_tab_runtime_cache( payload_base64: String, row_count: i64, column_count: i64, + owner_id: Option, ) -> Result<(), String> { let payload = BASE64.decode(payload_base64).map_err(|e| e.to_string())?; - state.storage.save_tab_runtime_cache(&key, payload, row_count, column_count).await + state.storage.save_tab_runtime_cache(&key, payload, row_count, column_count, owner_id).await } #[tauri::command] @@ -41,9 +45,47 @@ pub async fn load_tab_runtime_cache( column_count: entry.column_count, byte_size: entry.byte_size, updated_at: entry.updated_at, + created_at: entry.created_at, + last_accessed_at: entry.last_accessed_at, + owner_id: entry.owner_id, })) } +#[tauri::command] +pub async fn list_tab_runtime_cache_metadata( + state: State<'_, Arc>, +) -> Result, String> { + state.storage.list_tab_runtime_cache_metadata().await +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct PruneTabRuntimeCacheRequest { + pub live_keys: Vec, + pub max_bytes: i64, + pub orphan_grace_ms: i64, + pub max_age_ms: Option, +} + +#[tauri::command] +pub async fn prune_tab_runtime_cache( + state: State<'_, Arc>, + request: PruneTabRuntimeCacheRequest, +) -> Result { + state + .storage + .prune_tab_runtime_cache(request.live_keys, request.max_bytes, request.orphan_grace_ms, request.max_age_ms) + .await +} + +#[tauri::command] +pub async fn delete_tab_runtime_cache_owner( + state: State<'_, Arc>, + owner_id: String, +) -> Result { + state.storage.delete_tab_runtime_cache_owner(&owner_id).await +} + #[tauri::command] pub async fn delete_tab_runtime_cache(state: State<'_, Arc>, key: String) -> Result<(), String> { state.storage.delete_tab_runtime_cache(&key).await diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 972f8f7c85..6325739a47 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -10,7 +10,7 @@ use commands::connection::AppState; use dbx_core::storage::{maybe_import_user_data_db, DesktopIconTheme, DesktopSettings, Storage}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; #[cfg(target_os = "macos")] use tauri::menu::Menu; #[cfg(target_os = "macos")] @@ -1129,6 +1129,9 @@ pub fn run() { commands::schema_cache::delete_schema_cache_prefix, commands::tab_runtime_cache::save_tab_runtime_cache, commands::tab_runtime_cache::load_tab_runtime_cache, + commands::tab_runtime_cache::list_tab_runtime_cache_metadata, + commands::tab_runtime_cache::prune_tab_runtime_cache, + commands::tab_runtime_cache::delete_tab_runtime_cache_owner, commands::tab_runtime_cache::delete_tab_runtime_cache, commands::query::execute_query, commands::query::execute_multi, @@ -1457,6 +1460,8 @@ pub fn run() { if should_confirm_app_exit_request(std::env::consts::OS, *code, confirmed_exit) { api.prevent_exit(); request_app_close(app_handle, "quit"); + } else if let Some(state) = app_handle.try_state::>() { + tauri::async_runtime::block_on(state.shutdown_background_tasks(Duration::from_secs(3))); } }