perf(query): optimize result loading and connection reuse
This commit is contained in:
miracle
GitHub
parent
24c63b6012
commit
16e879e40e
|
|
@ -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<TableColumnTemplateGridRow[]>(tableColumnTemplateRowsFromSettings(settingsStore.editorSettings.tableColumnTemplateFields));
|
||||
const editTableColumnTemplateDatabaseType = ref<DatabaseType>(TABLE_COLUMN_TEMPLATE_DATABASE_TYPES[0] ?? "mysql");
|
||||
const editSqlVariableSyntaxOverrides = ref<SqlVariableSyntaxOverrides>(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"
|
||||
/>
|
||||
</div>
|
||||
<div class="flex items-center justify-between gap-4 rounded-md border bg-muted/20 px-3 py-2">
|
||||
<div class="space-y-1">
|
||||
<Label for="auto-calculate-total-rows">
|
||||
{{ t("settings.autoCalculateTotalRows") }}
|
||||
</Label>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
{{ t("settings.autoCalculateTotalRowsDescription") }}
|
||||
</p>
|
||||
</div>
|
||||
<Switch id="auto-calculate-total-rows" v-model="editAutoCalculateTotalRows" />
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<Separator />
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -249,6 +249,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
|
|||
await queryStore.executeTabSql(tab.id, sql, {
|
||||
pagination: { offset, limit },
|
||||
preserveResultDuringExecution: true,
|
||||
preserveTotalRowCountDuringExecution: true,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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).",
|
||||
|
|
|
|||
|
|
@ -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).",
|
||||
|
|
|
|||
|
|
@ -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).",
|
||||
|
|
|
|||
|
|
@ -3210,6 +3210,8 @@ export default withEnglishFallback({
|
|||
updateDownloadSourceCnb: "CNB",
|
||||
updateDownloadSourceAtomgit: "AtomGit",
|
||||
infiniteScroll: "無限スクロール読み込み",
|
||||
autoCalculateTotalRows: "総行数を自動計算",
|
||||
autoCalculateTotalRowsDescription: "クエリごとに COUNT(*) を自動実行し、一致する総行数を表示します。大きなクエリを高速に保つため既定はオフです。結果フッターから必要に応じて計算できます。",
|
||||
infiniteScrollDescription: "テーブルのスクロール時に次のページのデータを自動読み込みします。",
|
||||
infiniteScrollMaxRows: "無限スクロールの最大行数",
|
||||
infiniteScrollMaxRowsDescription: "無限スクロールモードで読み込む最大行数（1000〜50000）。",
|
||||
|
|
|
|||
|
|
@ -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).",
|
||||
|
|
|
|||
|
|
@ -3064,6 +3064,8 @@ export default withEnglishFallback({
|
|||
dataGridQuickEntry: "表格快捷录入",
|
||||
dataGridQuickEntryDescription: "开启后，编辑单元格或填写底部空行并离焦会立即保存到数据库。便于高频录入，但存在误操作风险。",
|
||||
infiniteScroll: "无限滚动加载",
|
||||
autoCalculateTotalRows: "自动统计总行数",
|
||||
autoCalculateTotalRowsDescription: "每次查询后自动执行 COUNT(*) 显示匹配的总行数。默认关闭以保证大查询速度 —— 可在结果栏按需手动统计。",
|
||||
infiniteScrollDescription: "滚动到表格底部时自动加载下一页数据，无需手动翻页。",
|
||||
infiniteScrollMaxRows: "无限滚动最大行数",
|
||||
infiniteScrollMaxRowsDescription: "无限滚动模式下最多加载的行数（1000–50000）。",
|
||||
|
|
|
|||
|
|
@ -2766,6 +2766,8 @@ export default withEnglishFallback({
|
|||
dataGridQuickEntry: "表格快速輸入",
|
||||
dataGridQuickEntryDescription: "開啟後，編輯儲存格或填寫底部空白列並離焦會立即儲存到資料庫。適合高頻輸入，但有誤操作風險。",
|
||||
infiniteScroll: "無限滾動載入",
|
||||
autoCalculateTotalRows: "自動統計總筆數",
|
||||
autoCalculateTotalRowsDescription: "每次查詢後自動執行 COUNT(*) 顯示符合的總筆數。預設關閉以確保大型查詢速度 —— 可在結果列按需手動統計。",
|
||||
infiniteScrollDescription: "滾動到表格底部時自動載入下一頁資料，無需手動翻頁。",
|
||||
infiniteScrollMaxRows: "無限滾動最大筆數",
|
||||
infiniteScrollMaxRowsDescription: "無限滾動模式下最多載入的筆數（1000–50000）。",
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ export const EDITOR_SETTINGS_DRAFT_KEYS = [
|
|||
"dataGridQuickEntry",
|
||||
"infiniteScroll",
|
||||
"infiniteScrollMaxRows",
|
||||
"autoCalculateTotalRows",
|
||||
"tableColumnTemplateFields",
|
||||
"shortcuts",
|
||||
"sqlFormatter",
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<TotalRowCountSqlTarget | undefined>;
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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<EditorSettings>, 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;
|
||||
|
|
|
|||
|
|
@ -2648,6 +2648,10 @@ fn normalize_client_session_id(client_session_id: Option<&str>) -> Option<String
|
|||
client_session_id.map(str::trim).filter(|session| !session.is_empty()).map(|session| session.replace(':', "_"))
|
||||
}
|
||||
|
||||
pub fn task_client_session_id(task_kind: &str, task_id: &str) -> 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);
|
||||
|
|
|
|||
|
|
@ -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<RwLock<HashSet<String>>> =
|
||||
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::<u64>().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}")
|
||||
|
|
|
|||
|
|
@ -985,7 +985,13 @@ pub async fn connect(url: &str, fallback_timeout: Duration) -> Result<Pool, Stri
|
|||
let pg_config = tokio_postgres::Config::from_str(&postgres_url.url)
|
||||
.map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?;
|
||||
|
||||
let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Verified };
|
||||
// Fast recycling only checks whether the connection is already closed
|
||||
// instead of issuing a validation query on every checkout, saving one
|
||||
// round-trip per query. Connections that went stale without being
|
||||
// observed are caught when the query runs and recovered by the
|
||||
// executor's ReconnectAndRetry path (see pool_error_action / do_execute
|
||||
// in query.rs).
|
||||
let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Fast };
|
||||
let tls_config = postgres_tls_config(
|
||||
&pg_config,
|
||||
&postgres_url.ssl_files,
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ use crate::connection::{AppState, PoolKind, TransactionSession, TxnConnection};
|
|||
use crate::database_capabilities;
|
||||
use crate::db;
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseType};
|
||||
use crate::query_execution_sql::is_write_sql;
|
||||
#[cfg(feature = "duckdb-bundled")]
|
||||
use crate::sql::starts_with_duckdb_result_sql_keyword;
|
||||
use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable_sql_keyword};
|
||||
|
|
@ -841,6 +842,16 @@ pub fn should_discard_pool_after_error(db_type: Option<DatabaseType>, err: &str)
|
|||
matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)
|
||||
}
|
||||
|
||||
fn query_pool_error_action(db_type: Option<DatabaseType>, 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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
|
|
|
|||
|
|
@ -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<db::Colu
|
|||
}
|
||||
|
||||
pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
|
||||
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<db::QueryResult, String> {
|
||||
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<usize>,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
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<usize>,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
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<DatabaseType>,
|
||||
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<String>, Option<String>, Option<DatabaseType>) {
|
||||
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<usize>,
|
||||
safety: TransferExecutionSafety,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
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<usize>,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
// 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");
|
||||
|
|
|
|||
|
|
@ -309,7 +309,7 @@ pub struct BuildDatabaseSqlExportRequest {
|
|||
pub async fn execute_query(
|
||||
State(state): State<Arc<WebState>>,
|
||||
Json(req): Json<ExecuteQueryRequest>,
|
||||
) -> Result<Json<serde_json::Value>, AppError> {
|
||||
) -> Result<Json<dbx_core::db::QueryResult>, 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<Arc<WebState>>,
|
||||
Json(req): Json<ExecuteQueryRequest>,
|
||||
) -> Result<Json<serde_json::Value>, AppError> {
|
||||
) -> Result<Json<Vec<dbx_core::db::QueryResult>>, 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<Arc<WebState>>,
|
||||
Json(req): Json<ExecuteBatchRequest>,
|
||||
) -> Result<Json<serde_json::Value>, AppError> {
|
||||
) -> Result<Json<dbx_core::db::QueryResult>, 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<Arc<WebState>>,
|
||||
Json(req): Json<ExecuteQueryRequest>,
|
||||
) -> Result<Json<serde_json::Value>, AppError> {
|
||||
) -> Result<Json<dbx_core::db::QueryResult>, 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<Arc<WebState>>,
|
||||
Json(req): Json<ExecuteBatchRequest>,
|
||||
) -> Result<Json<serde_json::Value>, AppError> {
|
||||
) -> Result<Json<dbx_core::db::QueryResult>, 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(
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue