use serde::{Deserialize, Serialize}; use serde_json::Value; use std::fs::File; use std::io::{BufWriter, Write}; use crate::connection::AppState; use crate::models::connection::DatabaseType; use crate::query::{execute_sql_statement_with_options, QueryExecutionOptions}; use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions}; const TABLE_DATA_EXPORT_PAGE_SIZE: usize = 10_000; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TableCsvExportOptions { pub file_path: String, pub connection_id: String, pub database: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub schema: Option, pub table_name: String, #[serde(default)] pub columns: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub page_size: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub timeout_secs: Option, } pub(crate) fn escape_csv(value: &str) -> String { format!("\"{}\"", value.replace('"', "\"\"")) } pub(crate) fn value_to_csv_text(value: &Value) -> String { match value { Value::Null => String::new(), Value::Bool(v) => v.to_string(), Value::Number(v) => v.to_string(), Value::String(v) => v.clone(), other => other.to_string(), } } pub(crate) fn value_to_query_result_csv_text(value: &Value) -> String { match value { Value::Null => "NULL".to_string(), other => value_to_csv_text(other), } } /// Format query-result rows as CSV text without a header row, using the /// query-result NULL semantics (NULL → "NULL"). Used by the streaming /// query-result export for batches after the first. pub fn format_query_result_csv_rows(rows: &[Vec]) -> String { rows.iter() .map(|row| { row.iter().map(|cell| escape_csv(&value_to_query_result_csv_text(cell))).collect::>().join(",") }) .collect::>() .join("\n") } fn format_csv_with_value_formatter( columns: &[String], rows: &[Vec], value_formatter: fn(&Value) -> String, ) -> String { let header = columns.iter().map(|col| escape_csv(col)).collect::>().join(","); let body = rows .iter() .map(|row| row.iter().map(|cell| escape_csv(&value_formatter(cell))).collect::>().join(",")) .collect::>() .join("\n"); format!("{header}\n{body}") } pub fn format_csv(columns: &[String], rows: &[Vec]) -> String { format_csv_with_value_formatter(columns, rows, value_to_csv_text) } pub fn format_query_result_csv(columns: &[String], rows: &[Vec]) -> String { format_csv_with_value_formatter(columns, rows, value_to_query_result_csv_text) } fn write_csv_row(writer: &mut impl Write, values: impl IntoIterator) -> 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())?; } Ok(()) } async fn connection_database_type(state: &AppState, connection_id: &str) -> Result { state .configs .read() .await .get(connection_id) .map(|config| config.db_type) .ok_or_else(|| format!("Connection config not found: {connection_id}")) } pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExportOptions) -> Result { let database_type = connection_database_type(state, &options.connection_id).await?; let page_size = options.page_size.unwrap_or(TABLE_DATA_EXPORT_PAGE_SIZE).max(1); let mut writer = BufWriter::new(File::create(&options.file_path).map_err(|err| format!("Failed to write CSV file: {err}"))?); writer.write_all("\u{FEFF}".as_bytes()).map_err(|err| err.to_string())?; let mut offset = 0usize; let mut rows_exported = 0u64; let mut wrote_header = false; loop { let sql = build_table_data_select_sql(TableDataSelectSqlOptions { database_type: Some(database_type), schema: options.schema.clone(), table_name: options.table_name.clone(), table_type: None, primary_keys: Vec::new(), columns: options.columns.clone(), fallback_order_columns: Vec::new(), order_by: None, limit: Some(page_size), offset: Some(offset), where_input: None, include_row_id: false, }); let result = execute_sql_statement_with_options( state, &options.connection_id, &options.database, &sql, options.schema.as_deref(), None, QueryExecutionOptions { max_rows: Some(page_size), timeout_secs: options.timeout_secs, ..Default::default() }, ) .await?; if !wrote_header { write_csv_row(&mut writer, result.columns)?; wrote_header = true; } let fetched = result.rows.len(); if fetched == 0 { break; } for row in result.rows { writer.write_all(b"\n").map_err(|err| err.to_string())?; write_csv_row(&mut writer, row.iter().map(value_to_csv_text))?; } rows_exported += fetched as u64; if fetched < page_size { break; } offset += fetched; } if rows_exported == 0 { writer.write_all(b"\n").map_err(|err| err.to_string())?; } writer.flush().map_err(|err| err.to_string())?; Ok(rows_exported) } #[cfg(test)] mod tests { use super::{format_csv, format_query_result_csv}; use serde_json::json; #[test] fn formats_csv_with_header_and_escaped_values() { let out = format_csv(&["id".to_string(), "name".to_string()], &[vec![json!(1), json!("Ada \"Lovelace\"")]]); assert_eq!(out, "\"id\",\"name\"\n\"1\",\"Ada \"\"Lovelace\"\"\""); } #[test] fn formats_null_as_empty_cell() { let out = format_csv(&["id".to_string(), "note".to_string()], &[vec![json!(1), Value::Null]]); assert_eq!(out, "\"id\",\"note\"\n\"1\",\"\""); } #[test] fn formats_query_result_null_as_null_text() { let out = format_query_result_csv(&["id".to_string(), "note".to_string()], &[vec![json!(1), Value::Null]]); assert_eq!(out, "\"id\",\"note\"\n\"1\",\"NULL\""); } use serde_json::Value; }