feat(export): configurable CSV field quoting mode
This commit is contained in:
32 files changed
+442
-56
No files matched your search
@@ -11,6 +11,14 @@ use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions}
|
||||
|
||||
const TABLE_DATA_EXPORT_PAGE_SIZE: usize = 10_000;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum CsvQuoteMode {
|
||||
#[default]
|
||||
All,
|
||||
Necessary,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TableCsvExportOptions {
|
||||
@@ -26,6 +34,8 @@ pub struct TableCsvExportOptions {
|
||||
pub page_size: Option<usize>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub timeout_secs: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub csv_quote_mode: CsvQuoteMode,
|
||||
}
|
||||
|
||||
/// CSV 转义直写目标 buffer:包引号 + 内部 `"` 翻倍。值不含 `"` 时整段拷贝,
|
||||
@@ -46,6 +56,18 @@ pub(crate) fn push_csv_escaped(out: &mut String, value: &str) {
|
||||
out.push('"');
|
||||
}
|
||||
|
||||
fn csv_field_needs_quotes(value: &str) -> bool {
|
||||
value.bytes().any(|byte| matches!(byte, b',' | b'"' | b'\n' | b'\r'))
|
||||
}
|
||||
|
||||
pub(crate) fn push_csv_field(out: &mut String, value: &str, quote_mode: CsvQuoteMode) {
|
||||
if quote_mode == CsvQuoteMode::All || csv_field_needs_quotes(value) {
|
||||
push_csv_escaped(out, value);
|
||||
} else {
|
||||
out.push_str(value);
|
||||
}
|
||||
}
|
||||
|
||||
struct CsvEscapedWriter<'a>(&'a mut String);
|
||||
|
||||
impl fmt::Write for CsvEscapedWriter<'_> {
|
||||
@@ -72,6 +94,26 @@ pub(crate) fn push_csv_text_value(out: &mut String, value: &Value) {
|
||||
out.push('"');
|
||||
}
|
||||
|
||||
fn push_csv_value_with_quote_mode(out: &mut String, value: &Value, quote_mode: CsvQuoteMode, quote_null: bool) {
|
||||
if quote_mode == CsvQuoteMode::All {
|
||||
if value.is_null() && !quote_null {
|
||||
return;
|
||||
}
|
||||
push_csv_text_value(out, value);
|
||||
return;
|
||||
}
|
||||
|
||||
match value {
|
||||
Value::Null => {}
|
||||
Value::String(value) => push_csv_field(out, value, quote_mode),
|
||||
Value::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
|
||||
Value::Number(value) => {
|
||||
fmt::write(out, format_args!("{value}")).expect("writing a number into a String cannot fail")
|
||||
}
|
||||
other => push_csv_field(out, &other.to_string(), quote_mode),
|
||||
}
|
||||
}
|
||||
|
||||
fn push_csv_value(out: &mut String, value: &Value) {
|
||||
if value.is_null() {
|
||||
return;
|
||||
@@ -118,12 +160,19 @@ pub(crate) fn estimated_rows_capacity(rows: &[Vec<Value>]) -> usize {
|
||||
cells.saturating_mul(12).min(ROWS_CAPACITY_ESTIMATE_MAX)
|
||||
}
|
||||
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
pub(crate) fn escape_csv(value: &str) -> String {
|
||||
let mut out = String::with_capacity(value.len() + 2);
|
||||
push_csv_escaped(&mut out, value);
|
||||
out
|
||||
}
|
||||
|
||||
fn escape_csv_with_quote_mode(value: &str, quote_mode: CsvQuoteMode) -> String {
|
||||
let mut out = String::with_capacity(value.len() + 2);
|
||||
push_csv_field(&mut out, value, quote_mode);
|
||||
out
|
||||
}
|
||||
|
||||
fn format_csv_value(value: &Value) -> String {
|
||||
let mut out = String::new();
|
||||
push_csv_value(&mut out, value);
|
||||
@@ -172,20 +221,28 @@ pub(crate) fn format_tsv(columns: &[String], rows: &[Vec<Value>]) -> String {
|
||||
/// use the same empty-cell representation as table-data exports. Used by the
|
||||
/// streaming query-result export for batches after the first.
|
||||
pub(crate) fn push_query_result_csv_row(out: &mut String, row: &[Value]) {
|
||||
push_query_result_csv_row_with_quote_mode(out, row, CsvQuoteMode::All);
|
||||
}
|
||||
|
||||
pub(crate) fn push_query_result_csv_row_with_quote_mode(out: &mut String, row: &[Value], quote_mode: CsvQuoteMode) {
|
||||
for (cell_index, cell) in row.iter().enumerate() {
|
||||
if cell_index > 0 {
|
||||
out.push(',');
|
||||
}
|
||||
push_csv_value(out, cell);
|
||||
push_csv_value_with_quote_mode(out, cell, quote_mode, false);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn push_table_csv_row(out: &mut String, row: &[Value]) {
|
||||
push_table_csv_row_with_quote_mode(out, row, CsvQuoteMode::All);
|
||||
}
|
||||
|
||||
pub(crate) fn push_table_csv_row_with_quote_mode(out: &mut String, row: &[Value], quote_mode: CsvQuoteMode) {
|
||||
for (cell_index, cell) in row.iter().enumerate() {
|
||||
if cell_index > 0 {
|
||||
out.push(',');
|
||||
}
|
||||
push_csv_text_value(out, cell);
|
||||
push_csv_value_with_quote_mode(out, cell, quote_mode, true);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -223,30 +280,73 @@ pub fn format_csv(columns: &[String], rows: &[Vec<Value>]) -> String {
|
||||
format_csv_with_value_formatter(columns, rows)
|
||||
}
|
||||
|
||||
pub fn format_csv_with_quote_mode(columns: &[String], rows: &[Vec<Value>], quote_mode: CsvQuoteMode) -> String {
|
||||
let mut out = String::with_capacity(
|
||||
estimated_rows_capacity(rows).saturating_add(columns.len().saturating_mul(12)).min(ROWS_CAPACITY_ESTIMATE_MAX),
|
||||
);
|
||||
for (index, column) in columns.iter().enumerate() {
|
||||
if index > 0 {
|
||||
out.push(',');
|
||||
}
|
||||
push_csv_field(&mut out, column, quote_mode);
|
||||
}
|
||||
out.push('\n');
|
||||
for (row_index, row) in rows.iter().enumerate() {
|
||||
if row_index > 0 {
|
||||
out.push('\n');
|
||||
}
|
||||
push_query_result_csv_row_with_quote_mode(&mut out, row, quote_mode);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn format_query_result_csv(columns: &[String], rows: &[Vec<Value>]) -> String {
|
||||
format_csv(columns, rows)
|
||||
}
|
||||
|
||||
fn write_csv_text_row(writer: &mut impl Write, values: impl IntoIterator<Item = String>) -> Result<(), String> {
|
||||
pub fn format_query_result_csv_with_quote_mode(
|
||||
columns: &[String],
|
||||
rows: &[Vec<Value>],
|
||||
quote_mode: CsvQuoteMode,
|
||||
) -> String {
|
||||
format_csv_with_quote_mode(columns, rows, quote_mode)
|
||||
}
|
||||
|
||||
fn write_csv_text_row(
|
||||
writer: &mut impl Write,
|
||||
values: impl IntoIterator<Item = String>,
|
||||
quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
let mut first = true;
|
||||
for value in values {
|
||||
if !first {
|
||||
writer.write_all(b",").map_err(|err| err.to_string())?;
|
||||
}
|
||||
first = false;
|
||||
writer.write_all(escape_csv(&value).as_bytes()).map_err(|err| err.to_string())?;
|
||||
writer.write_all(escape_csv_with_quote_mode(&value, quote_mode).as_bytes()).map_err(|err| err.to_string())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_csv_value_row(writer: &mut impl Write, values: impl IntoIterator<Item = Value>) -> Result<(), String> {
|
||||
fn write_csv_value_row(
|
||||
writer: &mut impl Write,
|
||||
values: impl IntoIterator<Item = Value>,
|
||||
quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
let mut first = true;
|
||||
for value in values {
|
||||
if !first {
|
||||
writer.write_all(b",").map_err(|err| err.to_string())?;
|
||||
}
|
||||
first = false;
|
||||
writer.write_all(format_csv_value(&value).as_bytes()).map_err(|err| err.to_string())?;
|
||||
let formatted = if quote_mode == CsvQuoteMode::All {
|
||||
format_csv_value(&value)
|
||||
} else {
|
||||
let mut formatted = String::new();
|
||||
push_csv_value_with_quote_mode(&mut formatted, &value, quote_mode, false);
|
||||
formatted
|
||||
};
|
||||
writer.write_all(formatted.as_bytes()).map_err(|err| err.to_string())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -304,7 +404,7 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor
|
||||
.await?;
|
||||
|
||||
if !wrote_header {
|
||||
write_csv_text_row(&mut writer, result.columns)?;
|
||||
write_csv_text_row(&mut writer, result.columns, options.csv_quote_mode)?;
|
||||
wrote_header = true;
|
||||
}
|
||||
|
||||
@@ -314,7 +414,7 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor
|
||||
}
|
||||
for row in result.rows {
|
||||
writer.write_all(b"\n").map_err(|err| err.to_string())?;
|
||||
write_csv_value_row(&mut writer, row)?;
|
||||
write_csv_value_row(&mut writer, row, options.csv_quote_mode)?;
|
||||
}
|
||||
|
||||
rows_exported += fetched as u64;
|
||||
@@ -333,7 +433,10 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{format_csv, format_query_result_csv, format_query_result_csv_rows, format_tsv};
|
||||
use super::{
|
||||
format_csv, format_csv_with_quote_mode, format_query_result_csv, format_query_result_csv_rows, format_tsv,
|
||||
CsvQuoteMode,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
@@ -354,6 +457,27 @@ mod tests {
|
||||
assert_eq!(out, "\"id\",\"note\"\n\"1\",");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn necessary_quote_mode_only_quotes_csv_special_characters() {
|
||||
let out = format_csv_with_quote_mode(
|
||||
&["id".to_string(), "district,name".to_string(), "note".to_string()],
|
||||
&[
|
||||
vec![json!(2085252644_u64), json!("延庆县"), json!("plain")],
|
||||
vec![json!(2085252645_u64), json!("门头沟区"), json!("line 1\n\"line 2\"")],
|
||||
],
|
||||
CsvQuoteMode::Necessary,
|
||||
);
|
||||
assert_eq!(
|
||||
out,
|
||||
"id,\"district,name\",note\n2085252644,延庆县,plain\n2085252645,门头沟区,\"line 1\n\"\"line 2\"\"\""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn csv_quote_mode_defaults_to_all_for_backward_compatibility() {
|
||||
assert_eq!(CsvQuoteMode::default(), CsvQuoteMode::All);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_streamed_query_result_null_as_empty_cell_and_preserves_literal_null() {
|
||||
let out = format_query_result_csv_rows(&[vec![Value::Null, json!("NULL"), json!("")]]);
|
||||
|
||||
@@ -10,7 +10,10 @@ use std::sync::{
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::connection::{AppState, PoolKind};
|
||||
use crate::csv_export::{format_query_result_csv, format_tsv, push_query_result_csv_row, push_tsv_row};
|
||||
use crate::csv_export::{
|
||||
format_query_result_csv_with_quote_mode, format_tsv, push_query_result_csv_row_with_quote_mode, push_tsv_row,
|
||||
CsvQuoteMode,
|
||||
};
|
||||
pub use crate::database_export::ExportStatus;
|
||||
use crate::database_export::{build_export_insert_statements, is_export_cancelled, BuildExportInsertStatementsOptions};
|
||||
use crate::models::connection::DatabaseType;
|
||||
@@ -79,6 +82,8 @@ pub struct QueryResultExportRequest {
|
||||
pub file_path: String,
|
||||
pub format: String,
|
||||
#[serde(default)]
|
||||
pub csv_quote_mode: CsvQuoteMode,
|
||||
#[serde(default)]
|
||||
pub include_sql_sheet: bool,
|
||||
pub page_size: usize,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -410,8 +415,12 @@ fn effective_row_limit(request: &QueryResultExportRequest) -> Option<usize> {
|
||||
request.row_limit
|
||||
}
|
||||
|
||||
fn format_text_export_header(format: &str, columns: &[String]) -> String {
|
||||
let content = if format == "csv" { format_query_result_csv(columns, &[]) } else { format_tsv(columns, &[]) };
|
||||
fn format_text_export_header(format: &str, columns: &[String], csv_quote_mode: CsvQuoteMode) -> String {
|
||||
let content = if format == "csv" {
|
||||
format_query_result_csv_with_quote_mode(columns, &[], csv_quote_mode)
|
||||
} else {
|
||||
format_tsv(columns, &[])
|
||||
};
|
||||
content.strip_suffix('\n').unwrap_or(&content).to_string()
|
||||
}
|
||||
|
||||
@@ -420,11 +429,12 @@ fn write_text_export_row<W: Write>(
|
||||
format: &str,
|
||||
row: &[Value],
|
||||
buffer: &mut String,
|
||||
csv_quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
buffer.clear();
|
||||
buffer.push('\n');
|
||||
if format == "csv" {
|
||||
push_query_result_csv_row(buffer, row);
|
||||
push_query_result_csv_row_with_quote_mode(buffer, row, csv_quote_mode);
|
||||
} else {
|
||||
push_tsv_row(buffer, row);
|
||||
}
|
||||
@@ -436,12 +446,13 @@ fn write_text_export_rows<W: Write>(
|
||||
format: &str,
|
||||
rows: &[Vec<Value>],
|
||||
buffer: &mut String,
|
||||
csv_quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
buffer.clear();
|
||||
for row in rows {
|
||||
buffer.push('\n');
|
||||
if format == "csv" {
|
||||
push_query_result_csv_row(buffer, row);
|
||||
push_query_result_csv_row_with_quote_mode(buffer, row, csv_quote_mode);
|
||||
} else {
|
||||
push_tsv_row(buffer, row);
|
||||
}
|
||||
@@ -905,14 +916,26 @@ async fn export_query_result_core_inner(
|
||||
if format == "csv" || format == "txt" {
|
||||
if let Some(file) = text_file.as_mut() {
|
||||
if !wrote_text_header {
|
||||
let header = format_text_export_header(&format, &columns);
|
||||
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}"))?;
|
||||
if row_count > 0 {
|
||||
write_text_export_rows(file, &format, formatted_rows.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_rows(
|
||||
file,
|
||||
&format,
|
||||
formatted_rows.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
}
|
||||
wrote_text_header = true;
|
||||
} else if row_count > 0 {
|
||||
write_text_export_rows(file, &format, formatted_rows.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_rows(
|
||||
file,
|
||||
&format,
|
||||
formatted_rows.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
}
|
||||
}
|
||||
} else if format == "sql" {
|
||||
@@ -981,7 +1004,7 @@ async fn export_query_result_core_inner(
|
||||
|
||||
if format == "csv" || format == "txt" {
|
||||
if !wrote_text_header {
|
||||
let header = format_text_export_header(&format, &columns);
|
||||
let header = format_text_export_header(&format, &columns, request.csv_quote_mode);
|
||||
if let Some(file) = text_file.as_mut() {
|
||||
file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?;
|
||||
}
|
||||
@@ -1092,7 +1115,7 @@ async fn try_export_postgres_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.set_columns(columns.clone(), &column_types, &[], request);
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
let header = format_text_export_header(format, &columns);
|
||||
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 {
|
||||
let xlsx_file =
|
||||
@@ -1114,7 +1137,13 @@ async fn try_export_postgres_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.write_row(formatted.into_owned(), None)?;
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
write_text_export_row(file, format, formatted.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_row(
|
||||
file,
|
||||
format,
|
||||
formatted.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
} 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 {
|
||||
@@ -1327,7 +1356,7 @@ async fn try_export_mysql_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.set_columns(columns.clone(), &column_types, &[], request);
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
let header = format_text_export_header(format, &columns);
|
||||
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 {
|
||||
let xlsx_file =
|
||||
@@ -1349,7 +1378,13 @@ async fn try_export_mysql_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.write_row(formatted.into_owned(), None)?;
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
write_text_export_row(file, format, formatted.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_row(
|
||||
file,
|
||||
format,
|
||||
formatted.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
} 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 {
|
||||
@@ -1541,7 +1576,7 @@ async fn try_export_clickhouse_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.set_columns(columns.clone(), &column_types, &[], request);
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
let header = format_text_export_header(format, &columns);
|
||||
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 {
|
||||
let xlsx_file =
|
||||
@@ -1563,7 +1598,13 @@ async fn try_export_clickhouse_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.write_row(formatted.into_owned(), None)?;
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
write_text_export_row(file, format, formatted.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_row(
|
||||
file,
|
||||
format,
|
||||
formatted.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
} 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 {
|
||||
@@ -1724,7 +1765,7 @@ async fn try_export_sqlserver_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.set_columns(columns.clone(), &temporal_column_types, &[], request);
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
let header = format_text_export_header(format, &columns);
|
||||
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 {
|
||||
let xlsx_file =
|
||||
@@ -1742,7 +1783,13 @@ async fn try_export_sqlserver_query_result_stream(
|
||||
if let Some(writer) = sql_writer.as_mut() {
|
||||
writer.write_row(formatted.into_owned(), None)?;
|
||||
} else if let Some(file) = text_file.as_mut() {
|
||||
write_text_export_row(file, format, formatted.as_ref(), &mut text_buffer)?;
|
||||
write_text_export_row(
|
||||
file,
|
||||
format,
|
||||
formatted.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
} 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 {
|
||||
@@ -1922,6 +1969,7 @@ mod tests {
|
||||
client_session_id: None,
|
||||
execution_id: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
export_table_name: None,
|
||||
export_column_types: None,
|
||||
numeric_column_right_align: false,
|
||||
@@ -1969,7 +2017,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn txt_export_header_keeps_columns_for_empty_results() {
|
||||
assert_eq!(format_text_export_header("txt", &["id".to_string(), "note".to_string()]), "id\tnote");
|
||||
assert_eq!(
|
||||
format_text_export_header("txt", &["id".to_string(), "note".to_string()], CsvQuoteMode::All),
|
||||
"id\tnote"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1978,7 +2029,7 @@ mod tests {
|
||||
let mut output = Vec::new();
|
||||
let mut buffer = String::new();
|
||||
|
||||
write_text_export_row(&mut output, "csv", &row, &mut buffer).expect("write csv row");
|
||||
write_text_export_row(&mut output, "csv", &row, &mut buffer, CsvQuoteMode::All).expect("write csv row");
|
||||
assert_eq!(String::from_utf8(output).expect("utf8 csv"), "\n,\"\",\"line\n\"\"two\"\"\"");
|
||||
}
|
||||
|
||||
|
||||
@@ -10,8 +10,13 @@ use tokio_util::sync::CancellationToken;
|
||||
use crate::connection::MysqlMode;
|
||||
use crate::connection::{config_for_pool_key, task_client_session_id, AppState, PoolKind};
|
||||
#[cfg(test)]
|
||||
use crate::csv_export::format_csv;
|
||||
#[cfg(test)]
|
||||
use crate::csv_export::format_tsv_rows;
|
||||
use crate::csv_export::{format_csv, format_tsv, push_table_csv_row, push_tsv_row};
|
||||
use crate::csv_export::{
|
||||
format_csv_with_quote_mode, format_tsv, push_table_csv_row, push_table_csv_row_with_quote_mode, push_tsv_row,
|
||||
CsvQuoteMode,
|
||||
};
|
||||
pub use crate::database_export::ExportStatus;
|
||||
use crate::database_export::{
|
||||
build_export_insert_statements, is_export_cancelled, is_internal_export_column, BuildExportInsertStatementsOptions,
|
||||
@@ -47,6 +52,8 @@ pub struct TableExportRequest {
|
||||
/// "csv", "xlsx", "json", "markdown", "sql", or "txt"
|
||||
pub format: String,
|
||||
#[serde(default)]
|
||||
pub csv_quote_mode: CsvQuoteMode,
|
||||
#[serde(default)]
|
||||
pub columns: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub column_types: Option<Vec<Option<String>>>,
|
||||
@@ -100,11 +107,17 @@ fn format_csv_rows(rows: &[Vec<Value>]) -> String {
|
||||
out
|
||||
}
|
||||
|
||||
fn write_table_text_row<W: Write>(file: &mut W, csv: bool, row: &[Value], buffer: &mut String) -> Result<(), String> {
|
||||
fn write_table_text_row<W: Write>(
|
||||
file: &mut W,
|
||||
csv: bool,
|
||||
row: &[Value],
|
||||
buffer: &mut String,
|
||||
csv_quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
buffer.clear();
|
||||
buffer.push('\n');
|
||||
if csv {
|
||||
push_table_csv_row(buffer, row);
|
||||
push_table_csv_row_with_quote_mode(buffer, row, csv_quote_mode);
|
||||
} else {
|
||||
push_tsv_row(buffer, row);
|
||||
}
|
||||
@@ -116,12 +129,13 @@ fn write_table_text_rows<W: Write>(
|
||||
csv: bool,
|
||||
rows: &[Vec<Value>],
|
||||
buffer: &mut String,
|
||||
csv_quote_mode: CsvQuoteMode,
|
||||
) -> Result<(), String> {
|
||||
buffer.clear();
|
||||
for row in rows {
|
||||
buffer.push('\n');
|
||||
if csv {
|
||||
push_table_csv_row(buffer, row);
|
||||
push_table_csv_row_with_quote_mode(buffer, row, csv_quote_mode);
|
||||
} else {
|
||||
push_tsv_row(buffer, row);
|
||||
}
|
||||
@@ -958,7 +972,7 @@ async fn try_export_native_table_stream(
|
||||
std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?,
|
||||
);
|
||||
file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?;
|
||||
let header = format_csv(col_names, &[]);
|
||||
let header = format_csv_with_quote_mode(col_names, &[], request.csv_quote_mode);
|
||||
let header = header.strip_suffix('\n').unwrap_or(&header);
|
||||
file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?;
|
||||
let mut row_buffer = String::new();
|
||||
@@ -977,7 +991,7 @@ async fn try_export_native_table_stream(
|
||||
column_types,
|
||||
request.date_time_format.as_deref(),
|
||||
);
|
||||
write_table_text_row(&mut file, true, formatted.as_ref(), &mut row_buffer)?;
|
||||
write_table_text_row(&mut file, true, formatted.as_ref(), &mut row_buffer, request.csv_quote_mode)?;
|
||||
rows_exported += 1;
|
||||
if rows_exported.is_multiple_of(progress_interval) {
|
||||
on_progress(TableExportProgress {
|
||||
@@ -1021,7 +1035,13 @@ async fn try_export_native_table_stream(
|
||||
column_types,
|
||||
request.date_time_format.as_deref(),
|
||||
);
|
||||
write_table_text_row(&mut file, false, formatted.as_ref(), &mut row_buffer)?;
|
||||
write_table_text_row(
|
||||
&mut file,
|
||||
false,
|
||||
formatted.as_ref(),
|
||||
&mut row_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
rows_exported += 1;
|
||||
if rows_exported.is_multiple_of(progress_interval) {
|
||||
on_progress(TableExportProgress {
|
||||
@@ -1578,12 +1598,19 @@ async fn export_table_data_core_inner(
|
||||
|
||||
if is_first_batch {
|
||||
// First batch: write header + rows via format_csv
|
||||
let csv_content = format_csv(&col_names, formatted_rows.as_ref());
|
||||
let csv_content =
|
||||
format_csv_with_quote_mode(&col_names, formatted_rows.as_ref(), request.csv_quote_mode);
|
||||
file.write_all(csv_content.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?;
|
||||
is_first_batch = false;
|
||||
} else {
|
||||
// Subsequent batches: write rows only (prepend newline for separation)
|
||||
write_table_text_rows(&mut file, true, formatted_rows.as_ref(), &mut text_buffer)?;
|
||||
write_table_text_rows(
|
||||
&mut file,
|
||||
true,
|
||||
formatted_rows.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
}
|
||||
|
||||
rows_exported += row_count as u64;
|
||||
@@ -1665,10 +1692,22 @@ async fn export_table_data_core_inner(
|
||||
);
|
||||
|
||||
if is_first_batch {
|
||||
write_table_text_rows(&mut file, false, formatted_rows.as_ref(), &mut text_buffer)?;
|
||||
write_table_text_rows(
|
||||
&mut file,
|
||||
false,
|
||||
formatted_rows.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
is_first_batch = false;
|
||||
} else {
|
||||
write_table_text_rows(&mut file, false, formatted_rows.as_ref(), &mut text_buffer)?;
|
||||
write_table_text_rows(
|
||||
&mut file,
|
||||
false,
|
||||
formatted_rows.as_ref(),
|
||||
&mut text_buffer,
|
||||
request.csv_quote_mode,
|
||||
)?;
|
||||
}
|
||||
|
||||
rows_exported += row_count as u64;
|
||||
@@ -2210,6 +2249,7 @@ mod tests {
|
||||
batch_size: Some(batch_size),
|
||||
row_limit,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2314,7 +2354,7 @@ mod tests {
|
||||
let mut output = Vec::new();
|
||||
let mut buffer = String::new();
|
||||
|
||||
write_table_text_row(&mut output, true, &row, &mut buffer).expect("write csv row");
|
||||
write_table_text_row(&mut output, true, &row, &mut buffer, CsvQuoteMode::All).expect("write csv row");
|
||||
assert_eq!(String::from_utf8(output).expect("utf8 csv"), "\n\"\",\"\",\"line\n\"\"two\"\"\"");
|
||||
}
|
||||
|
||||
@@ -2391,6 +2431,7 @@ mod tests {
|
||||
batch_size: Some(50),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2447,6 +2488,7 @@ mod tests {
|
||||
batch_size: Some(50),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2481,6 +2523,7 @@ mod tests {
|
||||
batch_size: Some(50),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2510,6 +2553,7 @@ mod tests {
|
||||
batch_size: Some(50),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2547,6 +2591,7 @@ mod tests {
|
||||
batch_size: Some(25),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2596,6 +2641,7 @@ mod tests {
|
||||
batch_size: Some(500),
|
||||
row_limit: Some(1000),
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2639,6 +2685,7 @@ mod tests {
|
||||
batch_size: Some(50),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2710,6 +2757,7 @@ mod tests {
|
||||
batch_size: Some(100),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2763,6 +2811,7 @@ mod tests {
|
||||
batch_size: Some(100),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
@@ -2818,6 +2867,7 @@ mod tests {
|
||||
batch_size: Some(100),
|
||||
row_limit: None,
|
||||
date_time_format: None,
|
||||
csv_quote_mode: CsvQuoteMode::All,
|
||||
numeric_column_right_align: false,
|
||||
column_comments: None,
|
||||
auto_filter: None,
|
||||
|
||||
Reference in new issue
Block a user