use std::fs::File;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use tempfile::NamedTempFile;
use crate::error::{SqawkError, SqawkResult};
use crate::table::{Table, Value};
pub struct CsvHandler {}
impl Default for CsvHandler {
fn default() -> Self {
Self::new()
}
}
impl CsvHandler {
pub fn new() -> Self {
CsvHandler {}
}
pub fn save_csv(&self, table: &Table, file_path: &Path) -> SqawkResult<()> {
let parent = file_path.parent().unwrap_or(Path::new("."));
let temp_file = NamedTempFile::new_in(parent).map_err(SqawkError::IoError)?;
{
let mut writer = csv::Writer::from_writer(temp_file.as_file());
writer
.write_record(table.columns())
.map_err(SqawkError::CsvError)?;
for row in table.rows() {
let string_values: Vec<String> =
row.iter().map(|value| value.to_string()).collect();
writer
.write_record(&string_values)
.map_err(SqawkError::CsvError)?;
}
writer.flush().map_err(SqawkError::IoError)?;
}
temp_file
.as_file()
.sync_all()
.map_err(SqawkError::IoError)?;
temp_file
.persist(file_path)
.map_err(|e| SqawkError::IoError(e.error))?;
Ok(())
}
pub fn load_csv(
&self,
file_spec: &str,
custom_columns: Option<Vec<String>>,
recover_errors: Option<bool>,
) -> SqawkResult<Table> {
let (table_name, file_path) = self.parse_file_spec(file_spec)?;
let file = File::open(&file_path)?;
let reader = BufReader::new(file);
let has_headers = true;
let mut csv_reader = csv::ReaderBuilder::new()
.has_headers(has_headers)
.comment(Some(b'#')) .flexible(true) .from_reader(reader);
let headers = if let Some(columns) = custom_columns {
columns
} else {
csv_reader
.headers()
.map_err(SqawkError::CsvError)?
.iter()
.map(|s| s.to_string())
.collect::<Vec<_>>()
};
let mut table = Table::new_with_delimiter(
&table_name,
headers,
Some(file_path.clone()),
",".to_string(),
);
let should_recover = recover_errors.unwrap_or(false);
let mut skipped_rows = 0;
let mut row_number = 0;
let column_count = table.column_count();
let mut row_buffer = Vec::with_capacity(column_count);
for result in csv_reader.records() {
row_number += 1;
match result {
Ok(record) => {
row_buffer.clear();
if should_recover && record.len() != column_count {
for i in 0..column_count {
if i < record.len() {
row_buffer.push(Value::from(record.get(i).unwrap_or("")));
} else {
row_buffer.push(Value::Null);
}
}
table.add_row_from_slice(&row_buffer)?;
} else {
for field in record.iter() {
row_buffer.push(Value::from(field));
}
if let Err(e) = table.add_row_from_slice(&row_buffer) {
if should_recover {
skipped_rows += 1;
eprintln!(
"Warning: Skipping row at line {} with inconsistent field count: {}",
row_number + 1, e
);
} else {
return Err(e);
}
}
}
}
Err(csv_err) if should_recover => {
skipped_rows += 1;
eprintln!(
"Warning: Skipping malformed row at line {}: {}",
row_number + 1, csv_err
);
}
Err(csv_err) => {
return Err(SqawkError::CsvParseError {
file: file_path.to_string_lossy().to_string(),
line: row_number + 1, error: format!("{}", csv_err),
});
}
}
}
if should_recover && skipped_rows > 0 {
eprintln!(
"Note: Skipped {} malformed rows while loading {}",
skipped_rows,
file_path.to_string_lossy()
);
}
Ok(table)
}
fn parse_file_spec(&self, file_spec: &str) -> SqawkResult<(String, PathBuf)> {
if let Some((table_name, file_path)) = file_spec.split_once('=') {
Ok((table_name.to_string(), PathBuf::from(file_path)))
} else {
let path = PathBuf::from(file_spec);
path.file_name().ok_or_else(|| {
SqawkError::InvalidFileSpec(format!("Invalid file specification: {}", file_spec))
})?;
let stem = path.file_stem().ok_or_else(|| {
SqawkError::InvalidFileSpec(format!("Invalid file specification: {}", file_spec))
})?;
Ok((stem.to_string_lossy().to_string(), path))
}
}
}