csv2xlsx_line/

lib.rs

use csv::StringRecord;
use std::error::Error;
use std::fs::File;
use std::env;

use rust_xlsxwriter::{Workbook, XlsxError};
use std::io::{BufReader, BufRead};
use csv::ReaderBuilder;
use chrono::Local;

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use std::path::{Path, PathBuf};

/// 读取 CSV 文件的表头（假设第一行为表头）
///
/// # 参数
///
/// * `path` - CSV 文件的路径
///
/// # 返回
///
/// 返回一个 `Result<StringRecord, Box<dyn Error>>`，其中 `StringRecord` 包含 CSV 表头。
pub fn read_columns_headers(path: &str) -> Result<StringRecord, Box<dyn Error>> {
    // 打开 CSV 文件
    let file = File::open(path)?;
    
    // 创建 CSV 读取器
    let mut rdr = csv::Reader::from_reader(file);
    
    // 读取 CSV 的表头（假定第一行是表头）
    let headers = rdr.headers()?.clone();
    
    // 返回 headers
    Ok(headers)
}

// 假设每个 Excel 文件的最大行数是一个常量
const ROWS_PER_FILE: usize = 1_000_000;

// 转化函数，增加 cancel_flag 参数，用于检测取消信号
pub fn convert_from_path(file_path: &str,  output_dir: &str, cancel_flag: &Arc<AtomicBool>) -> Result<(), Box<dyn Error>> {
    println!("开始时间：{}", Local::now().format("%Y-%m-%d %H:%M:%S"));

    // 确保输出目录存在
    std::fs::create_dir_all(output_dir)?;
    
    // 从文件路径获取基本文件名
    // let base_name = file_path.split('.').next().unwrap_or("not_found_file");
    let path = Path::new(file_path);
    let base_name = path
        .file_stem()
        .and_then(|os| os.to_str())
        .unwrap_or("not_found_file");

    // 在开始计数前检测取消标志
    if cancel_flag.load(Ordering::SeqCst) {
        println!("取消操作：在计数阶段");
        return Err("Operation cancelled".into());
    }
    let total_rows = count_csv_rows(file_path)?;
    println!("文件总行数（不含表头）：{}", total_rows);
    let total_files = (total_rows + ROWS_PER_FILE - 1) / ROWS_PER_FILE;
    println!("预计生成文件数：{}", total_files);

    // 打开 CSV 文件，并使用 8MB 缓冲区读取
    let file = File::open(file_path)?;
    let buf_reader = BufReader::with_capacity(8 * 1024 * 1024, file);
    let mut reader = ReaderBuilder::new().has_headers(true).from_reader(buf_reader);

    // 获取表头
    let headers: Vec<String> = reader
        .headers()?
        .iter()
        .map(|s| s.to_string())
        .collect();

    let mut file_index = 1;
    let mut current_chunk: Vec<Vec<String>> = Vec::with_capacity(ROWS_PER_FILE);

    // 逐条处理 CSV 记录
    for result in reader.records() {
        // 每次读取记录前检查取消信号
        if cancel_flag.load(Ordering::SeqCst) {
            println!("检测到取消信号，提前退出转换过程");
            return Err("Operation cancelled".into());
        }

        let record: StringRecord = result?;
        let row_data: Vec<String> = record.iter().map(|s| s.to_string()).collect();
        current_chunk.push(row_data);

        // 当累计达到设定行数时，写入一个 Excel 文件
        if current_chunk.len() == ROWS_PER_FILE {
            // 写入前再次检查取消信号
            if cancel_flag.load(Ordering::SeqCst) {
                println!("检测到取消信号，在写入文件前取消");
                return Err("Operation cancelled".into());
            }
            // let output_file_name = format!("{}-part{}.xlsx", base_name, file_index);
            // write_to_excel(&output_file_name, &headers, &current_chunk,  Arc::clone(&cancel_flag))?;
            // 构建完整输出路径
            let output_path = Path::new(output_dir)
                .join(format!("{}-part{}.xlsx", base_name, file_index));
            println!("写入文件：{}", output_path.to_string_lossy());
            write_to_excel(&output_path.to_string_lossy(), &headers, &current_chunk,  Arc::clone(&cancel_flag))?;
            current_chunk.clear();
            file_index += 1;
        }
    }

    // 如果剩余未满一批，则写入最后一个文件之前检查取消
    if !current_chunk.is_empty() {
        if cancel_flag.load(Ordering::SeqCst) {
            println!("检测到取消信号，在写入剩余文件前取消");
            return Err("Operation cancelled".into());
        }
        // let output_file_name = format!("{}-part{}.xlsx", base_name, file_index);
        // write_to_excel(&output_file_name, &headers, &current_chunk,  Arc::clone(&cancel_flag))?;
        let output_path = Path::new(output_dir)
            .join(format!("{}-part{}.xlsx", base_name, file_index));
        println!("写入文件：{}", output_path.to_string_lossy());
        write_to_excel(&output_path.to_string_lossy(), &headers, &current_chunk,  Arc::clone(&cancel_flag))?;
    }

    println!("结束时间：{}", Local::now().format("%Y-%m-%d %H:%M:%S"));
    Ok(())
}

