From b70d6da227b6e468eb63b9ee739c2bb21cfe6779 Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 14 Sep 2026 20:04:11 +0800 Subject: [PATCH] fix(export): stream JSON query exports and allow web TXT --- apps/desktop/src/components/grid/DataGrid.vue | 2 +- .../src/components/layout/ContentArea.vue | 2 +- .../useDataGridExport.sqlProgress.spec.ts | 8 ++ .../src/composables/useDataGridExport.ts | 5 +- apps/desktop/src/lib/backend/tauri.ts | 2 +- apps/desktop/src/stores/queryStore.ts | 2 +- crates/dbx-core/src/query_result_export.rs | 126 +++++++++++++++++- .../dbx-web/src/routes/query_result_export.rs | 4 + 8 files changed, 140 insertions(+), 11 deletions(-) diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index e0dcbfebc1..9e8a5d585f 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -513,7 +513,7 @@ interface DataGridProps { queryResultExportRequest?: (options: { exportId: string; filePath: string; - format: "csv" | "xlsx" | "txt" | "sql"; + format: "csv" | "xlsx" | "json" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 5dbef2f9b9..3ef362649f 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -1963,7 +1963,7 @@ defineExpose({ :on-execute-sql="async (sql: string) => emit('executeSql', activeTab.id, sql)" :full-export-result="(onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => queryStore.fetchTabResultForExport(activeTab.id, onProgress)" :query-result-export-request=" - (options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' | 'txt' | 'sql'; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; insertMode?: SqlInsertMode }) => + (options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' | 'json' | 'txt' | 'sql'; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; insertMode?: SqlInsertMode }) => queryStore.buildQueryResultExportRequest(activeTab.id, options) " :all-export-results="allResultExportSheets" diff --git a/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts index 5bb3fd75b8..1baabcb313 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts @@ -139,4 +139,12 @@ describe("query result SQL export progress", () => { expect(cancelHandler.value).toBeNull(); expect(api.startQueryResultExport).toHaveBeenCalledWith(expect.objectContaining({ format: "sql", exportTableName: "users", exportColumnTypes: ["int4", "text"] }), expect.any(Function)); }); + + it("routes incomplete query-result JSON exports through the streaming backend", async () => { + const state = useDataGridExport(createOptions()); + + await state.exportJson(); + + expect(api.startQueryResultExport).toHaveBeenCalledWith(expect.objectContaining({ format: "json", sql: "SELECT id, name FROM users" }), expect.any(Function)); + }); }); diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index b6216a779e..3784a825bd 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -116,7 +116,7 @@ export interface UseDataGridExportOptions { queryResultExportRequest?: (options: { exportId: string; filePath: string; - format: "csv" | "xlsx" | "txt" | "sql"; + format: "csv" | "xlsx" | "json" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; @@ -878,6 +878,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { await runExclusiveExport(async () => { try { if (await exportFullTableDataViaBackend("json", rowIds)) return; + if (await exportQueryResultViaBackend("json", rowIds)) return; let outputPath = exportFileName("export", "json"); if (isTauriRuntime()) { @@ -1267,7 +1268,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return true; } - async function exportQueryResultViaBackend(format: "csv" | "xlsx" | "txt" | "sql", rowIds?: number[], includeSqlSheet = false, headerMode: XlsxHeaderMode = "name", autoFilter = true, insertMode?: SqlInsertMode): Promise { + async function exportQueryResultViaBackend(format: "csv" | "xlsx" | "json" | "txt" | "sql", rowIds?: number[], includeSqlSheet = false, headerMode: XlsxHeaderMode = "name", autoFilter = true, insertMode?: SqlInsertMode): Promise { if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) { return false; } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index b1fddc4111..f07d0692b5 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -5327,7 +5327,7 @@ export interface QueryResultExportRequest { databaseType: DatabaseType; useAgentCursor: boolean; filePath: string; - format: "csv" | "xlsx" | "txt" | "sql"; + format: "csv" | "xlsx" | "json" | "txt" | "sql"; insertMode?: SqlInsertMode; csvQuoteMode?: CsvQuoteMode; includeSqlSheet?: boolean; diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index de0ba3f8ef..b472d05701 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -157,7 +157,7 @@ function cloneTabDraft(value: T): T { interface BuildQueryResultExportRequestOptions { exportId: string; filePath: string; - format: "csv" | "xlsx" | "txt" | "sql"; + format: "csv" | "xlsx" | "json" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 452fbf5c73..6a268e4275 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -324,6 +324,63 @@ struct SqlInsertWriter { identifier_quote: Option, } +/// Streams query-result JSON rows without retaining the complete result set. +/// The staged target is committed only after the closing array delimiter is +/// written, so cancellation or an execution error cannot replace an existing +/// destination with invalid JSON. +struct JsonStreamWriter { + file: Option>, + target: Option, + wrote_row: bool, +} + +impl JsonStreamWriter { + fn create(request: &QueryResultExportRequest) -> Result { + let target = StagedExportTarget::new(&request.file_path)?; + let mut file = BufWriter::new( + File::create(target.path()).map_err(|e| format!("Failed to create JSON export temp file: {e}"))?, + ); + file.write_all(b"[\n").map_err(|e| format!("Failed to start JSON export: {e}"))?; + Ok(Self { file: Some(file), target: Some(target), wrote_row: false }) + } + + fn write_row(&mut self, columns: &[String], row: &[Value]) -> Result<(), String> { + let file = self.file.as_mut().ok_or_else(|| "JSON export file already closed".to_string())?; + if self.wrote_row { + file.write_all(b",\n").map_err(|e| format!("Failed to write JSON separator: {e}"))?; + } + let mut object = serde_json::Map::new(); + for (index, column) in columns.iter().enumerate() { + if let Some(value) = row.get(index) { + object.insert(column.clone(), value.clone()); + } + } + serde_json::to_writer_pretty(&mut *file, &Value::Object(object)) + .map_err(|e| format!("Failed to write JSON row: {e}"))?; + self.wrote_row = true; + Ok(()) + } + + fn write_rows(&mut self, columns: &[String], rows: &[Vec]) -> Result<(), String> { + for row in rows { + self.write_row(columns, row)?; + } + Ok(()) + } + + fn finish(mut self) -> Result<(), String> { + let file = self.file.as_mut().ok_or_else(|| "JSON export file already closed".to_string())?; + file.write_all(b"\n]\n").map_err(|e| format!("Failed to finish JSON export: {e}"))?; + file.flush().map_err(|e| format!("Failed to flush JSON export: {e}"))?; + self.file.take(); + self.target + .take() + .ok_or_else(|| "JSON export target already finalized".to_string())? + .commit() + .map_err(|error| format!("Failed to finalize JSON export file: {error}")) + } +} + impl SqlInsertWriter { /// Create the writer and open the temp file. Column metadata is supplied later /// via [`SqlInsertWriter::set_columns`] once the executed result is known. @@ -730,7 +787,7 @@ async fn export_query_result_core_inner( session_id: &mut Option, ) -> Result<(), String> { let format = request.format.to_lowercase(); - if format != "csv" && format != "xlsx" && format != "txt" && format != "sql" { + if format != "csv" && format != "xlsx" && format != "json" && format != "txt" && format != "sql" { return Err(format!("Unsupported streaming query-result export format: {format}")); } @@ -791,6 +848,7 @@ async fn export_query_result_core_inner( return Err(error.to_string()); } + let mut json_writer = if format == "json" { Some(JsonStreamWriter::create(request)?) } else { None }; let mut sql_writer: Option = if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; @@ -945,6 +1003,9 @@ async fn export_query_result_core_inner( )?; } } + } else if format == "json" { + let writer = json_writer.as_mut().ok_or_else(|| "JSON export writer missing".to_string())?; + writer.write_rows(&columns, formatted_rows.as_ref())?; } else if format == "sql" { let writer = sql_writer.as_mut().ok_or_else(|| "SQL export writer missing".to_string())?; for (row_index, row) in formatted_rows.into_owned().into_iter().enumerate() { @@ -1019,6 +1080,10 @@ async fn export_query_result_core_inner( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } + } else if format == "json" { + if let Some(writer) = json_writer.take() { + writer.finish()?; + } } else if format == "sql" { if let Some(writer) = sql_writer.take() { writer.finish()?; @@ -1099,6 +1164,7 @@ async fn try_export_postgres_query_result_stream( }; let mut xlsx = None; let mut text_buffer = String::new(); + let mut json_writer = if format == "json" { Some(JsonStreamWriter::create(request)?) } else { None }; let mut sql_writer: Option = if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let budget = operation_budget_for_pool_key(state, &pool_key, query_export_timeout(request.timeout_secs)).await; @@ -1124,7 +1190,7 @@ async fn try_export_postgres_query_result_stream( } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns, request.csv_quote_mode); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; - } else { + } else if json_writer.is_none() { let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; xlsx = Some(start_query_result_xlsx_workbook( @@ -1151,6 +1217,8 @@ async fn try_export_postgres_query_result_stream( &mut text_buffer, request.csv_quote_mode, )?; + } else if let Some(writer) = json_writer.as_mut() { + writer.write_row(&columns, formatted.as_ref())?; } else if let Some(writer) = xlsx.as_mut() { writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; } else { @@ -1210,6 +1278,8 @@ async fn try_export_postgres_query_result_stream( } if let Some(writer) = sql_writer.take() { writer.finish()?; + } else if let Some(writer) = json_writer.take() { + writer.finish()?; } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; @@ -1290,6 +1360,7 @@ async fn try_export_mysql_query_result_stream( }; let mut xlsx = None; let mut text_buffer = String::new(); + let mut json_writer = if format == "json" { Some(JsonStreamWriter::create(request)?) } else { None }; let mut sql_writer: Option = if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); @@ -1365,7 +1436,7 @@ async fn try_export_mysql_query_result_stream( } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns, request.csv_quote_mode); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; - } else { + } else if json_writer.is_none() { let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; xlsx = Some(start_query_result_xlsx_workbook( @@ -1392,6 +1463,8 @@ async fn try_export_mysql_query_result_stream( &mut text_buffer, request.csv_quote_mode, )?; + } else if let Some(writer) = json_writer.as_mut() { + writer.write_row(&columns, formatted.as_ref())?; } else if let Some(writer) = xlsx.as_mut() { writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; } else { @@ -1483,6 +1556,8 @@ async fn try_export_mysql_query_result_stream( } if let Some(writer) = sql_writer.take() { writer.finish()?; + } else if let Some(writer) = json_writer.take() { + writer.finish()?; } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; @@ -1559,6 +1634,7 @@ async fn try_export_clickhouse_query_result_stream( }; let mut xlsx = None; let mut text_buffer = String::new(); + let mut json_writer = if format == "json" { Some(JsonStreamWriter::create(request)?) } else { None }; let mut sql_writer: Option = if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); @@ -1585,7 +1661,7 @@ async fn try_export_clickhouse_query_result_stream( } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns, request.csv_quote_mode); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; - } else { + } else if json_writer.is_none() { let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; xlsx = Some(start_query_result_xlsx_workbook( @@ -1612,6 +1688,8 @@ async fn try_export_clickhouse_query_result_stream( &mut text_buffer, request.csv_quote_mode, )?; + } else if let Some(writer) = json_writer.as_mut() { + writer.write_row(&columns, formatted.as_ref())?; } else if let Some(writer) = xlsx.as_mut() { writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; } else { @@ -1677,6 +1755,8 @@ async fn try_export_clickhouse_query_result_stream( } if let Some(writer) = sql_writer.take() { writer.finish()?; + } else if let Some(writer) = json_writer.take() { + writer.finish()?; } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; @@ -1734,6 +1814,7 @@ async fn try_export_sqlserver_query_result_stream( }; let mut xlsx = None; let mut text_buffer = String::new(); + let mut json_writer = if format == "json" { Some(JsonStreamWriter::create(request)?) } else { None }; let mut sql_writer: Option = if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); @@ -1774,7 +1855,7 @@ async fn try_export_sqlserver_query_result_stream( } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns, request.csv_quote_mode); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; - } else { + } else if json_writer.is_none() { let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; xlsx = @@ -1797,6 +1878,8 @@ async fn try_export_sqlserver_query_result_stream( &mut text_buffer, request.csv_quote_mode, )?; + } else if let Some(writer) = json_writer.as_mut() { + writer.write_row(&columns, formatted.as_ref())?; } else if let Some(writer) = xlsx.as_mut() { writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; } else { @@ -1867,6 +1950,8 @@ async fn try_export_sqlserver_query_result_stream( } if let Some(writer) = sql_writer.take() { writer.finish()?; + } else if let Some(writer) = json_writer.take() { + writer.finish()?; } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; @@ -2080,6 +2165,37 @@ mod tests { ); } + #[test] + fn json_stream_writer_emits_a_valid_array_without_retaining_rows() { + let dir = tempfile::tempdir().expect("temp dir"); + let destination = dir.path().join("result.json"); + let mut req = request("json", None, None); + req.file_path = destination.to_string_lossy().into_owned(); + + let mut writer = JsonStreamWriter::create(&req).expect("create JSON writer"); + writer + .write_rows( + &["id".to_string(), "name".to_string()], + &[vec![serde_json::json!(1), serde_json::json!("Ada")]], + ) + .expect("write first batch"); + writer + .write_rows(&["id".to_string(), "name".to_string()], &[vec![serde_json::json!(2), serde_json::Value::Null]]) + .expect("write second batch"); + writer.finish().expect("finish JSON writer"); + + let output: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(destination).expect("read JSON output")) + .expect("parse JSON output"); + assert_eq!( + output, + serde_json::json!([ + {"id": 1, "name": "Ada"}, + {"id": 2, "name": null} + ]) + ); + } + #[test] fn reusable_text_row_buffer_preserves_query_null_semantics() { let row = vec![Value::Null, serde_json::json!(""), serde_json::json!("line\n\"two\"")]; diff --git a/crates/dbx-web/src/routes/query_result_export.rs b/crates/dbx-web/src/routes/query_result_export.rs index 82ce87408a..91c17c2096 100644 --- a/crates/dbx-web/src/routes/query_result_export.rs +++ b/crates/dbx-web/src/routes/query_result_export.rs @@ -44,6 +44,8 @@ pub async fn start_query_result_export( let ext = match req.format.as_str() { "csv" => "csv", "xlsx" => "xlsx", + "json" => "json", + "txt" => "txt", _ => return Err(AppError::from(format!("Unsupported query result export format: {}", req.format))), }; let tmp_dir = state.data_dir.join("tmp"); @@ -162,6 +164,8 @@ pub async fn query_result_export_download( let content_type = match export_file.format.as_str() { "csv" => "text/csv; charset=utf-8", "xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "json" => "application/json; charset=utf-8", + "txt" => "text/plain; charset=utf-8", format => return Err(AppError::from(format!("Unknown format: {format}"))), };