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 fn format_csv(columns: &[String], rows: &[Vec]) -> 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_to_csv_text(cell))).collect::>().join(",")) .collect::>() .join("\n"); format!("{header}\n{body}") } 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(), 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; 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\",\"\""); } use serde_json::Value; }