From 16e879e40ec39b718c83295a7f97b9664a4aa680 Mon Sep 17 00:00:00 2001 From: miracle Date: Sun, 12 Jul 2026 11:18:52 +0800 Subject: [PATCH] perf(query): optimize result loading and connection reuse --- .../editor/EditorSettingsDialog.vue | 16 ++ apps/desktop/src/components/grid/DataGrid.vue | 2 +- .../src/components/layout/ContentArea.vue | 2 + .../src/composables/useDataGridActions.ts | 1 + apps/desktop/src/i18n/locales/en.ts | 2 + apps/desktop/src/i18n/locales/es.ts | 2 + apps/desktop/src/i18n/locales/it.ts | 2 + apps/desktop/src/i18n/locales/ja.ts | 2 + apps/desktop/src/i18n/locales/pt-BR.ts | 2 + apps/desktop/src/i18n/locales/zh-CN.ts | 2 + apps/desktop/src/i18n/locales/zh-TW.ts | 2 + .../src/lib/settings/editorSettingsDraft.ts | 1 + .../queryStore.hiddenPrimaryKey.spec.ts | 151 ++++++++++++++++-- apps/desktop/src/stores/queryStore.ts | 76 +++++++-- apps/desktop/src/stores/settingsStore.ts | 4 + crates/dbx-core/src/connection.rs | 15 +- crates/dbx-core/src/database_export.rs | 14 +- crates/dbx-core/src/db/postgres.rs | 8 +- crates/dbx-core/src/query.rs | 29 +++- crates/dbx-core/src/table_export.rs | 17 +- crates/dbx-core/src/table_import.rs | 6 +- crates/dbx-core/src/transfer.rs | 132 ++++++++++++++- crates/dbx-web/src/routes/query.rs | 20 +-- packages/app-tests/queryStore.test.ts | 6 + src-tauri/src/commands/database_export.rs | 4 + src-tauri/src/commands/table_export.rs | 4 + src-tauri/src/commands/table_import.rs | 10 +- 27 files changed, 476 insertions(+), 56 deletions(-) diff --git a/apps/desktop/src/components/editor/EditorSettingsDialog.vue b/apps/desktop/src/components/editor/EditorSettingsDialog.vue index 8adc1f939..65f58e5a2 100644 --- a/apps/desktop/src/components/editor/EditorSettingsDialog.vue +++ b/apps/desktop/src/components/editor/EditorSettingsDialog.vue @@ -283,6 +283,7 @@ const editCompactColumnHeaderActions = ref(settingsStore.editorSettings.compactC const editDataGridQuickEntry = ref(settingsStore.editorSettings.dataGridQuickEntry); const editInfiniteScroll = ref(settingsStore.editorSettings.infiniteScroll); const editInfiniteScrollMaxRows = ref(settingsStore.editorSettings.infiniteScrollMaxRows); +const editAutoCalculateTotalRows = ref(settingsStore.editorSettings.autoCalculateTotalRows); const editTableColumnTemplateRows = ref(tableColumnTemplateRowsFromSettings(settingsStore.editorSettings.tableColumnTemplateFields)); const editTableColumnTemplateDatabaseType = ref(TABLE_COLUMN_TEMPLATE_DATABASE_TYPES[0] ?? "mysql"); const editSqlVariableSyntaxOverrides = ref(normalizeSqlVariableSyntaxOverrides(settingsStore.editorSettings.sqlVariableSyntaxOverrides)); @@ -400,6 +401,7 @@ function currentEditorSettingsDraft(): EditorSettingsDraft { dataGridQuickEntry: editDataGridQuickEntry.value, infiniteScroll: editInfiniteScroll.value, infiniteScrollMaxRows: editInfiniteScrollMaxRows.value, + autoCalculateTotalRows: editAutoCalculateTotalRows.value, tableColumnTemplateFields: normalizedEditTableColumnTemplateFields.value, shortcuts: editShortcuts.value, sqlFormatter: normalizeSqlFormatterSettings(editSqlFormatter.value), @@ -672,6 +674,7 @@ function syncEditorSettingsDraftFromStore() { editDataGridQuickEntry.value = settingsStore.editorSettings.dataGridQuickEntry; editInfiniteScroll.value = settingsStore.editorSettings.infiniteScroll; editInfiniteScrollMaxRows.value = settingsStore.editorSettings.infiniteScrollMaxRows; + editAutoCalculateTotalRows.value = settingsStore.editorSettings.autoCalculateTotalRows; editTableColumnTemplateRows.value = tableColumnTemplateRowsFromSettings(settingsStore.editorSettings.tableColumnTemplateFields); editShortcuts.value = normalizeShortcutSettings(settingsStore.editorSettings.shortcuts); editSqlFormatter.value = normalizeSqlFormatterSettings(settingsStore.editorSettings.sqlFormatter); @@ -886,6 +889,7 @@ function resetDefaultsForTab(tab: SettingsCategory) { editDataGridQuickEntry.value = DEFAULT_EDITOR_SETTINGS.dataGridQuickEntry; editInfiniteScroll.value = DEFAULT_EDITOR_SETTINGS.infiniteScroll; editInfiniteScrollMaxRows.value = DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows; + editAutoCalculateTotalRows.value = DEFAULT_EDITOR_SETTINGS.autoCalculateTotalRows; editDuckDbWorkerProcessIsolation.value = DEFAULT_DESKTOP_SETTINGS.duckdb_worker_process_isolation; editDuckDbWorkerMaxProcesses.value = DEFAULT_DESKTOP_SETTINGS.duckdb_worker_max_processes; editTableColumnTemplateRows.value = tableColumnTemplateRowsFromSettings(DEFAULT_EDITOR_SETTINGS.tableColumnTemplateFields); @@ -938,6 +942,7 @@ function resetAllDefaults() { editDataGridQuickEntry.value = DEFAULT_EDITOR_SETTINGS.dataGridQuickEntry; editInfiniteScroll.value = DEFAULT_EDITOR_SETTINGS.infiniteScroll; editInfiniteScrollMaxRows.value = DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows; + editAutoCalculateTotalRows.value = DEFAULT_EDITOR_SETTINGS.autoCalculateTotalRows; editTableColumnTemplateRows.value = tableColumnTemplateRowsFromSettings(DEFAULT_EDITOR_SETTINGS.tableColumnTemplateFields); editShortcuts.value = normalizeShortcutSettings(DEFAULT_EDITOR_SETTINGS.shortcuts); editSqlFormatter.value = normalizeSqlFormatterSettings(DEFAULT_EDITOR_SETTINGS.sqlFormatter); @@ -3242,6 +3247,17 @@ onUnmounted(cleanupPreviewEditor); class="h-7 w-24 px-2 text-xs tabular-nums [appearance:textfield] [&::-webkit-inner-spin-button]:appearance-none [&::-webkit-outer-spin-button]:appearance-none" /> +
+
+ +

+ {{ t("settings.autoCalculateTotalRowsDescription") }} +

+
+ +
diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 2938ddc59..9c8a69d06 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -3618,7 +3618,7 @@ const canGoNextPage = computed(() => { }); const canJumpLastPage = computed(() => canGoNextPage.value && (hasKnownTotalRowCount.value || allRowsLoaded.value || !!props.tableMeta || !!props.countSql)); const totalRowCountBusy = computed(() => props.totalRowCountLoading === true || manualTotalRowCountLoading.value); -const canCalculateTotalRowCount = computed(() => !isResultsContext.value && !!props.connectionId && (!!props.tableMeta || !!props.countSql)); +const canCalculateTotalRowCount = computed(() => !!props.connectionId && (!!props.tableMeta || !!props.countSql)); // When a refresh/rollback completes and the current page exceeds the last // available page (e.g. data was deleted while viewing), auto-navigate to the // last available page instead of showing an empty page. diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 42323f233..972ddc2ec 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -1440,6 +1440,8 @@ defineExpose({ focusSearch, refreshData, handleModRTarget, requestQueryEditorExe :table-info-tab="activeTab.tableInfoTab" :page-offset="activeTab.resultPageOffset" :page-limit="activeTab.resultPageLimit" + :total-row-count="activeTab.resultTotalRowCount" + :total-row-count-loading="activeTab.resultTotalRowCountLoading" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" :full-export-result="(onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => queryStore.fetchTabResultForExport(activeTab.id, onProgress)" :export-file-base-name="activeTab.title" diff --git a/apps/desktop/src/composables/useDataGridActions.ts b/apps/desktop/src/composables/useDataGridActions.ts index dd9960c9f..ae21c2f9e 100644 --- a/apps/desktop/src/composables/useDataGridActions.ts +++ b/apps/desktop/src/composables/useDataGridActions.ts @@ -249,6 +249,7 @@ export function useDataGridActions(activeTab: ComputedRef) await queryStore.executeTabSql(tab.id, sql, { pagination: { offset, limit }, preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, }); } diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 7b019e2ec..d36e49994 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -3065,6 +3065,8 @@ export default { dataGridQuickEntry: "Quick grid entry", dataGridQuickEntryDescription: "When enabled, editing a cell or filling the bottom blank row saves to the database as soon as focus leaves. Useful for frequent entry, but it can cause accidental writes.", infiniteScroll: "Infinite scroll loading", + autoCalculateTotalRows: "Auto-calculate total row count", + autoCalculateTotalRowsDescription: "Run COUNT(*) automatically after each query to show the total matching rows. Off by default to keep large queries fast — calculate it on demand from the result footer.", infiniteScrollDescription: "Automatically load the next page of data when scrolling to the bottom of the table.", infiniteScrollMaxRows: "Infinite scroll max rows", infiniteScrollMaxRowsDescription: "Maximum number of rows to load in infinite scroll mode (1000–50000).", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index c6a2b4459..0d1a17709 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2917,6 +2917,8 @@ export default withEnglishFallback({ dataGridQuickEntry: "Entrada rápida en grilla", dataGridQuickEntryDescription: "Al activarlo, editar una celda o llenar la fila vacía inferior se guarda en la base de datos al perder el foco. Agiliza la captura frecuente, pero puede causar escrituras accidentales.", infiniteScroll: "Carga de desplazamiento infinito", + autoCalculateTotalRows: "Calcular automáticamente el total de filas", + autoCalculateTotalRowsDescription: "Ejecuta COUNT(*) automáticamente tras cada consulta para mostrar el total de filas coincidentes. Desactivado por defecto para mantener rápidas las consultas grandes; puedes calcularlo cuando quieras desde el pie de resultados.", infiniteScrollDescription: "Carga automáticamente la siguiente página de datos al desplazarse hasta el final de la tabla.", infiniteScrollMaxRows: "Máximo de filas en desplazamiento infinito", infiniteScrollMaxRowsDescription: "Número máximo de filas a cargar en modo desplazamiento infinito (1000–50000).", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 3cec6e5b6..585280e90 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2915,6 +2915,8 @@ export default withEnglishFallback({ dataGridQuickEntry: "Inserimento rapido griglia", dataGridQuickEntryDescription: "Se attivo, la modifica di una cella o la compilazione della riga vuota in fondo viene salvata nel database appena il focus esce. Utile per inserimenti frequenti, ma può causare scritture accidentali.", infiniteScroll: "Caricamento a scorrimento infinito", + autoCalculateTotalRows: "Calcola automaticamente il totale delle righe", + autoCalculateTotalRowsDescription: "Esegue COUNT(*) automaticamente dopo ogni query per mostrare il totale delle righe corrispondenti. Disattivato per impostazione predefinita per mantenere veloci le query grandi; puoi calcolarlo all'occorrenza dal piè di pagina dei risultati.", infiniteScrollDescription: "Carica automaticamente la pagina successiva dei dati quando scorri fino in fondo alla tabella.", infiniteScrollMaxRows: "Righe max a scorrimento infinito", infiniteScrollMaxRowsDescription: "Numero massimo di righe da caricare in modalità scorrimento infinito (1000–50000).", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 190871ed5..5699f8d42 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -3210,6 +3210,8 @@ export default withEnglishFallback({ updateDownloadSourceCnb: "CNB", updateDownloadSourceAtomgit: "AtomGit", infiniteScroll: "無限スクロール読み込み", + autoCalculateTotalRows: "総行数を自動計算", + autoCalculateTotalRowsDescription: "クエリごとに COUNT(*) を自動実行し、一致する総行数を表示します。大きなクエリを高速に保つため既定はオフです。結果フッターから必要に応じて計算できます。", infiniteScrollDescription: "テーブルのスクロール時に次のページのデータを自動読み込みします。", infiniteScrollMaxRows: "無限スクロールの最大行数", infiniteScrollMaxRowsDescription: "無限スクロールモードで読み込む最大行数（1000〜50000）。", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index 810989943..01e72d4c8 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2917,6 +2917,8 @@ export default withEnglishFallback({ dataGridQuickEntry: "Entrada rápida na grade", dataGridQuickEntryDescription: "Quando ativado, editar uma célula ou preencher a linha vazia inferior salva no banco de dados assim que o foco sai. Ajuda na entrada frequente, mas pode causar gravações acidentais.", infiniteScroll: "Carregamento por rolagem infinita", + autoCalculateTotalRows: "Calcular automaticamente o total de linhas", + autoCalculateTotalRowsDescription: "Executa COUNT(*) automaticamente após cada consulta para mostrar o total de linhas correspondentes. Desativado por padrão para manter consultas grandes rápidas; você pode calculá-lo quando quiser no rodapé dos resultados.", infiniteScrollDescription: "Carregar automaticamente a próxima página de dados ao rolar até o final da tabela.", infiniteScrollMaxRows: "Máximo de linhas em rolagem infinita", infiniteScrollMaxRowsDescription: "Número máximo de linhas a carregar no modo de rolagem infinita (1000–50000).", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 9d6821c3b..fe5c5398f 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -3064,6 +3064,8 @@ export default withEnglishFallback({ dataGridQuickEntry: "表格快捷录入", dataGridQuickEntryDescription: "开启后，编辑单元格或填写底部空行并离焦会立即保存到数据库。便于高频录入，但存在误操作风险。", infiniteScroll: "无限滚动加载", + autoCalculateTotalRows: "自动统计总行数", + autoCalculateTotalRowsDescription: "每次查询后自动执行 COUNT(*) 显示匹配的总行数。默认关闭以保证大查询速度 —— 可在结果栏按需手动统计。", infiniteScrollDescription: "滚动到表格底部时自动加载下一页数据，无需手动翻页。", infiniteScrollMaxRows: "无限滚动最大行数", infiniteScrollMaxRowsDescription: "无限滚动模式下最多加载的行数（1000–50000）。", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 4b6039a86..7c5ed5c07 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2766,6 +2766,8 @@ export default withEnglishFallback({ dataGridQuickEntry: "表格快速輸入", dataGridQuickEntryDescription: "開啟後，編輯儲存格或填寫底部空白列並離焦會立即儲存到資料庫。適合高頻輸入，但有誤操作風險。", infiniteScroll: "無限滾動載入", + autoCalculateTotalRows: "自動統計總筆數", + autoCalculateTotalRowsDescription: "每次查詢後自動執行 COUNT(*) 顯示符合的總筆數。預設關閉以確保大型查詢速度 —— 可在結果列按需手動統計。", infiniteScrollDescription: "滾動到表格底部時自動載入下一頁資料，無需手動翻頁。", infiniteScrollMaxRows: "無限滾動最大筆數", infiniteScrollMaxRowsDescription: "無限滾動模式下最多載入的筆數（1000–50000）。", diff --git a/apps/desktop/src/lib/settings/editorSettingsDraft.ts b/apps/desktop/src/lib/settings/editorSettingsDraft.ts index 8d85d0623..6043871bc 100644 --- a/apps/desktop/src/lib/settings/editorSettingsDraft.ts +++ b/apps/desktop/src/lib/settings/editorSettingsDraft.ts @@ -26,6 +26,7 @@ export const EDITOR_SETTINGS_DRAFT_KEYS = [ "dataGridQuickEntry", "infiniteScroll", "infiniteScrollMaxRows", + "autoCalculateTotalRows", "tableColumnTemplateFields", "shortcuts", "sqlFormatter", diff --git a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts index 613e45c09..ec2b40e6e 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts @@ -2,28 +2,37 @@ import { createPinia, setActivePinia } from "pinia"; import { beforeEach, describe, expect, it, vi } from "vitest"; const executeMulti = vi.fn(); +const executeQuery = vi.fn(); const analyzeEditableQueryEditability = vi.fn(); const getColumns = vi.fn(); const listIndexes = vi.fn(); const getConnectionConfig = vi.fn(); const buildSortedQuerySql = vi.fn(); +const buildDataGridCountSql = vi.fn(); +const prepareQueryPaginationExecutionPlan = vi.fn(async (options) => ({ + sqlToExecute: options.sql, + pageSql: undefined, + pageLimit: undefined, + pageOffset: undefined, + countSql: undefined, + useAgentResultSession: false, +})); +const editorSettings = { + pageSize: 100, + autoCalculateTotalRows: false, +}; vi.mock("@/lib/backend/api", () => ({ analyzeEditableQueryEditability, + buildDataGridCountSql, buildSortedQuerySql, closeClientConnectionSession: vi.fn().mockResolvedValue(undefined), closeQuerySession: vi.fn().mockResolvedValue(undefined), executeMulti, + executeQuery, getColumns, listIndexes, - prepareQueryPaginationExecutionPlan: vi.fn(async (options) => ({ - sqlToExecute: options.sql, - pageSql: undefined, - pageLimit: undefined, - pageOffset: undefined, - countSql: undefined, - useAgentResultSession: false, - })), + prepareQueryPaginationExecutionPlan, saveOpenTabsState: vi.fn().mockResolvedValue(undefined), })); @@ -37,7 +46,7 @@ vi.mock("@/stores/connectionStore", () => ({ vi.mock("@/stores/settingsStore", () => ({ useSettingsStore: () => ({ - editorSettings: { pageSize: 100 }, + editorSettings, }), })); @@ -68,6 +77,23 @@ describe("queryStore hidden primary key editing", () => { listIndexes.mockResolvedValue([]); analyzeEditableQueryEditability.mockImplementation(async (sql: string) => queryAnalysis(sql)); buildSortedQuerySql.mockImplementation(async (options) => ({ ok: true, sql: `${options.originalSql} ORDER BY ${options.column} ${options.direction.toUpperCase()}` })); + buildDataGridCountSql.mockResolvedValue("SELECT COUNT(*) FROM `users`"); + prepareQueryPaginationExecutionPlan.mockImplementation(async (options) => ({ + sqlToExecute: options.sql, + pageSql: undefined, + pageLimit: undefined, + pageOffset: undefined, + countSql: undefined, + useAgentResultSession: false, + })); + editorSettings.pageSize = 100; + editorSettings.autoCalculateTotalRows = false; + executeQuery.mockResolvedValue({ + columns: ["row_count"], + rows: [[0]], + affected_rows: 0, + execution_time_ms: 1, + }); executeMulti.mockResolvedValue([ { columns: ["name", "__DBX_PK_0"], @@ -339,4 +365,111 @@ describe("queryStore hidden primary key editing", () => { await vi.waitFor(() => expect(tab.queryEditabilityReason).toBe("primary-key-not-returned")); expect(tab.queryAnalysis).toBeUndefined(); }); + + it("records the returned row count when a page is known to be incomplete without count sql", async () => { + prepareQueryPaginationExecutionPlan.mockResolvedValue({ + sqlToExecute: "SELECT name FROM users LIMIT 100 OFFSET 0", + pageSql: "SELECT name FROM users LIMIT 100 OFFSET 0", + pageLimit: 100, + pageOffset: 0, + countSql: undefined, + useAgentResultSession: false, + }); + executeMulti.mockResolvedValue([ + { + columns: ["name"], + rows: Array.from({ length: 42 }, (_, index) => [`user-${index}`]), + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.resultTotalRowCount).toBe(42); + expect(tab.resultTotalRowCountLoading).toBe(false); + expect(executeQuery).not.toHaveBeenCalled(); + }); + + it("does not treat an empty later page as the total row count", async () => { + prepareQueryPaginationExecutionPlan.mockResolvedValue({ + sqlToExecute: "SELECT name FROM users LIMIT 100 OFFSET 200", + pageSql: "SELECT name FROM users LIMIT 100 OFFSET 200", + pageLimit: 100, + pageOffset: 200, + countSql: undefined, + useAgentResultSession: false, + }); + executeMulti.mockResolvedValue([ + { + columns: ["name"], + rows: [], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.resultTotalRowCount).toBeUndefined(); + expect(tab.resultTotalRowCountLoading).toBe(false); + expect(executeQuery).not.toHaveBeenCalled(); + }); + + it("automatically counts table data totals when the setting is enabled", async () => { + editorSettings.autoCalculateTotalRows = true; + executeMulti.mockResolvedValue([ + { + columns: ["id", "name"], + rows: Array.from({ length: 100 }, (_, index) => [index + 1, `user-${index + 1}`]), + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + executeQuery.mockResolvedValue({ + columns: ["row_count"], + rows: [[123]], + affected_rows: 0, + execution_time_ms: 1, + }); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "users", "data", "public"); + store.setTableMeta(tabId, { + schema: "public", + tableName: "users", + columns: [ + { name: "id", data_type: "int", is_nullable: false, is_primary_key: true, column_default: null, extra: null }, + { name: "name", data_type: "varchar", is_nullable: true, is_primary_key: false, column_default: null, extra: null }, + ], + primaryKeys: ["id"], + }); + + await store.executeTabSql(tabId, "SELECT id, name FROM users LIMIT 100", { + pagination: { limit: 100, offset: 0 }, + }); + + expect(buildDataGridCountSql).toHaveBeenCalledWith({ + databaseType: "mysql", + catalog: undefined, + schema: "public", + tableName: "users", + whereInput: undefined, + }); + await vi.waitFor(() => expect(executeQuery).toHaveBeenCalledWith("mysql-1", "app", "SELECT COUNT(*) FROM `users`", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 }))); + const tab = store.tabs.find((item) => item.id === tabId)!; + await vi.waitFor(() => expect(tab.resultTotalRowCount).toBe(123)); + expect(tab.resultTotalRowCountLoading).toBe(false); + }); }); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 8127893e1..f221101b9 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2339,32 +2339,61 @@ export const useQueryStore = defineStore("query", () => { function setQueryTotalRowCountIfCurrent(tabId: string, executionId: string, result: QueryResult, totalRowCount: number | undefined) { const current = tabs.value.find((t) => t.id === tabId); - if (current?.mode !== "query") return; + if (!current || (current.mode !== "query" && current.mode !== "data")) return; if (current.executionId !== executionId && current.result !== result) return; current.resultTotalRowCount = totalRowCount; current.resultTotalRowCountLoading = false; syncActiveResultRunFromDisplayed(current); } - function countQueryTotalRowsInBackground(options: { tabId: string; connectionId: string; database: string; schema?: string; countSql?: string; result: QueryResult; pageLimit?: number; pageOffset?: number; executionId: string; traceId: string; elapsed: () => string; timeoutSecs: number }) { + type TotalRowCountSqlTarget = { sql: string; schema?: string }; + + function countQueryTotalRowsInBackground(options: { + tabId: string; + connectionId: string; + database: string; + schema?: string; + countSql?: string; + countSqlTarget?: () => Promise; + result: QueryResult; + pageLimit?: number; + pageOffset?: number; + executionId: string; + traceId: string; + elapsed: () => string; + timeoutSecs: number; + }) { const resultRowCount = options.result.rows.length; - if (!options.countSql || resultRowCount <= 0) { + if (resultRowCount <= 0) { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); return; } - const countSql = options.countSql; - const clientSessionId = tabClientSessionId({ id: options.tabId }, "count"); - const countExecutionId = `${options.executionId}:count`; - if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount); return; } + // A full page was returned, so more rows may exist and determining the true + // total requires a potentially expensive COUNT(*) over the user's query. + // Only run it automatically when the user opted in; otherwise leave the + // total unknown and let them trigger it on demand from the result grid + // (matches DBeaver's default of not counting large result sets). + if (!useSettingsStore().editorSettings.autoCalculateTotalRows) { + setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); + return; + } + + const clientSessionId = tabClientSessionId({ id: options.tabId }, "count"); + const countExecutionId = `${options.executionId}:count`; void (async () => { try { + const countTarget = options.countSql ? { sql: options.countSql, schema: options.schema } : await options.countSqlTarget?.(); + if (!countTarget?.sql) { + setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); + return; + } console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() }); - const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, countExecutionId, { + const countResult = await api.executeQuery(options.connectionId, options.database, countTarget.sql, countTarget.schema, countExecutionId, { clientSessionId, timeoutSecs: options.timeoutSecs, }); @@ -2911,24 +2940,49 @@ export const useQueryStore = defineStore("query", () => { if (!options?.preserveTotalRowCountDuringExecution) { current.resultTotalRowCount = undefined; } - current.resultTotalRowCountLoading = current.mode === "query" && !!current.result && !!countSql; + const resultRowCount = current.result?.rows.length ?? 0; + const totalKnownFromIncompletePage = !!current.result && typeof pageLimit === "number" && resultRowCount < pageLimit; + const dataCountTarget = + current.mode === "data" + ? (() => { + const tableMeta = tableMetaForDataTab(current); + if (!tableMeta?.tableName) return undefined; + return { + databaseType: effectiveDbType, + catalog: tableMeta.catalog, + schema: tableMeta.schema, + tableName: tableMeta.tableName, + whereInput: current.whereInput?.trim() || undefined, + }; + })() + : undefined; + const canAutoCalculateTotalRows = !!current.result && resultRowCount > 0 && !totalKnownFromIncompletePage && settingsStore.editorSettings.autoCalculateTotalRows && ((current.mode === "query" && !!countSql) || (current.mode === "data" && !!dataCountTarget)); + current.resultTotalRowCountLoading = canAutoCalculateTotalRows; // Server-side pagination without a countSql: the backend (currently // the Elasticsearch driver) already reports the true match total via // affected_rows. Use it directly so the result-grid can compute the // page count without issuing a separate COUNT query. - if (current.result && current.mode === "query" && typeof pageLimit === "number" && !countSql && typeof current.result.affected_rows === "number") { + let totalRowCountResolved = false; + if (current.result && current.mode === "query" && typeof pageLimit === "number" && !countSql && typeof current.result.affected_rows === "number" && current.result.affected_rows > current.result.rows.length) { current.resultTotalRowCount = current.result.affected_rows; current.resultTotalRowCountLoading = false; + totalRowCountResolved = true; } touchResult(current); syncDisplayedResultRun(current, queryBaseSql); - if (current.mode === "query" && current.result) { + if (!totalRowCountResolved && (current.mode === "query" || current.mode === "data") && current.result) { countQueryTotalRowsInBackground({ tabId: id, connectionId: current.connectionId, database: current.database, schema: current.schema, countSql, + countSqlTarget: dataCountTarget + ? async () => ({ + sql: await api.buildDataGridCountSql(dataCountTarget), + schema: undefined, + }) + : undefined, result: current.result, pageLimit, pageOffset, diff --git a/apps/desktop/src/stores/settingsStore.ts b/apps/desktop/src/stores/settingsStore.ts index 26de55f92..493525a66 100644 --- a/apps/desktop/src/stores/settingsStore.ts +++ b/apps/desktop/src/stores/settingsStore.ts @@ -388,6 +388,7 @@ export interface EditorSettings { pageSize: number; infiniteScroll: boolean; infiniteScrollMaxRows: number; + autoCalculateTotalRows: boolean; mongoViewMode: "document" | "table"; showColumnCommentsInHeader: boolean; showColumnTypesInHeader: boolean; @@ -521,6 +522,7 @@ export const DEFAULT_EDITOR_SETTINGS: EditorSettings = { pageSize: 100, infiniteScroll: false, infiniteScrollMaxRows: 5000, + autoCalculateTotalRows: false, mongoViewMode: "document", showColumnCommentsInHeader: true, showColumnTypesInHeader: true, @@ -754,6 +756,7 @@ export function normalizeEditorSettings(settings: Partial, exist pageSize: normalizeResultPageSize(settings.pageSize), infiniteScroll: settings.infiniteScroll ?? DEFAULT_EDITOR_SETTINGS.infiniteScroll, infiniteScrollMaxRows: typeof settings.infiniteScrollMaxRows === "number" && settings.infiniteScrollMaxRows >= 1000 && settings.infiniteScrollMaxRows <= 50000 ? Math.round(settings.infiniteScrollMaxRows) : DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows, + autoCalculateTotalRows: settings.autoCalculateTotalRows ?? DEFAULT_EDITOR_SETTINGS.autoCalculateTotalRows, mongoViewMode: settings.mongoViewMode === "table" ? "table" : DEFAULT_EDITOR_SETTINGS.mongoViewMode, showColumnCommentsInHeader: settings.showColumnCommentsInHeader ?? DEFAULT_EDITOR_SETTINGS.showColumnCommentsInHeader, showColumnTypesInHeader: settings.showColumnTypesInHeader ?? DEFAULT_EDITOR_SETTINGS.showColumnTypesInHeader, @@ -999,6 +1002,7 @@ export const useSettingsStore = defineStore("settings", () => { if (partial.infiniteScroll !== undefined) editorSettings.value.infiniteScroll = partial.infiniteScroll; if (partial.infiniteScrollMaxRows !== undefined) editorSettings.value.infiniteScrollMaxRows = typeof partial.infiniteScrollMaxRows === "number" && partial.infiniteScrollMaxRows >= 1000 && partial.infiniteScrollMaxRows <= 50000 ? Math.round(partial.infiniteScrollMaxRows) : DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows; + if (partial.autoCalculateTotalRows !== undefined) editorSettings.value.autoCalculateTotalRows = partial.autoCalculateTotalRows === true; if (partial.mongoViewMode !== undefined) editorSettings.value.mongoViewMode = partial.mongoViewMode; if (partial.showColumnCommentsInHeader !== undefined) editorSettings.value.showColumnCommentsInHeader = partial.showColumnCommentsInHeader; if (partial.showColumnTypesInHeader !== undefined) editorSettings.value.showColumnTypesInHeader = partial.showColumnTypesInHeader; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 5438d7e7f..098eb5d05 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2648,6 +2648,10 @@ fn normalize_client_session_id(client_session_id: Option<&str>) -> Option String { + format!("{task_kind}:{task_id}") +} + fn mysql_pool_max_connections_for_session(client_session_id: Option<&str>) -> usize { if normalize_client_session_id(client_session_id).is_some() { 1 @@ -3078,8 +3082,8 @@ mod tests { metadata_connection_config, mysql_metadata_fallback_url, oceanbase_mysql_query_timeout_sql, oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint, redis_sentinel_transport_id, redis_sentinel_transport_prefix, sqlserver_legacy_agent_config, - sqlserver_legacy_agent_error, uses_bare_mysql_pool, uses_tcp_probe, validate_h2_database_path, AppState, - MysqlMode, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS, + sqlserver_legacy_agent_error, task_client_session_id, uses_bare_mysql_pool, uses_tcp_probe, + validate_h2_database_path, AppState, MysqlMode, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS, }; use crate::agent_connection::{ agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver, @@ -3149,6 +3153,13 @@ mod tests { } } + #[test] + fn task_client_session_ids_are_stable_and_isolated() { + assert_eq!(task_client_session_id("table-export", "job-1"), "table-export:job-1"); + assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("database-export", "job-1")); + assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("table-export", "job-2")); + } + #[test] fn access_agent_connect_timeout_has_longer_default_floor() { let mut config = mysql_config(None); diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index fdab47ff4..061770e14 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -4,6 +4,7 @@ use std::collections::HashSet; use std::io::Write; use tokio::sync::RwLock; +use crate::connection::task_client_session_id; use crate::models::connection::DatabaseType; use crate::object_source_sql::build_export_object_source_sql; use crate::sql_dialect::{qualified_table_name, quote_table_identifier, uses_single_row_insert_statements}; @@ -16,6 +17,10 @@ use crate::transfer::{ static EXPORT_CANCELLED: std::sync::LazyLock>> = std::sync::LazyLock::new(|| RwLock::new(HashSet::new())); +pub fn database_export_client_session_id(export_id: &str) -> String { + task_client_session_id("database-export", export_id) +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct DatabaseExportRequest { @@ -979,7 +984,10 @@ pub async fn export_database_sql_core( .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?; // 2. Get pool - let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; + let client_session_id = database_export_client_session_id(&request.export_id); + let pool_key = state + .get_or_create_pool_for_session(&request.connection_id, Some(&request.database), Some(&client_session_id)) + .await?; // 3. List tables let all_tables = crate::schema::list_tables_core( @@ -1213,7 +1221,7 @@ pub async fn export_database_sql_core( if !col_names.is_empty() { // Get row count let count_query = crate::transfer::count_sql(table_name, &request.schema, &db_type); - let total_rows = match crate::transfer::execute_on_pool(state, &pool_key, &count_query).await { + let total_rows = match crate::transfer::execute_read_on_pool(state, &pool_key, &count_query).await { Ok(result) => result.rows.first().and_then(|r| r.first()).and_then(|v| match v { serde_json::Value::Number(n) => n.as_u64(), serde_json::Value::String(s) => s.parse::().ok(), @@ -1251,7 +1259,7 @@ pub async fn export_database_sql_core( batch_size, ); - let result = match crate::transfer::execute_on_pool(state, &pool_key, &sql).await { + let result = match crate::transfer::execute_read_on_pool(state, &pool_key, &sql).await { Ok(r) => r, Err(e) => { writeln!(file, "-- ERROR exporting data for table {table_name}: {e}") diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 00e255d1d..c7c9f438a 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -985,7 +985,13 @@ pub async fn connect(url: &str, fallback_timeout: Duration) -> Result, err: &str) matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) } +fn query_pool_error_action(db_type: Option, sql: &str, err: &str) -> PoolErrorAction { + match pool_error_action(db_type, err) { + // A connection error does not prove that the database did not receive + // a write. Only replay statements already accepted by the read-only + // protection classifier; writes discard the stale pool without retry. + PoolErrorAction::ReconnectAndRetry if is_write_sql(sql) => PoolErrorAction::Discard, + action => action, + } +} + fn is_os_connection_error(lower: &str) -> bool { let os_error_codes = ["10053", "10054", "10057", "10058", "10060", "10061"]; if let Some(pos) = lower.find("os error ") { @@ -1577,7 +1588,7 @@ pub async fn execute_sql_statement_with_options( do_execute(state, &pool_key, mysql_dialect, Some(database), sql, schema, cancel_token.clone(), options.clone()) .await; - let action = result.as_ref().err().map(|e| pool_error_action(db_type, e)); + let action = result.as_ref().err().map(|e| query_pool_error_action(db_type, sql, e)); match action { Some(PoolErrorAction::ReconnectAndRetry) if !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; @@ -3067,6 +3078,22 @@ mod tests { assert!(!is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_query")); } + #[test] + fn query_pool_error_policy_retries_reads_but_not_writes() { + assert_eq!( + query_pool_error_action(Some(DatabaseType::Postgres), "SELECT * FROM users", "connection reset by peer"), + PoolErrorAction::ReconnectAndRetry + ); + assert_eq!( + query_pool_error_action( + Some(DatabaseType::Postgres), + "UPDATE users SET active = true", + "connection reset by peer" + ), + PoolErrorAction::Discard + ); + } + #[test] fn external_driver_method_unsupported_detects_legacy_plugin_errors() { assert!(is_external_driver_method_unsupported( diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index df9f0dd5b..15919bd2c 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -7,14 +7,14 @@ use std::time::Duration; use tokio_util::sync::CancellationToken; use crate::connection::MysqlMode; -use crate::connection::{AppState, PoolKind}; +use crate::connection::{task_client_session_id, AppState, PoolKind}; use crate::csv_export::{escape_csv, format_csv, value_to_csv_text}; pub use crate::database_export::ExportStatus; use crate::database_export::{build_export_insert_statements, is_export_cancelled, BuildExportInsertStatementsOptions}; use crate::db::agent_driver::AgentTableReadStartParams; use crate::models::connection::DatabaseType; use crate::transfer::{ - count_sql_with_where, execute_on_pool, execute_on_pool_with_max_rows, keyset_pagination_sql, + count_sql_with_where, execute_read_on_pool, execute_read_on_pool_with_max_rows, keyset_pagination_sql, pagination_sql_with_filter_order, qualified_table, quote_identifier, }; use crate::types::QueryResult; @@ -23,6 +23,10 @@ use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_wo const DEFAULT_BATCH_SIZE: usize = 10_000; const SQL_INSERT_BATCH_SIZE: usize = 100; +pub fn table_export_client_session_id(export_id: &str) -> String { + task_client_session_id("table-export", export_id) +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TableExportRequest { @@ -353,7 +357,7 @@ async fn fetch_paginated_table_export_batch( offset, active_batch_size, ); - execute_on_pool_with_max_rows(state, pool_key, &sql, Some(active_batch_size)).await + execute_read_on_pool_with_max_rows(state, pool_key, &sql, Some(active_batch_size)).await } async fn close_table_read_session_if_open( @@ -773,7 +777,10 @@ pub async fn export_table_data_core( .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?; // 2. Get pool - let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; + let client_session_id = table_export_client_session_id(&request.export_id); + let pool_key = state + .get_or_create_pool_for_session(&request.connection_id, Some(&request.database), Some(&client_session_id)) + .await?; // 3. Resolve columns. Data grid exports can provide columns/primary keys // directly, which avoids expensive metadata round-trips on JDBC drivers. @@ -830,7 +837,7 @@ pub async fn export_table_data_core( &db_type, request.where_input.as_deref(), ); - match execute_on_pool(state, &pool_key, &count_query).await { + match execute_read_on_pool(state, &pool_key, &count_query).await { Ok(result) => result .rows .first() diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index 91f42f30a..72b3c728e 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -6,7 +6,7 @@ use calamine::{open_workbook_auto, Data, Reader}; use chrono::{DateTime, NaiveDate, NaiveDateTime}; use serde::{Deserialize, Serialize}; -use crate::connection::AppState; +use crate::connection::{task_client_session_id, AppState}; use crate::models::connection::DatabaseType; use crate::transfer::{ execute_on_pool, generate_insert_typed, get_columns_for_transfer, qualified_table, quote_identifier, @@ -17,6 +17,10 @@ pub const DEFAULT_BATCH_SIZE: usize = 500; pub const CREATE_TABLE_INFERENCE_ROWS: usize = 100; pub const MAX_NON_STREAMING_IMPORT_BYTES: u64 = 100 * 1024 * 1024; +pub fn table_import_client_session_id(import_id: &str) -> String { + task_client_session_id("table-import", import_id) +} + #[derive(Debug, Clone)] pub struct ParsedImportFile { pub columns: Vec, diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 6dfef9222..1df26d659 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -3,12 +3,12 @@ use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; use tokio::sync::RwLock; -use crate::connection::{AppState, PoolKind}; +use crate::connection::{config_for_pool_key, AppState, PoolKind}; use crate::db; use crate::db::mongo_driver::MongoDocumentResult; use crate::models::connection::DatabaseType; use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput}; -use crate::query::{agent_execute_query_params, should_discard_pool_after_error, QueryExecutionOptions}; +use crate::query::{agent_execute_query_params, pool_error_action, PoolErrorAction, QueryExecutionOptions}; use crate::sql::split_sql_statements; #[cfg(feature = "duckdb-bundled")] use crate::sql::starts_with_executable_sql_keyword; @@ -2513,7 +2513,20 @@ fn mongo_columns_from_documents(documents: &[serde_json::Value]) -> Vec Result { - execute_on_pool_with_max_rows(state, pool_key, sql, None).await + execute_on_pool_with_options(state, pool_key, sql, None, TransferExecutionSafety::WriteNoReplay).await +} + +pub async fn execute_read_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result { + execute_read_on_pool_with_max_rows(state, pool_key, sql, None).await +} + +pub async fn execute_read_on_pool_with_max_rows( + state: &AppState, + pool_key: &str, + sql: &str, + max_rows: Option, +) -> Result { + execute_on_pool_with_options(state, pool_key, sql, max_rows, TransferExecutionSafety::ReadOnlyRetryable).await } async fn execute_transfer_ddl_on_pool( @@ -2596,6 +2609,91 @@ pub async fn execute_on_pool_with_max_rows( pool_key: &str, sql: &str, max_rows: Option, +) -> Result { + execute_on_pool_with_options(state, pool_key, sql, max_rows, TransferExecutionSafety::WriteNoReplay).await +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TransferExecutionSafety { + ReadOnlyRetryable, + WriteNoReplay, +} + +fn transfer_pool_error_action( + safety: TransferExecutionSafety, + db_type: Option, + err: &str, +) -> PoolErrorAction { + match (safety, pool_error_action(db_type, err)) { + (TransferExecutionSafety::WriteNoReplay, PoolErrorAction::ReconnectAndRetry) => PoolErrorAction::Discard, + (_, action) => action, + } +} + +async fn transfer_pool_context( + state: &AppState, + pool_key: &str, +) -> (Option, Option, Option) { + let configs = state.configs.read().await; + let config = config_for_pool_key(pool_key, &configs); + ( + config.map(|config| config.id.clone()), + database_from_pool_key(pool_key).map(str::to_string), + config.map(|config| config.db_type), + ) +} + +fn client_session_id_from_pool_key(pool_key: &str) -> Option<&str> { + pool_key.split_once(":session:").map(|(_, session)| session).filter(|session| !session.is_empty()) +} + +async fn execute_on_pool_with_options( + state: &AppState, + pool_key: &str, + sql: &str, + max_rows: Option, + safety: TransferExecutionSafety, +) -> Result { + let (connection_id, database, db_type) = transfer_pool_context(state, pool_key).await; + let client_session_id = client_session_id_from_pool_key(pool_key).map(str::to_string); + let mut current_pool_key = pool_key.to_string(); + + for attempt in 0..2 { + let result = execute_on_pool_once(state, ¤t_pool_key, sql, max_rows).await; + let Some(error) = result.as_ref().err() else { + return result; + }; + + match transfer_pool_error_action(safety, db_type, error) { + PoolErrorAction::Keep => return result, + PoolErrorAction::Discard => { + state.remove_pool_by_key(¤t_pool_key).await; + return result; + } + PoolErrorAction::ReconnectAndRetry if attempt == 0 => { + let Some(connection_id) = connection_id.as_deref() else { + state.remove_pool_by_key(¤t_pool_key).await; + return result; + }; + current_pool_key = state + .reconnect_pool_for_session(connection_id, database.as_deref(), client_session_id.as_deref()) + .await?; + } + PoolErrorAction::ReconnectAndRetry => { + state.remove_pool_by_key(¤t_pool_key).await; + return result; + } + } + } + + unreachable!("transfer pool execution retry loop runs at most twice") +} + +async fn execute_on_pool_once( + state: &AppState, + pool_key: &str, + sql: &str, + max_rows: Option, ) -> Result { // Read-only check: block transfer operations in readonly mode crate::query::check_read_only_for_connection(state, pool_key, sql).await?; @@ -2631,10 +2729,6 @@ pub async fn execute_on_pool_with_max_rows( let mut client = client.lock().await; let result = db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).await; drop(client); - if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(Some(DatabaseType::SqlServer), err)) - { - state.remove_pool_by_key(pool_key).await; - } result } PoolKind::Agent(client) => { @@ -6248,6 +6342,30 @@ SELECT 1 FROM dual"# assert_eq!(database_from_pool_key("conn"), None); } + #[test] + fn transfer_read_retry_policy_retries_connection_errors() { + assert_eq!( + transfer_pool_error_action( + TransferExecutionSafety::ReadOnlyRetryable, + Some(DatabaseType::Postgres), + "connection reset by peer" + ), + PoolErrorAction::ReconnectAndRetry + ); + } + + #[test] + fn transfer_write_retry_policy_discards_without_replaying_batch() { + assert_eq!( + transfer_pool_error_action( + TransferExecutionSafety::WriteNoReplay, + Some(DatabaseType::Postgres), + "connection reset by peer" + ), + PoolErrorAction::Discard + ); + } + #[test] fn map_column_type_preserves_longtext_for_mysql_target() { assert_eq!(map_column_type("longtext", &DatabaseType::Mysql, &DatabaseType::Mysql), "longtext"); diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 5293cba94..e9ca5c309 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -309,7 +309,7 @@ pub struct BuildDatabaseSqlExportRequest { pub async fn execute_query( State(state): State>, Json(req): Json, -) -> Result, AppError> { +) -> Result, AppError> { let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let registered = state.app.running_queries.register_task( @@ -341,13 +341,13 @@ pub async fn execute_query( .map_err(AppError)?; drop(registered); - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) + Ok(Json(result)) } pub async fn execute_multi( State(state): State>, Json(req): Json, -) -> Result, AppError> { +) -> Result>, AppError> { let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let registered = state.app.running_queries.register_task( @@ -379,13 +379,13 @@ pub async fn execute_multi( .map_err(AppError)?; drop(registered); - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) + Ok(Json(result)) } pub async fn execute_batch( State(state): State>, Json(req): Json, -) -> Result, AppError> { +) -> Result, AppError> { let result = dbx_core::query::execute_statements( &state.app, &req.connection_id, @@ -397,7 +397,7 @@ pub async fn execute_batch( .await .map_err(AppError)?; - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) + Ok(Json(result)) } pub async fn cancel_query( @@ -442,7 +442,7 @@ pub async fn close_client_connection_session( pub async fn execute_script( State(state): State>, Json(req): Json, -) -> Result, AppError> { +) -> Result, AppError> { let db_type = { let configs = state.app.configs.read().await; configs.get(&req.connection_id).map(|config| config.db_type) @@ -461,13 +461,13 @@ pub async fn execute_script( .await .map_err(AppError)?; - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) + Ok(Json(result)) } pub async fn execute_in_transaction( State(state): State>, Json(req): Json, -) -> Result, AppError> { +) -> Result, AppError> { let result = dbx_core::query::execute_statements_in_transaction( &state.app, &req.connection_id, @@ -478,7 +478,7 @@ pub async fn execute_in_transaction( .await .map_err(AppError)?; - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) + Ok(Json(result)) } pub async fn analyze_sql_references( diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index ca3ffe44b..363726c5f 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -4083,9 +4083,12 @@ test("query execution keeps automatically counting total rows in the background" const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); const store = useQueryStore(); const originalFetch = globalThis.fetch; + settingsStore.updateEditorSettings({ autoCalculateTotalRows: true }); + connectionStore.addEphemeralConnection(conn("conn-1")); const tabId = store.createTab("conn-1", "db", "Query", "query", "public"); const tab = store.tabs.find((item) => item.id === tabId); @@ -4168,9 +4171,12 @@ test("paginated query execution keeps the previous total while refreshing it in const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); const store = useQueryStore(); const originalFetch = globalThis.fetch; + settingsStore.updateEditorSettings({ autoCalculateTotalRows: true }); + connectionStore.addEphemeralConnection(conn("conn-1")); const tabId = store.createTab("conn-1", "db", "Query", "query", "public"); const tab = store.tabs.find((item) => item.id === tabId); diff --git a/src-tauri/src/commands/database_export.rs b/src-tauri/src/commands/database_export.rs index d0b7d9f06..d4119a0b0 100644 --- a/src-tauri/src/commands/database_export.rs +++ b/src-tauri/src/commands/database_export.rs @@ -24,6 +24,10 @@ pub async fn export_database_sql( }) .await; + let client_session_id = dbx_core::database_export::database_export_client_session_id(&export_id); + let _ = + state.close_client_session_pool(&request.connection_id, Some(&request.database), &client_session_id).await; + if let Err(e) = result { emit_progress( &app, diff --git a/src-tauri/src/commands/table_export.rs b/src-tauri/src/commands/table_export.rs index ec0ce3bf6..9cdde5ca5 100644 --- a/src-tauri/src/commands/table_export.rs +++ b/src-tauri/src/commands/table_export.rs @@ -23,6 +23,10 @@ pub async fn start_table_export( dbx_core::table_export::export_table_data_core(&state, &request, |progress| emit_progress(&app, progress)) .await; + let client_session_id = dbx_core::table_export::table_export_client_session_id(&export_id); + let _ = + state.close_client_session_pool(&request.connection_id, Some(&request.database), &client_session_id).await; + if let Err(e) = result { emit_progress( &app, diff --git a/src-tauri/src/commands/table_import.rs b/src-tauri/src/commands/table_import.rs index 5b6678f19..9e282007c 100644 --- a/src-tauri/src/commands/table_import.rs +++ b/src-tauri/src/commands/table_import.rs @@ -45,11 +45,10 @@ pub async fn import_table_file( // Reject import early if the connection is read-only — importing is inherently a write operation ensure_connection_writable(&state, &request.connection_id, "Import").await?; let db_type = get_db_type(&state, &request.connection_id).await?; - let pool_key = if request.database.is_empty() { - request.connection_id.clone() - } else { - state.get_or_create_pool(&request.connection_id, Some(&request.database)).await? - }; + let database = (!request.database.trim().is_empty()).then_some(request.database.as_str()); + let client_session_id = dbx_core::table_import::table_import_client_session_id(&request.import_id); + let pool_key = + state.get_or_create_pool_for_session(&request.connection_id, database, Some(&client_session_id)).await?; let result = dbx_core::table_import::import_table_file_core( &state, @@ -61,6 +60,7 @@ pub async fn import_table_file( ) .await; + let _ = state.close_client_session_pool(&request.connection_id, database, &client_session_id).await; clear_cancelled(&request.import_id).await; result }