diff --git a/agents/common/src/test/java/com/dbx/agent/JsonRpcLargeValueTest.java b/agents/common/src/test/java/com/dbx/agent/JsonRpcLargeValueTest.java index 7af2b2b12c..a09f6cfdda 100644 --- a/agents/common/src/test/java/com/dbx/agent/JsonRpcLargeValueTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JsonRpcLargeValueTest.java @@ -46,4 +46,27 @@ class JsonRpcLargeValueTest { JsonParser.parseString("{}").getAsJsonObject()); assertEquals(List.of("invalidate", "commit"), calls); } + + @Test + void invalidatesBeforeFailedWriteAndRollbackAndOnCancellation() throws Exception { + List calls = new ArrayList<>(); + DatabaseAgent agent = (DatabaseAgent) Proxy.newProxyInstance(DatabaseAgent.class.getClassLoader(), + new Class[]{DatabaseAgent.class}, (object, method, args) -> { + if (method.getName().equals("supportsQueryTiming")) return false; + if (method.getName().equals("invalidateLargeValues")) { calls.add("invalidate"); return null; } + if (method.getName().equals("executeBatch")) { calls.add("write"); throw new java.sql.SQLException("concurrent row removed"); } + if (method.getName().equals("rollbackManualTransaction")) { calls.add("rollback"); return Map.of("ok", true); } + return null; + }); + JsonRpcServer server = new JsonRpcServer(agent); + assertThrows(Exception.class, () -> server.dispatchForRuntime(AgentProtocol.METHOD_EXECUTE_BATCH, + JsonParser.parseString("{\"statements\":[\"UPDATE T SET C='full original value'\"]}").getAsJsonObject())); + assertEquals(List.of("invalidate", "write"), calls); + calls.clear(); + server.dispatchForRuntime(AgentProtocol.METHOD_ROLLBACK_MANUAL_TRANSACTION, JsonParser.parseString("{}").getAsJsonObject()); + assertEquals(List.of("invalidate", "rollback"), calls); + calls.clear(); + server.cancelActiveStatements(); + assertEquals(List.of("invalidate"), calls); + } } diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index c3cddd2296..eb2338e01e 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -8663,6 +8663,7 @@ const { queryResultExportRequest: props.queryResultExportRequest, hasCompleteLocalResult, completeLocalResult: computed(() => (hasCompleteLocalResult.value ? props.result : undefined)), + snapshotResult: computed(() => props.result), allExportResults: computed(() => props.allExportResults), currentResultLabel: computed(() => props.result.sourceLabel), exportFileBaseName: computed(() => props.exportFileBaseName), diff --git a/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts index 0abcecf613..f242ac8754 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts @@ -2,13 +2,13 @@ import { computed, ref } from "vue"; import { createPinia, setActivePinia } from "pinia"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { useDataGridExport, type UseDataGridExportOptions } from "@/composables/useDataGridExport"; -import type { DatabaseType } from "@/types/database"; +import type { DatabaseType, QueryResult } from "@/types/database"; import { buildDataGridCopyUpdateStatements } from "@/lib/dataGrid/dataGridSql"; import { copyToClipboard } from "@/lib/common/clipboard"; import type { DataGridTableMeta } from "@/lib/dataGrid/dataGridSql"; import type { CellSelectionMatrix, SelectionData } from "@/lib/dataGrid/gridSelection"; import type { CellValue } from "@/lib/dataGrid/cellValue"; -import { extractDataGridSelection, exportQueryResultCsv, exportQueryResultJson } from "@/lib/backend/api"; +import { extractDataGridSelection, exportQueryResultCsv, exportQueryResultJson, exportQueryResultMarkdown, exportSnapshotResult, readLargeValueChunk } from "@/lib/backend/api"; import { DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS } from "@/lib/dataGrid/dataGridCopyExtractor"; import { clearDataGridClipboardCopy, parseDataGridClipboard } from "@/lib/dataGrid/dataGridClipboard"; import { MONGO_DOCUMENT_GRID_NULL, mongoDocumentGridExternalValue } from "@/lib/mongo/mongoDocumentValues"; @@ -52,6 +52,10 @@ vi.mock("@/lib/backend/api", async (importOriginal) => { extractDataGridSelection: vi.fn(), exportQueryResultCsv: vi.fn(), exportQueryResultJson: vi.fn(), + exportQueryResultMarkdown: vi.fn(), + exportSnapshotResult: vi.fn(), + readLargeValueChunk: vi.fn(), + cancelQuery: vi.fn(), }; }); @@ -143,6 +147,7 @@ function createExportState( contextColumn?: number, databaseType: DatabaseType = "mysql", displayValue?: (value: CellValue, columnIndex: number) => string, + overrides: Partial = {}, ) { const rows = (rowDataList ?? [rowData ?? columns.map((column, index) => (column === "id" ? 1 : `value-${index}`))]).map((data, index) => ({ ...row(data), id: index + 1 })); const resolvedContextRowId = contextRowId === undefined ? (rows[0]?.id ?? null) : contextRowId; @@ -175,9 +180,56 @@ function createExportState( selectedRowIds, hasRowSelection: computed(() => selectedRowIds.value.size > 0), }; - return useDataGridExport(options); + return useDataGridExport({ ...options, ...overrides }); } +describe("original LOB result exports", () => { + beforeEach(() => { + vi.clearAllMocks(); + setActivePinia(createPinia()); + vi.mocked(exportSnapshotResult).mockResolvedValue(undefined); + }); + function snapshotState() { + const original = { + columns: ["id", "name"], rows: [[1, "preview"]], column_types: ["NUMBER", "CLOB"], + large_value_context: { connectionId: "original-connection", database: "original-database", clientSessionId: "original-session" }, + large_value_cells: [{ row_index: 0, column_index: 1, value_ref: "original-locator", value_kind: "text" }], + } as QueryResult; + return { original, state: createExportState(editableTable, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, "oracle", undefined, { + snapshotResult: computed(() => original), hasCompleteLocalResult: computed(() => true), completeLocalResult: computed(() => original), + }) }; + } + it.each(["csv", "json"] as const)("streams export-all %s from captured locators without frontend hydration or requery", async (format) => { + const { state } = snapshotState(); + await (format === "csv" ? state.exportCsv() : state.exportJson()); + expect(exportSnapshotResult).toHaveBeenCalledWith(expect.objectContaining({ + format, context: expect.objectContaining({ connectionId: "original-connection", clientSessionId: "original-session" }), + cells: [{ rowIndex: 0, columnIndex: 1, valueRef: "original-locator" }], + }), expect.any(String)); + expect(readLargeValueChunk).not.toHaveBeenCalled(); + expect(exportQueryResultCsv).not.toHaveBeenCalled(); + expect(exportQueryResultJson).not.toHaveBeenCalled(); + }); + it("does not publish a legacy preview export when the snapshot backend fails", async () => { + const { state, original } = snapshotState(); + vi.mocked(exportSnapshotResult).mockRejectedValue(new Error("LOB snapshot expired")); + await state.exportJson(); + expect(exportQueryResultJson).not.toHaveBeenCalled(); + expect(original.rows[0]![1]).toBe("preview"); + expect(toast).toHaveBeenCalledWith(expect.stringContaining("LOB snapshot expired"), 5000); + }); + it("materializes export-all for bounded formats rather than passing preview rows", async () => { + const { original } = snapshotState(); + const localState = createExportState(editableTable, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined, "oracle", undefined, { + context: computed(() => "results"), snapshotResult: computed(() => original), hasCompleteLocalResult: computed(() => true), completeLocalResult: computed(() => original), + }); + vi.mocked(readLargeValueChunk).mockResolvedValue({ status: "ok", data: "complete original text", next_offset: 22, eof: true, value_kind: "text" }); + await localState.exportMarkdown(); + expect(exportQueryResultMarkdown).toHaveBeenCalledWith(expect.any(String), original.columns, [[1, "complete original text"]]); + expect(original.rows[0]![1]).toBe("preview"); + }); +}); + const editableTable: DataGridTableMeta = { tableName: "users", primaryKeys: ["id"], diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index abc9068fb2..d75f6ce776 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -1,4 +1,4 @@ -import { computed, type ComputedRef, type Ref, createApp } from "vue"; +import { computed, type ComputedRef, type Ref, createApp, watch, getCurrentScope, onScopeDispose } from "vue"; import { useI18n } from "vue-i18n"; import { useDataGridExtractor } from "@/composables/useDataGridExtractor"; import { isTauriRuntime } from "@/lib/backend/tauriRuntime"; @@ -14,6 +14,7 @@ import { useExportTracker } from "@/composables/useExportTracker"; import { clipboardCellValue, type CellValue } from "@/lib/dataGrid/cellValue"; import { binaryCellClipboardText } from "@/lib/dataGrid/binaryCellDownload"; import { materializeSnapshotResultRows } from "@/lib/dataGrid/largeValueSnapshot"; +import { snapshotExportSelection, assertSnapshotXlsxCellLengths } from "@/lib/dataGrid/snapshotExport"; import { tryStartExclusiveActivation, type ActionActivationGuard } from "@/lib/connection/actionActivation"; import { clipboardLineEndings, copyToClipboard } from "@/lib/common/clipboard"; import { clearDataGridClipboardCopy, rememberDataGridClipboardCopy } from "@/lib/dataGrid/dataGridClipboard"; @@ -170,6 +171,7 @@ export interface UseDataGridExportOptions { * silently change what "export all data" produces. */ completeLocalResult?: ComputedRef; + snapshotResult?: ComputedRef; allExportResults?: ComputedRef | undefined>; currentResultLabel?: ComputedRef; exportFileBaseName?: ComputedRef; @@ -198,6 +200,8 @@ interface CopyInsertData { } export function useDataGridExport(options: UseDataGridExportOptions) { + let activeSnapshotCancel: (() => Promise) | undefined; + if (getCurrentScope()) onScopeDispose(() => { void activeSnapshotCancel?.().catch(() => undefined); }); const { t } = useI18n(); const { toast } = useToast(); const tracker = useExportTracker(); @@ -524,7 +528,10 @@ export function useDataGridExport(options: UseDataGridExportOptions) { // client-side filters/search and unsaved edits, which would silently // change what the export contains. if (useFullExport && rowIds === undefined && hasCompleteLocalResult?.value && completeLocalResult?.value) { - const normalized = hasColumnSubset || databaseType.value === "mongodb" ? normalizeCompleteLocalResult(completeLocalResult.value, targetColumns, projectionSourceIndexes) : normalizeCompleteLocalResult(completeLocalResult.value); + const original = completeLocalResult.value; + const rows = await materializeSnapshotResultRows(original, () => completeLocalResult.value === original); + const materialized = { ...original, rows }; + const normalized = hasColumnSubset || databaseType.value === "mongodb" ? normalizeCompleteLocalResult(materialized, targetColumns, projectionSourceIndexes) : normalizeCompleteLocalResult(materialized); const columnComments = buildXlsxHeaderOverrides(normalized.columns, normalized.columnComments, headerMode); return { ...applyGlobalDateTimeExportFormat( @@ -590,6 +597,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { } async function writeXlsxResult(outputPath: string, result: { columns: string[]; columnTypes: string[]; columnComments?: (string | null)[]; rows: CellValue[][] }, includeSqlSheet: boolean, autoFilter: boolean, sqlOverride?: string) { + if (options.snapshotResult?.value.large_value_refs?.length || options.snapshotResult?.value.large_value_cells?.some((cell) => cell.value_ref)) assertSnapshotXlsxCellLengths(result.rows); const effectiveSql = sqlOverride ?? currentExportSql(); const sqlWorksheet = includeSqlSheet ? buildXlsxSqlWorksheet([{ sql: effectiveSql || "" }]) : undefined; const rightAlign = useSettingsStore().editorSettings.numericColumnRightAlign; @@ -930,10 +938,76 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return exported; } + async function exportLocalSnapshot(format: "csv" | "json", rowIds: number[] | undefined, columnIndexes: number[] | undefined, full: boolean): Promise { + const original = options.snapshotResult?.value; + const originalConnectionId = connectionId.value; + const originalDatabase = database.value; + const isCurrent = () => options.snapshotResult?.value === original && connectionId.value === originalConnectionId && database.value === originalDatabase; + if (!original || full && rowIds === undefined && !hasCompleteLocalResult?.value) return false; + const exportAll = full && rowIds === undefined && hasCompleteLocalResult?.value; + const visibleIndexes = visibleColumnIndexesOption?.value ?? columns.value.map((_, index) => index); + const indexes = columnIndexes?.filter((index) => index >= 0 && index < columns.value.length); + const sourceIndexes = exportAll && !indexes ? original.columns.map((_, index) => index) : (indexes ?? columns.value.map((_, index) => index)).map((index) => visibleIndexes[index] ?? index); + const names = sourceIndexes.map((index) => original.columns[index]!); + const items = exportAll + ? original.rows.map((row, sourceIndex) => ({ sourceIndex, data: sourceIndexes.map((index) => row[index]) })) + : rowsToExport(rowIds).map((item) => ({ + sourceIndex: item.sourceIndex, isNew: item.isNew, + data: externalizeRows([item.data], indexes)[0]!, + isDirtyCol: (indexes ?? columns.value.map((_, index) => index)).map((index) => item.isDirtyCol[index] ?? false), + })); + const selected = snapshotExportSelection(original, names, sourceIndexes, items); + if (!selected) return false; + const types = sourceIndexes.map((index) => original.column_types?.[index] ?? ""); + selected.rows = applyGlobalDateTimeExportFormat({ columns: names, columnTypes: types, rows: selected.rows as CellValue[][] }, true).rows; + let outputPath = exportFileName(full ? "export" : "export-page", format, { page: !full }); + if (isTauriRuntime()) { + const path = await promptExportSavePath({ defaultFileName: outputPath, filters: [{ name: FORMAT_META[format]!.label, extensions: [format] }], preferredPath: useSettingsStore().editorSettings.preferredExportPath }); + if (!path) return true; + outputPath = path; + } + if (!isCurrent()) throw new Error("LOB result context changed"); + const executionId = uuid(); + const cancel = async () => { await api.cancelQuery(executionId); }; + activeSnapshotCancel = cancel; + const title = currentExportTitle(); + tracker.addTask(title, format, outputPath, executionId); + tracker.registerTaskCancelHandler(executionId, cancel); + if (exportProgressState) exportProgressState.value = { + title: t("exportProgress.title"), tableName: title, format, rowsExported: 0, totalRows: selected.rows.length, + status: "Running", errorMessage: null, filePath: outputPath, startedAt: Date.now(), finishedAt: undefined, + }; + if (exportProgressDialog) exportProgressDialog.value = true; + if (exportCanMinimize) exportCanMinimize.value = true; + const previousCancel = exportCancelHandler?.value; + if (exportCancelHandler) exportCancelHandler.value = cancel; + const stop = watch(() => [options.snapshotResult?.value, connectionId.value, database.value], () => { if (!isCurrent()) void cancel().catch(() => undefined); }); + try { + await api.exportSnapshotResult({ ...selected, context: { ...selected.context, executionId }, format, quoteMode: useSettingsStore().editorSettings.csvQuoteMode, nullLiteral: csvNullLiteralForMode(useSettingsStore().editorSettings.csvNullMode) }, outputPath); + if (!isCurrent()) throw new Error("LOB result context changed"); + tracker.updateTableExportTask(executionId, { exportId: executionId, tableName: title, rowsExported: selected.rows.length, totalRows: selected.rows.length, status: "Done" }); + if (exportProgressState) exportProgressState.value = { ...exportProgressState.value, status: "Done", rowsExported: selected.rows.length, finishedAt: Date.now() }; + notifyExportSuccess(outputPath); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + tracker.updateTableExportTask(executionId, { exportId: executionId, tableName: title, rowsExported: 0, totalRows: selected.rows.length, status: "Error", errorMessage }); + if (exportProgressState) exportProgressState.value = { ...exportProgressState.value, status: "Error", errorMessage, finishedAt: Date.now() }; + throw error; + } finally { + stop(); + if (activeSnapshotCancel === cancel) activeSnapshotCancel = undefined; + tracker.unregisterTaskCancelHandler(executionId); + if (exportCanMinimize) exportCanMinimize.value = false; + if (exportCancelHandler) exportCancelHandler.value = previousCancel ?? null; + } + return true; + } + async function exportCsv(target?: ExportTargetParam, columnIndexesParam?: number[]) { const { rowIds, columnIndexes } = normalizeExportTarget(target, columnIndexesParam); await runExclusiveExport(async () => { try { + if (await exportLocalSnapshot("csv", rowIds, columnIndexes, true)) return; if (await exportQueryResultViaBackend("csv", rowIds, false, "name", true, undefined, false, columnIndexes)) return; if (await exportFullTableDataViaBackend("csv", rowIds, "name", true, undefined, false, columnIndexes)) return; @@ -1016,6 +1090,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportCurrentPageCsv(columnIndexes?: number[]) { await runExclusiveExport(async () => { try { + if (await exportLocalSnapshot("csv", undefined, columnIndexes, false)) return; let outputPath = exportFileName("export-page", "csv", { page: true }); if (isTauriRuntime()) { const path = await promptExportSavePath({ @@ -1039,6 +1114,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { const { rowIds, columnIndexes } = normalizeExportTarget(target, columnIndexesParam); await runExclusiveExport(async () => { try { + if (await exportLocalSnapshot("json", rowIds, columnIndexes, true)) return; if (await exportFullTableDataViaBackend("json", rowIds, "name", true, undefined, false, columnIndexes)) return; if (await exportQueryResultViaBackend("json", rowIds, false, "name", true, undefined, false, columnIndexes)) return; @@ -1064,6 +1140,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportCurrentPageJson(columnIndexes?: number[]) { await runExclusiveExport(async () => { try { + if (await exportLocalSnapshot("json", undefined, columnIndexes, false)) return; let outputPath = exportFileName("export-page", "json", { page: true }); if (isTauriRuntime()) { const path = await promptExportSavePath({ @@ -1375,6 +1452,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { for (const sheet of sheets) { const sourceRows = sheet.result.rows; const rows = await materializeSnapshotResultRows(sheet.result, () => sheet.result.rows === sourceRows); + if (sheet.result.large_value_refs?.length || sheet.result.large_value_cells?.some((cell) => cell.value_ref)) assertSnapshotXlsxCellLengths(rows); worksheets.push({ sheetName: sheet.sheetName, columns: sheet.result.columns, diff --git a/apps/desktop/src/composables/useDataGridLargeValues.ts b/apps/desktop/src/composables/useDataGridLargeValues.ts index 0363f1c8df..07f6680e9a 100644 --- a/apps/desktop/src/composables/useDataGridLargeValues.ts +++ b/apps/desktop/src/composables/useDataGridLargeValues.ts @@ -80,6 +80,7 @@ export interface UseDataGridLargeValuesOptions { } type ResolvedLargeValueCells = Map>; +const materializedSnapshotBytes = new WeakMap>(); type LargeValueCellRequest = { item: LargeValueRowItem; sourceIndex: number; @@ -513,6 +514,7 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { const isCurrent = () => options.result.value === sourceResult && options.connectionId.value === connectionId && options.executionDatabase.value === executionDatabase && options.resultLifecycle.isCurrent(operation); const resolved: ResolvedLargeValueCells = new Map(); + let snapshotBytes = 0; const requestedColumns = new Set(columnIndexes); const requestsByColumn = new Map(); for (const rowId of new Set(rowIds)) { @@ -531,6 +533,8 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { snapshotExecutionIds.add(executionId); try { const value = await materializeLargeValueSnapshot({ ...context, valueRef: metadata.value_ref, executionId }, isCurrent); + snapshotBytes += new TextEncoder().encode(value).length; + if (snapshotBytes > 64 * 1024 * 1024) throw new Error("LOB selection exceeds the 64 MiB view/copy limit; use CSV/JSON export or download individual complete values"); const rowValues = resolved.get(item.id) ?? new Map(); rowValues.set(columnIndex, value); resolved.set(item.id, rowValues); @@ -598,6 +602,7 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { if (!isLargeValuePreview(item, columnIndex) || item?.sourceIndex === undefined) return true; const sourceResult = options.result.value; const hydrationKey = largeValueCellKey(item.sourceIndex, columnIndex); + const ownsSnapshot = !!largeValueCellsByKey.value.get(hydrationKey)?.value_ref; const operation = options.resultLifecycle.beginOperation(); return pendingLargeValueHydrations.run(hydrationKey, sourceResult, async () => { try { @@ -605,6 +610,14 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { if (!options.resultLifecycle.isCurrent(operation) || options.result.value !== sourceResult) return false; const value = resolved.get(rowId)?.get(columnIndex); if (value === undefined && !resolved.get(rowId)?.has(columnIndex)) return false; + if (ownsSnapshot && typeof value === "string") { + const sizes = materializedSnapshotBytes.get(sourceResult) ?? new Map(); + const bytes = new TextEncoder().encode(value).length; + const total = [...sizes.values()].reduce((sum, size) => sum + size, 0) - (sizes.get(hydrationKey) ?? 0) + bytes; + if (total > 64 * 1024 * 1024) throw new Error("LOB result exceeds the 64 MiB view/edit limit; use CSV/JSON export or download individual complete values"); + sizes.set(hydrationKey, bytes); + materializedSnapshotBytes.set(sourceResult, sizes); + } const row = [...(sourceResult.rows[item.sourceIndex!] ?? [])]; row[columnIndex] = value ?? null; const rows = sourceResult.rows.slice(); diff --git a/apps/desktop/src/lib/__tests__/dataGrid/snapshotExport.spec.ts b/apps/desktop/src/lib/__tests__/dataGrid/snapshotExport.spec.ts new file mode 100644 index 0000000000..ffc4f560bb --- /dev/null +++ b/apps/desktop/src/lib/__tests__/dataGrid/snapshotExport.spec.ts @@ -0,0 +1,47 @@ +import { describe, expect, it } from "vitest"; +import { snapshotExportSelection, assertSnapshotXlsxCellLengths } from "@/lib/dataGrid/snapshotExport"; +import type { QueryResult } from "@/types/database"; + +const result = { + columns: ["ID", "Value", "value"], rows: [[1, "first preview", "second preview"], [2, "third preview", null]], + large_value_context: { connectionId: "original", database: "oracletest", clientSessionId: "original-session" }, + large_value_cells: [ + { row_index: 0, column_index: 1, value_ref: "first", value_kind: "text" }, + { row_index: 0, column_index: 2, value_ref: "second", value_kind: "text" }, + { row_index: 1, column_index: 1, value_ref: "third", value_kind: "text" }, + ], +} as QueryResult; + +describe("snapshot export selection", () => { + it("refuses XLSX truncation while retaining the complete value for other formats", () => { + expect(() => assertSnapshotXlsxCellLengths([["a".repeat(32767)]])).not.toThrow(); + const value = "δΈ­πŸ™‚".repeat(20000); + expect(() => assertSnapshotXlsxCellLengths([[value]])).toThrow("XLSX cannot preserve"); + expect(Array.from(value)).toHaveLength(40000); + }); + it("remaps sorted rows and projected quoted columns by source positions", () => { + const selected = snapshotExportSelection(result, ["value", "Value"], [2, 1], [ + { sourceIndex: 1, data: [null, "third preview"] }, + { sourceIndex: 0, data: ["second preview", "first preview"] }, + ])!; + expect(selected.cells).toEqual([ + { rowIndex: 0, columnIndex: 1, valueRef: "third" }, + { rowIndex: 1, columnIndex: 0, valueRef: "second" }, + { rowIndex: 1, columnIndex: 1, valueRef: "first" }, + ]); + expect(selected.context).toEqual({ ...result.large_value_context, valueRef: "" }); + }); + it("preserves complete dirty values and new rows rather than substituting their original snapshots", () => { + const selected = snapshotExportSelection(result, ["Value", "value"], [1, 2], [ + { sourceIndex: 0, data: ["edited complete", "second preview"], isDirtyCol: [true, false] }, + { isNew: true, data: ["new complete", "other complete"] }, + ])!; + expect(selected.cells).toEqual([{ rowIndex: 0, columnIndex: 1, valueRef: "second" }]); + expect(selected.rows[0]![0]).toBe("edited complete"); + expect(result.rows[0]![1]).toBe("first preview"); + }); + it("fails before export when original connection context is missing", () => { + expect(() => snapshotExportSelection({ ...result, large_value_context: undefined }, ["Value"], [1], [{ sourceIndex: 0, data: ["preview"] }])).toThrow("connection is unavailable"); + expect(snapshotExportSelection(result, ["ID"], [0], [{ sourceIndex: 0, data: [1] }])).toBeUndefined(); + }); +}); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index fcab597f7f..80395f4ea3 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -284,6 +284,7 @@ export const closeQuerySession = forward("closeQuerySession"); export const readLargeValueChunk = forward("readLargeValueChunk"); export const releaseLargeValue = forward("releaseLargeValue"); export const downloadLargeValue = forward("downloadLargeValue"); +export const exportSnapshotResult = forward("exportSnapshotResult"); export const closeClientConnectionSession = forward("closeClientConnectionSession"); export const analyzeSqlReferences = forward("analyzeSqlReferences"); export const findStatementAtCursor = forward("findStatementAtCursor"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 501b712c63..a3c523f57a 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -1624,6 +1624,25 @@ export async function downloadLargeValue(request: LargeValueRequest, filePath: s anchor.click(); } +export interface SnapshotExportRequest { + context: LargeValueRequest; + format: "csv" | "json"; + columns: string[]; + rows: unknown[][]; + cells: Array<{ rowIndex: number; columnIndex: number; valueRef: string }>; + quoteMode?: "all" | "necessary"; + nullLiteral?: string | null; +} + +export async function exportSnapshotResult(request: SnapshotExportRequest, filePath: string): Promise { + const fileName = filePath.split(/[\\/]/).pop() || `result.${request.format}`; + const prepared = await post<{ downloadId: string }>("/api/query/large-value/export", { request, fileName }); + const anchor = document.createElement("a"); + anchor.href = apiUrl(`/api/query/large-value/download/${encodeURIComponent(prepared.downloadId)}`); + anchor.download = fileName; + anchor.click(); +} + export async function closeQuerySession(connectionId: string, database: string, sessionId: string, clientSessionId?: string, catalog?: string): Promise { return post("/api/query/close-session", { connectionId, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 3476cc0df2..9614e8ff1d 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -2036,6 +2036,10 @@ export async function downloadLargeValue(request: import("./http").LargeValueReq await invoke("download_large_value", { request, filePath }); } +export async function exportSnapshotResult(request: import("./http").SnapshotExportRequest, filePath: string): Promise { + await invoke("export_snapshot_result", { request, filePath }); +} + export async function closeQuerySession(connectionId: string, database: string, sessionId: string, clientSessionId?: string, catalog?: string): Promise { return invoke("close_query_session", { connectionId, diff --git a/apps/desktop/src/lib/dataGrid/snapshotExport.ts b/apps/desktop/src/lib/dataGrid/snapshotExport.ts new file mode 100644 index 0000000000..9fb8f19269 --- /dev/null +++ b/apps/desktop/src/lib/dataGrid/snapshotExport.ts @@ -0,0 +1,41 @@ +import type { QueryResult } from "@/types/database"; +import type { SnapshotExportRequest } from "@/lib/backend/http"; + +interface SnapshotExportRow { + sourceIndex?: number; + data: unknown[]; + isNew?: boolean; + isDirtyCol?: boolean[]; +} + +/** Projection uses source positions, never labels, so quoted and repeated names remain distinct. */ +export function snapshotExportSelection( + result: QueryResult, + columns: string[], + sourceColumns: number[], + items: SnapshotExportRow[], +): Pick | undefined { + const refs = new Map((result.large_value_cells ?? []).filter((cell) => cell.value_ref).map((cell) => [`${cell.row_index}:${cell.column_index}`, cell.value_ref!])); + const cells: SnapshotExportRequest["cells"] = []; + items.forEach((item, rowIndex) => { + if (item.isNew || item.sourceIndex === undefined) return; + sourceColumns.forEach((sourceColumn, columnIndex) => { + if (item.isDirtyCol?.[columnIndex]) return; + const valueRef = refs.get(`${item.sourceIndex}:${sourceColumn}`); + if (valueRef) cells.push({ rowIndex, columnIndex, valueRef }); + }); + }); + if (!cells.length) return undefined; + if (!result.large_value_context) throw new Error("LOB result connection is unavailable; execute the query again"); + return { columns, rows: items.map((item) => item.data), cells, context: { ...result.large_value_context, valueRef: "" } }; +} + +export function assertSnapshotXlsxCellLengths(rows: unknown[][]): void { + for (const row of rows) for (const value of row) { + if (typeof value !== "string") continue; + let count = 0; + for (const _character of value) if (++count > 32767) { + throw new Error("XLSX cannot preserve a LOB longer than 32,767 characters; use CSV/JSON or download the complete value"); + } + } +} diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index 189c0189b4..2f0f9288ce 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -7,6 +7,7 @@ pub mod plugin_data; pub mod plugin_plan; pub mod query_cancel; pub mod redis_ops; +pub mod snapshot_export; mod sqlserver_agent_batch; pub mod sqlserver_manual_transaction; pub mod two_phase_commit; diff --git a/crates/dbx-core/src/query/snapshot_export.rs b/crates/dbx-core/src/query/snapshot_export.rs new file mode 100644 index 0000000000..bec2543017 --- /dev/null +++ b/crates/dbx-core/src/query/snapshot_export.rs @@ -0,0 +1,321 @@ +use super::{is_canceled, canceled_error, request_large_value_with_cancel, LargeValueRequest}; +use crate::connection::AppState; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; +use std::io::{BufReader, Read, Seek, SeekFrom, Write}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SnapshotExportCell { + pub row_index: usize, + pub column_index: usize, + pub value_ref: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SnapshotExportRequest { + pub context: LargeValueRequest, + pub format: String, + pub columns: Vec, + pub rows: Vec>, + pub cells: Vec, + #[serde(default)] + pub quote_mode: dbx_formats::csv_export::CsvQuoteMode, + pub null_literal: Option, +} + +/// Original result references only. A temporary cell file bounds memory even for +/// formula-guard inspection; callers publish the output only after success. +pub async fn write_snapshot_export(state: &AppState, request: SnapshotExportRequest, output: &mut std::fs::File) -> Result<(), String> { + let registered = request.context.execution_id.as_ref().map(|id| state.running_queries.register_task( + id.clone(), crate::query_cancel::RunningTaskMetadata::query( + request.context.connection_id.clone(), request.context.database.clone(), request.context.client_session_id.clone(), + ), + )); + let cancel = registered.as_ref().map(|task| task.token()); + write_snapshot_export_with_fetch(request, output, cancel.clone(), |chunk_request| { + request_large_value_with_cancel(state, chunk_request, false, cancel.clone()) + }).await +} + +async fn write_snapshot_export_with_fetch(request: SnapshotExportRequest, output: &mut std::fs::File, cancel: Option, mut fetch: F) -> Result<(), String> +where F: FnMut(LargeValueRequest) -> Fut, Fut: std::future::Future> { + let mut buffered = std::io::BufWriter::new(output); + let output = &mut buffered; + if request.format != "csv" && request.format != "json" { return Err("Unsupported snapshot export format".into()); } + if request.rows.iter().any(|row| row.len() != request.columns.len()) { return Err("Invalid snapshot export row".into()); } + let mut refs = HashMap::new(); + for cell in &request.cells { + if cell.row_index >= request.rows.len() || cell.column_index >= request.columns.len() + || cell.value_ref.is_empty() || cell.value_ref.len() > 128 + || refs.insert((cell.row_index, cell.column_index), cell.value_ref.as_str()).is_some() { + return Err("Invalid snapshot export cell".into()); + } + } + // JSON objects cannot preserve duplicate labels. Reject rather than dropping a column. + if request.format == "json" { + let mut names = std::collections::HashSet::new(); + if request.columns.iter().any(|name| !names.insert(name)) { return Err("JSON export requires unique column names; use CSV".into()); } + } + let json = request.format == "json"; + if json { output.write_all(b"[").map_err(|e| e.to_string())?; } + else { + let mut header = String::new(); + for (index, name) in request.columns.iter().enumerate() { + if index > 0 { header.push(','); } + dbx_formats::csv_export::push_csv_field(&mut header, name, request.quote_mode); + } + output.write_all(header.as_bytes()).map_err(|e| e.to_string())?; + } + for (row_index, row) in request.rows.iter().enumerate() { + if is_canceled(&cancel) { return Err(canceled_error()); } + output.write_all(if json { if row_index == 0 { b"{" } else { b",{" } } else { b"\n" }).map_err(|e| e.to_string())?; + for (column_index, value) in row.iter().enumerate() { + if column_index > 0 { output.write_all(b",").map_err(|e| e.to_string())?; } + if json { + serde_json::to_writer(&mut *output, &request.columns[column_index]).map_err(|e| e.to_string())?; + output.write_all(b":").map_err(|e| e.to_string())?; + } + if let Some(value_ref) = refs.get(&(row_index, column_index)) { + let mut cell_file = tempfile::tempfile().map_err(|e| e.to_string())?; + let mut chunk_request = request.context.clone(); + chunk_request.value_ref = (*value_ref).to_string(); + chunk_request.offset = 0; + chunk_request.limit = 4096; + let mut original_kind: Option = None; + loop { + if is_canceled(&cancel) { return Err(canceled_error()); } + let chunk = fetch(chunk_request.clone()).await?; + let (data, next, eof, kind) = checked_chunk(&chunk, chunk_request.offset)?; + if original_kind.as_deref().is_some_and(|original| original != kind) { return Err("LOB chunk type changed".into()); } + if original_kind.is_none() && kind == "binary" { cell_file.write_all(b"0x").map_err(|e| e.to_string())?; } + original_kind = Some(kind.to_string()); + cell_file.write_all(data.as_bytes()).map_err(|e| e.to_string())?; + if eof { break; } + chunk_request.offset = next; + } + cell_file.seek(SeekFrom::Start(0)).map_err(|e| e.to_string())?; + let guard = !json && stream_formula_guard(BufReader::new(CancelableReader { reader: &mut cell_file, cancel: &cancel }))?; + cell_file.seek(SeekFrom::Start(0)).map_err(|e| e.to_string())?; + let quoted = json || request.quote_mode == dbx_formats::csv_export::CsvQuoteMode::All || csv_needs_quotes(CancelableReader { reader: &mut cell_file, cancel: &cancel })?; + cell_file.seek(SeekFrom::Start(0)).map_err(|e| e.to_string())?; + if quoted { output.write_all(b"\"").map_err(|e| e.to_string())?; } + if guard { output.write_all(b"'").map_err(|e| e.to_string())?; } + let mut buffer = [0_u8; 8192]; + loop { + if is_canceled(&cancel) { return Err(canceled_error()); } + let count = cell_file.read(&mut buffer).map_err(|e| e.to_string())?; + if count == 0 { break; } + write_escaped(output, &buffer[..count], json)?; + } + if quoted { output.write_all(b"\"").map_err(|e| e.to_string())?; } + } else if json { serde_json::to_writer(&mut *output, value).map_err(|e| e.to_string())?; } + else { + let mut cell = String::new(); + dbx_formats::csv_export::push_query_result_csv_row_with_options(&mut cell, std::slice::from_ref(value), request.quote_mode, request.null_literal.as_deref()); + output.write_all(cell.as_bytes()).map_err(|e| e.to_string())?; + } + } + if json { output.write_all(b"}").map_err(|e| e.to_string())?; } + } + if is_canceled(&cancel) { return Err(canceled_error()); } + if json { output.write_all(b"]").map_err(|e| e.to_string())?; } + output.flush().map_err(|e| e.to_string()) +} + +struct CancelableReader<'a, R> { reader: R, cancel: &'a Option } +impl Read for CancelableReader<'_, R> { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + if is_canceled(self.cancel) { return Err(std::io::Error::other(canceled_error())); } + self.reader.read(buffer) + } +} + +fn checked_chunk(chunk: &Value, offset: u64) -> Result<(&str, u64, bool, &str), String> { + if chunk.get("status").and_then(Value::as_str) != Some("ok") { return Err("LOB snapshot expired; execute the query again".into()); } + let data = chunk.get("data").and_then(Value::as_str).ok_or("Invalid LOB data")?; + let next = chunk.get("next_offset").and_then(Value::as_u64).ok_or("Invalid LOB offset")?; + let eof = chunk.get("eof").and_then(Value::as_bool).ok_or("Invalid LOB EOF")?; + let kind = chunk.get("value_kind").and_then(Value::as_str).ok_or("Invalid LOB type")?; + if next < offset || next - offset > 4096 || (!eof && next == offset) { return Err("Invalid LOB offset".into()); } + let amount = match kind { + "text" => data.chars().count(), + "binary" if data.len() <= 8192 && data.len() % 2 == 0 && data.bytes().all(|b| b.is_ascii_hexdigit()) => data.len() / 2, + _ => return Err("Invalid LOB chunk encoding".into()), + }; + if amount as u64 != next - offset { return Err("Invalid LOB chunk length".into()); } + Ok((data, next, eof, kind)) +} + +fn write_escaped(output: &mut impl Write, bytes: &[u8], json: bool) -> Result<(), String> { + for &byte in bytes { + match byte { + b'"' => output.write_all(if json { b"\\\"" } else { b"\"\"" }), + b'\\' if json => output.write_all(b"\\\\"), + 0..=31 if json => write!(output, "\\u{byte:04x}"), + _ => output.write_all(&[byte]), + }.map_err(|e| e.to_string())?; + } + Ok(()) +} + +fn csv_needs_quotes(mut reader: impl Read) -> Result { + let mut buffer = [0_u8; 8192]; + loop { + let count = reader.read(&mut buffer).map_err(|e| e.to_string())?; + if count == 0 { return Ok(false); } + if buffer[..count].iter().any(|&byte| matches!(byte, b',' | b'"' | b'\n' | b'\r')) { return Ok(true); } + } +} + +// Same cell-level guard as dbx-formats, including negative decimals and literal +// apostrophes, without retaining leading spaces/zeroes or the complete cell. +fn stream_formula_guard(reader: impl Read) -> Result { + let mut bytes = reader.bytes(); + let mut next = || bytes.next().transpose().map_err(|e| e.to_string()); + let mut byte = next()?; + if byte == Some(b'\'') { + byte = next()?; + if byte == Some(b'\'') { return Ok(true); } + } + while byte == Some(b' ') { byte = next()?; } + if byte != Some(b'-') { return Ok(byte.is_some_and(|b| b"=+@\t\r".contains(&b))); } + byte = next()?; + if matches!(byte, Some(b'+') | Some(b'-')) { byte = next()?; } + let mut digits = 0_usize; + let mut significant = 0; + let mut dot = false; + loop { + match byte { + Some(b'0'..=b'9') => { + digits += 1; + if byte != Some(b'0') || significant > 0 { significant += 1; } + if significant > 15 { return Ok(true); } + } + Some(b'.') if !dot && digits > 0 => { dot = true; digits = 0; significant = 0; } + None => return Ok(digits == 0), + _ => return Ok(true), + } + byte = next()?; + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn streamed_guard_matches_existing_cell_semantics() { + for value in ["", "plain", "=cmd", " +86", "-1", "-1.25", "-00000000000000000000000000001", "-1.00000000000000001", "-2+cmd", "-", "'-1", "'=cmd", "''text", "'plain", " --1", "-.1", "\ttext"] { + let mut expected = String::new(); + dbx_formats::csv_export::push_formula_guard(&mut expected, value); + assert_eq!(stream_formula_guard(value.as_bytes()).unwrap(), !expected.is_empty(), "{value}"); + } + } + #[test] + fn escapes_chunk_boundaries_without_changing_utf8_or_binary_hex() { + let text = "δΈ­πŸ™‚\"\\\n\0"; + for split in 0..=text.len() { + let mut out = Vec::new(); + write_escaped(&mut out, &text.as_bytes()[..split], true).unwrap(); + write_escaped(&mut out, &text.as_bytes()[split..], true).unwrap(); + assert_eq!(serde_json::from_slice::(&[b"\"".as_slice(), &out, b"\""].concat()).unwrap(), text); + } + } + #[test] + fn rejects_expiry_invalid_offsets_and_incomplete_hex() { + for chunk in [serde_json::json!({"status":"expired"}), serde_json::json!({"status":"ok","data":"0","next_offset":1,"eof":true,"value_kind":"binary"}), serde_json::json!({"status":"ok","data":"a","next_offset":5000,"eof":true,"value_kind":"text"})] { + assert!(checked_chunk(&chunk, 0).is_err()); + } + let chunk = serde_json::json!({"status":"ok","data":"πŸ™‚","next_offset":1,"eof":true,"value_kind":"text"}); + assert_eq!(checked_chunk(&chunk, 0).unwrap().1, 1); + } + + fn request(format: &str) -> SnapshotExportRequest { + serde_json::from_value(serde_json::json!({ + "context": {"connectionId":"original-connection","database":"oracletest","valueRef":"","clientSessionId":"original-session"}, + "format":format,"columns":["ID","PAYLOAD"],"rows":[[1,"preview"]], + "cells":[{"rowIndex":0,"columnIndex":1,"valueRef":"original-locator"}],"nullLiteral":"\\N" + })).unwrap() + } + + #[tokio::test] + async fn production_writer_preserves_unicode_quotes_and_original_context_in_multiple_chunks() { + let original = format!(" =cmd{}\"\\\n", "δΈ­πŸ™‚".repeat(3000)); + let chars: Vec = original.chars().collect(); + for format in ["csv", "json"] { + let mut file = tempfile::tempfile().unwrap(); + let mut offsets = Vec::new(); + write_snapshot_export_with_fetch(request(format), &mut file, None, |chunk_request| { + assert_eq!(chunk_request.value_ref, "original-locator"); + assert_eq!(chunk_request.connection_id, "original-connection"); + assert_eq!(chunk_request.client_session_id.as_deref(), Some("original-session")); + assert_eq!(chunk_request.limit, 4096); + offsets.push(chunk_request.offset); + let begin = chunk_request.offset as usize; + let end = (begin + 4096).min(chars.len()); + std::future::ready(Ok(serde_json::json!({"status":"ok","data":chars[begin..end].iter().collect::(),"next_offset":end,"eof":end==chars.len(),"value_kind":"text"}))) + }).await.unwrap(); + assert_eq!(offsets, vec![0,4096]); + file.seek(SeekFrom::Start(0)).unwrap(); + let mut text = String::new(); file.read_to_string(&mut text).unwrap(); + if format == "json" { assert_eq!(serde_json::from_str::(&text).unwrap()[0]["PAYLOAD"], original); } + else { assert_eq!(text, dbx_formats::csv_export::format_query_result_csv_with_options(&["ID".into(),"PAYLOAD".into()], &[vec![Value::from(1), Value::from(original.clone())]], dbx_formats::csv_export::CsvQuoteMode::All, Some("\\N"))); } + } + } + + #[tokio::test] + async fn production_writer_never_completes_on_mid_read_expiry_or_cancel() { + for cancel_read in [false, true] { + let token = tokio_util::sync::CancellationToken::new(); + let mut file = tempfile::tempfile().unwrap(); + let mut calls = 0; + let error = write_snapshot_export_with_fetch(request("json"), &mut file, Some(token.clone()), |chunk_request| { + calls += 1; + if cancel_read { token.cancel(); } + std::future::ready(Ok(if chunk_request.offset == 0 { + serde_json::json!({"status":"ok","data":"first","next_offset":5,"eof":false,"value_kind":"text"}) + } else { serde_json::json!({"status":"expired"}) })) + }).await.unwrap_err(); + assert_eq!(calls, if cancel_read {1} else {2}); + assert!(!error.is_empty()); + } + } + + #[tokio::test] + async fn production_writer_keeps_blob_hex_prefix_once_and_all_byte_values() { + let hex: String = (0..=255_u8).map(|byte| format!("{byte:02x}")).collect(); + let mut file = tempfile::tempfile().unwrap(); + write_snapshot_export_with_fetch(request("json"), &mut file, None, |chunk_request| { + let begin = chunk_request.offset as usize; + let end = (begin + 127).min(256); + std::future::ready(Ok(serde_json::json!({"status":"ok","data":&hex[begin*2..end*2],"next_offset":end,"eof":end==256,"value_kind":"binary"}))) + }).await.unwrap(); + file.seek(SeekFrom::Start(0)).unwrap(); + let result: Value = serde_json::from_reader(file).unwrap(); + assert_eq!(result[0]["PAYLOAD"], format!("0x{hex}")); + } + + #[tokio::test] + async fn necessary_quotes_and_formula_guard_match_existing_export() { + for value in ["plain", "", "-000000000000000000001", " =cmd", "'plain", "'=cmd", "a,b", "a\"b"] { + let mut req = request("csv"); + req.quote_mode = dbx_formats::csv_export::CsvQuoteMode::Necessary; + let mut file = tempfile::tempfile().unwrap(); + write_snapshot_export_with_fetch(req, &mut file, None, |_| std::future::ready(Ok(serde_json::json!({ + "status":"ok","data":value,"next_offset":value.chars().count(),"eof":true,"value_kind":"text" + })))).await.unwrap(); + file.seek(SeekFrom::Start(0)).unwrap(); + let mut output = String::new(); file.read_to_string(&mut output).unwrap(); + assert_eq!(output, dbx_formats::csv_export::format_query_result_csv_with_options(&["ID".into(),"PAYLOAD".into()], &[vec![Value::from(1), Value::from(value)]], dbx_formats::csv_export::CsvQuoteMode::Necessary, Some("\\N"))); + } + } + + #[test] + fn cancelled_formula_scan_stops_instead_of_retrying_forever() { + let token = tokio_util::sync::CancellationToken::new(); token.cancel(); + let cancel = Some(token); + assert!(stream_formula_guard(BufReader::new(CancelableReader { reader: b" -0000".as_slice(), cancel: &cancel })).is_err()); + } +} diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index bba528a353..aa8d6dd495 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -916,6 +916,7 @@ async fn serve() -> Result<(), String> { .route("/query/large-value/chunk", post(routes::query::read_large_value_chunk)) .route("/query/large-value/release", post(routes::query::release_large_value)) .route("/query/large-value/download", post(routes::query::prepare_large_value_download)) + .route("/query/large-value/export", post(routes::query::prepare_snapshot_export)) .route("/query/large-value/download/{id}", get(routes::query::download_large_value)) .route("/query/close-client-session", post(routes::query::close_client_connection_session)) .route("/export/query-result-json", post(routes::text_export::export_query_result_json)) diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index d5bef8581d..42bf24bbb5 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -732,6 +732,36 @@ pub async fn download_large_value( super::mongodb_import_export::export_file_response(file).await } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SnapshotExportDownloadRequest { + pub request: dbx_core::query::snapshot_export::SnapshotExportRequest, + pub file_name: String, +} + +pub async fn prepare_snapshot_export( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let mut temporary = tempfile::NamedTempFile::new().map_err(|error| AppError::from(error.to_string()))?; + let format = req.request.format.clone(); + dbx_core::query::snapshot_export::write_snapshot_export(&state.app, req.request, temporary.as_file_mut()).await.map_err(AppError::from)?; + let (_, path) = temporary.keep().map_err(|error| AppError::from(error.to_string()))?; + let download_id = uuid::Uuid::new_v4().to_string(); + state.export_files.write().await.insert(download_id.clone(), crate::state::WebExportFile { + file_path: path.to_string_lossy().into_owned(), download_filename: req.file_name, format, + }); + let cleanup_state = state.clone(); + let cleanup_id = download_id.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(300)).await; + if let Some(file) = cleanup_state.export_files.write().await.remove(&cleanup_id) { + let _ = tokio::fs::remove_file(file.file_path).await; + } + }); + Ok(Json(serde_json::json!({ "downloadId": download_id }))) +} + pub async fn close_query_session( State(state): State>, Json(req): Json, diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 45116e50a8..5365ce5a6b 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -376,6 +376,20 @@ pub async fn close_client_connection_session( state.close_client_session_pool(&connection_id, database, &client_session_id).await } +#[tauri::command] +pub async fn export_snapshot_result( + state: State<'_, Arc>, + request: dbx_core::query::snapshot_export::SnapshotExportRequest, + file_path: String, +) -> Result<(), String> { + let path = std::path::Path::new(&file_path); + let parent = path.parent().filter(|parent| !parent.as_os_str().is_empty()).unwrap_or(std::path::Path::new(".")); + let mut temporary = tempfile::NamedTempFile::new_in(parent).map_err(|error| error.to_string())?; + dbx_core::query::snapshot_export::write_snapshot_export(&state, request, temporary.as_file_mut()).await?; + temporary.persist(path).map_err(|error| error.to_string())?; + Ok(()) +} + fn query_session_database<'a>(database: &'a str, catalog: Option<&str>) -> Option<&'a str> { if database.trim().is_empty() || catalog.is_some() { None diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index b6a3c19401..21be43f45e 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2079,6 +2079,7 @@ pub fn run() { commands::query::read_large_value_chunk, commands::query::release_large_value, commands::query::download_large_value, + commands::query::export_snapshot_result, commands::query::close_client_connection_session, commands::query::execute_batch, commands::query::execute_script,