github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 04:07:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(export): stream JSON query exports and allow web TXT

Browse Source
This commit is contained in:
zipg
2026-09-14 20:04:11 +08:00
committed by GitHub
parent af424b782f
commit b70d6da227
8 changed files with 140 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/desktop/src/components/grid/DataGrid.vue
+1 -1
View File
@@ -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<string | null | undefined>;
apps/desktop/src/components/layout/ContentArea.vue
+1 -1
View File
@@ -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<string | null | undefined>; insertMode?: SqlInsertMode }) =>
(options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' | 'json' | 'txt' | 'sql'; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array<string | null | undefined>; insertMode?: SqlInsertMode }) =>
queryStore.buildQueryResultExportRequest(activeTab.id, options)
"
:all-export-results="allResultExportSheets"
apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts
Vendored
+8
View File
@@ -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));
});
});
apps/desktop/src/composables/useDataGridExport.ts
+3 -2
View File
@@ -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<string | null | undefined>;
@@ -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<boolean> {
async function exportQueryResultViaBackend(format: "csv" | "xlsx" | "json" | "txt" | "sql", rowIds?: number[], includeSqlSheet = false, headerMode: XlsxHeaderMode = "name", autoFilter = true, insertMode?: SqlInsertMode): Promise<boolean> {
if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) {
return false;
}
apps/desktop/src/lib/backend/tauri.ts
+1 -1
View File
@@ -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;
apps/desktop/src/stores/queryStore.ts
+1 -1
View File
@@ -157,7 +157,7 @@ function cloneTabDraft<T>(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<string | null | undefined>;
crates/dbx-core/src/query_result_export.rs
+121 -5
View File
@@ -324,6 +324,63 @@ struct SqlInsertWriter {
identifier_quote: Option<String>,
}
/// 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<BufWriter<File>>,
target: Option<StagedExportTarget>,
wrote_row: bool,
}
impl JsonStreamWriter {
fn create(request: &QueryResultExportRequest) -> Result<Self, String> {
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<Value>]) -> 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<String>,
) -> 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<SqlInsertWriter> =
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<SqlInsertWriter> =
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<SqlInsertWriter> =
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<SqlInsertWriter> =
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<SqlInsertWriter> =
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\"")];
crates/dbx-web/src/routes/query_result_export.rs
+4
View File
@@ -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}"))),
};