// 将指定数据写入 Excel 文件的函数（这里暂不增加内部取消检查，如有需要可在循环中添加）
fn write_to_excel(
    file_path: &str,
    headers: &[String],
    data: &[Vec<String>],
    cancel_flag: Arc<AtomicBool>
) -> Result<(), Box<dyn Error>> {
    // 确保目标目录存在
    if let Some(parent) = Path::new(file_path).parent() {
        std::fs::create_dir_all(parent)?;
    }

    // 创建 Excel 工作簿对象（此处以 Workbook 为例，你可以使用相应的 Excel 写入库）
    let mut workbook = Workbook::new();
    let worksheet = workbook.add_worksheet_with_constant_memory();

    // 写入表头前检查是否被取消
    if cancel_flag.load(Ordering::SeqCst) {
        println!("取消信号：在写入表头前取消");
        // return Err(XlsxError("Operation cancelled".to_string()));
        return Err("Operation cancelled".into());
    }

    // 写入表头
    worksheet.write_row(0, 0, headers)?;

    // 逐行写入数据
    for (row_index, row) in data.iter().enumerate() {
        if cancel_flag.load(Ordering::SeqCst) {
            println!("取消信号：在写入第 {} 行时取消", row_index + 1);
            return Err("Operation cancelled".into());
        }
        worksheet.write_row((row_index + 1) as u32, 0, row)?;
    }
    workbook.save(file_path)?;
    Ok(())
}

// 计算 CSV 文件的行数（不包含表头）
fn count_csv_rows(file_path: &str) -> Result<usize, Box<dyn Error>> {
    let file = File::open(file_path)?;
    let reader = BufReader::with_capacity(8 * 1024 * 1024, file);
    // 计算行数并减去表头
    let total_rows = reader.lines().count().saturating_sub(1);
    Ok(total_rows)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    // #[test]
    // fn test_read_columns_headers() {
    //     // 创建一个临时 CSV 文件，写入测试数据
    //     let mut temp_file = NamedTempFile::new().expect("创建临时文件失败");
    //     writeln!(temp_file, "name,age,city").expect("写入测试数据失败");
    //     writeln!(temp_file, "Alice,30,New York").expect("写入测试数据失败");
        
    //     // 调用库函数
    //     let file = "data.csv";
    //     let headers = read_columns_headers(file);
    //     // println!("headers, {:?}", headers);
    //     println!("-------------------");
    //     // let headers = read_columns_headers(temp_file.path().to_str().unwrap()).expect("读取 CSV 表头失败");
        
    //     // 检查表头内容
    //     let expected = csv::StringRecord::from(vec!["name", "age", "city"]);
    //     // assert_eq!(headers, expected);
    // }

    #[test]
    fn test_convert_from_path() {
        // 调用库函数
        // let file = "data.csv";
        let file_path = "D:\\oracle-export-test\\csv\\aaa.csv";
        let output_dir = "D:\\oracle-export-test\\xlsx";
        let flag = Arc::new(AtomicBool::new(false));
        convert_from_path(file_path, output_dir, &flag);
    }
}