123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210 |
- package csvwriter
- import (
- "encoding/csv"
- "fmt"
- "os"
- "path/filepath"
- "strconv"
- )
- 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
- }
- 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 := 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
- defer writer.Flush()
-
- 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
- defer writer.Flush()
- }
-
- 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
- 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
- }
- 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
- }
- 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
- }
|