GTool_xgfetl/module/csv/csvwriter/csvwriter.go

package csvwriter

import (
	"encoding/csv"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
)

// WriteRecordsToCSV 将数据写入 CSV 文件
// 参数：
// - filePath: CSV 文件路径
// - headers: 表头字段列表
// - records: 要写入的记录列表，每个记录是一个字符串切片
// - delimiter: CSV 文件的分隔符
// - appendToFile: 是否追加写入，true 表示在文件末尾追加，false 表示新建文件
// 返回值：
// - error: 如果发生错误，返回错误信息
func WriteRecordsToCSV(filePath string, headers []string, records [][]string, delimiter rune, appendToFile bool) error {
	var file *os.File
	var err error

	if appendToFile {
		// 追加写入
		file, err = os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			return fmt.Errorf("failed to open file for appending: %v", err)
		}
	} else {
		// 新建文件
		file, err = os.Create(filePath)
		if err != nil {
			return fmt.Errorf("failed to create file: %v", err)
		}
	}
	defer file.Close()

	writer := csv.NewWriter(file)
	writer.Comma = delimiter
	defer writer.Flush()

	// 如果是新文件且有表头，需要写入表头
	if !appendToFile && headers != nil {
		if err := writer.Write(headers); err != nil {
			return fmt.Errorf("failed to write headers to CSV: %v", err)
		}
	}

	// 写入记录
	if err := writer.WriteAll(records); err != nil {
		return fmt.Errorf("failed to write records to CSV: %v", err)
	}

	return nil
}

/*
WriteRecordsToCSVWithFileSizeLimit 将数据写入 CSV 文件，并基于文件大小进行分片
参数:
- baseFilePath: CSV 文件的基本路径，用于生成不同分片的文件名
- headers: 表的列名，用于 CSV 表头
- records: 要写入的表记录
- delimiter: CSV 的分隔符
- maxFileSizeMB: 单个 CSV 文件的最大大小 MB
- fileIndex: 当前文件的索引号，用于管理分片文件的命名
- currentFileSize: 当前文件的大小，用于判断是否需要创建新文件
- currentFilePath: 当前正在写入的文件路径，用于管理文件的创建和打开
返回值:
- error: 如果写入过程中有错误，返回相关的错误信息
*/
func WriteRecordsToCSVWithFileSizeLimit(baseFilePath string, headers []string, records [][]string, delimiter rune, maxFileSizeMB int64, fileIndex *int, currentFileSize *int64, currentFilePath *string) error {
	var file *os.File
	var writer *csv.Writer

	// 将 maxFileSize 从 MB 转换为字节
	maxFileSize := maxFileSizeMB * 1024 * 1024

	// 检查是否需要创建新文件
	if *currentFileSize == 0 {
		// 如果是第一次写入，创建新文件
		*currentFilePath = filepath.Join(baseFilePath + "_" + strconv.Itoa(*fileIndex) + ".csv")
		var err error
		file, err = os.Create(*currentFilePath)
		if err != nil {
			return fmt.Errorf("failed to create file: %v", err)
		}
		defer file.Close()

		writer = csv.NewWriter(file)
		writer.Comma = delimiter // 设置 CSV 分隔符
		defer writer.Flush()

		// 写入表头到 CSV 文件
		if err := writer.Write(headers); err != nil {
			return fmt.Errorf("failed to write header to CSV: %v", err)
		}
		*currentFileSize += estimateRecordSize(headers, delimiter)
	} else {
		// 如果文件已存在，追加写入
		var err error
		file, err = os.OpenFile(*currentFilePath, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
		if err != nil {
			return fmt.Errorf("failed to open file: %v", err)
		}
		defer file.Close()

		writer = csv.NewWriter(file)
		writer.Comma = delimiter // 设置 CSV 分隔符
		defer writer.Flush()
	}

	// 遍历记录逐行写入 CSV 文件
	for _, record := range records {
		// 如果文件大小超过限制，创建新文件
		if *currentFileSize > maxFileSize {
			writer.Flush()
			file.Close()

			// 增加文件索引，创建新的文件
			*fileIndex++
			*currentFilePath = filepath.Join(baseFilePath + "_" + strconv.Itoa(*fileIndex) + ".csv")
			var err error
			file, err = os.Create(*currentFilePath)
			if err != nil {
				return fmt.Errorf("failed to create file: %v", err)
			}

			writer = csv.NewWriter(file)
			writer.Comma = delimiter // 设置 CSV 分隔符
			defer writer.Flush()

			// 写入表头到新文件
			if err := writer.Write(headers); err != nil {
				return fmt.Errorf("failed to write header to new CSV: %v", err)
			}
			*currentFileSize = estimateRecordSize(headers, delimiter)
		}

		// 写入当前记录
		if err := writer.Write(record); err != nil {
			return fmt.Errorf("failed to write row to CSV: %v", err)
		}

		// 更新当前文件大小
		*currentFileSize += estimateRecordSize(record, delimiter)
	}

	return nil
}

// WriteRecordsAppendToCSVFile 将数据写入 CSV 文件
func WriteRecordsAppendToCSVFile(filePath string, headers []string, records [][]string, delimiter rune, appendToFile bool) error {
	var file *os.File
	var err error

	if appendToFile {
		// 追加写入
		file, err = os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
		if err != nil {
			return fmt.Errorf("failed to open file for appending: %v", err)
		}
	} else {
		// 新建文件
		file, err = os.Create(filePath)
		if err != nil {
			return fmt.Errorf("failed to create file: %v", err)
		}
		// 写入表头
		writer := csv.NewWriter(file)
		writer.Comma = delimiter
		defer writer.Flush()

		if err := writer.Write(headers); err != nil {
			return fmt.Errorf("failed to write header to CSV: %v", err)
		}
	}

	defer file.Close()
	writer := csv.NewWriter(file)
	writer.Comma = delimiter
	defer writer.Flush()

	// 写入记录
	if err := writer.WriteAll(records); err != nil {
		return fmt.Errorf("failed to write records to CSV: %v", err)
	}

	return nil
}

// EstimateRecordSize 估算一条记录在 CSV 中的大小
// 参数:
// - record: 要估算的 CSV 记录
// - delimiter: CSV 的分隔符
// 返回值:
// - 记录的估算大小，以字节为单位
func estimateRecordSize(record []string, delimiter rune) int64 {
	size := int64(0)
	for i, field := range record {
		size += int64(len(field))
		// 每个字段之间的分隔符也会占用空间
		if i < len(record)-1 {
			size++
		}
	}
	// 加上换行符
	size++
	return size
}