csvwriter.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. package csvwriter
  2. import (
  3. "encoding/csv"
  4. "fmt"
  5. "os"
  6. "path/filepath"
  7. "strconv"
  8. )
  9. // WriteRecordsToCSV 将数据写入 CSV 文件
  10. // 参数:
  11. // - filePath: CSV 文件路径
  12. // - headers: 表头字段列表
  13. // - records: 要写入的记录列表,每个记录是一个字符串切片
  14. // - delimiter: CSV 文件的分隔符
  15. // - appendToFile: 是否追加写入,true 表示在文件末尾追加,false 表示新建文件
  16. // 返回值:
  17. // - error: 如果发生错误,返回错误信息
  18. func WriteRecordsToCSV(filePath string, headers []string, records [][]string, delimiter rune, appendToFile bool) error {
  19. var file *os.File
  20. var err error
  21. if appendToFile {
  22. // 追加写入
  23. file, err = os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, 0644)
  24. if err != nil {
  25. return fmt.Errorf("failed to open file for appending: %v", err)
  26. }
  27. } else {
  28. // 新建文件
  29. file, err = os.Create(filePath)
  30. if err != nil {
  31. return fmt.Errorf("failed to create file: %v", err)
  32. }
  33. }
  34. defer file.Close()
  35. writer := csv.NewWriter(file)
  36. writer.Comma = delimiter
  37. defer writer.Flush()
  38. // 如果是新文件且有表头,需要写入表头
  39. if !appendToFile && headers != nil {
  40. if err := writer.Write(headers); err != nil {
  41. return fmt.Errorf("failed to write headers to CSV: %v", err)
  42. }
  43. }
  44. // 写入记录
  45. if err := writer.WriteAll(records); err != nil {
  46. return fmt.Errorf("failed to write records to CSV: %v", err)
  47. }
  48. return nil
  49. }
  50. /*
  51. WriteRecordsToCSVWithFileSizeLimit 将数据写入 CSV 文件,并基于文件大小进行分片
  52. 参数:
  53. - baseFilePath: CSV 文件的基本路径,用于生成不同分片的文件名
  54. - headers: 表的列名,用于 CSV 表头
  55. - records: 要写入的表记录
  56. - delimiter: CSV 的分隔符
  57. - maxFileSizeMB: 单个 CSV 文件的最大大小 MB
  58. - fileIndex: 当前文件的索引号,用于管理分片文件的命名
  59. - currentFileSize: 当前文件的大小,用于判断是否需要创建新文件
  60. - currentFilePath: 当前正在写入的文件路径,用于管理文件的创建和打开
  61. 返回值:
  62. - error: 如果写入过程中有错误,返回相关的错误信息
  63. */
  64. func WriteRecordsToCSVWithFileSizeLimit(baseFilePath string, headers []string, records [][]string, delimiter rune, maxFileSizeMB int64, fileIndex *int, currentFileSize *int64, currentFilePath *string) error {
  65. var file *os.File
  66. var writer *csv.Writer
  67. // 将 maxFileSize 从 MB 转换为字节
  68. maxFileSize := maxFileSizeMB * 1024 * 1024
  69. // 检查是否需要创建新文件
  70. if *currentFileSize == 0 {
  71. // 如果是第一次写入,创建新文件
  72. *currentFilePath = filepath.Join(baseFilePath + "_" + strconv.Itoa(*fileIndex) + ".csv")
  73. var err error
  74. file, err = os.Create(*currentFilePath)
  75. if err != nil {
  76. return fmt.Errorf("failed to create file: %v", err)
  77. }
  78. defer file.Close()
  79. writer = csv.NewWriter(file)
  80. writer.Comma = delimiter // 设置 CSV 分隔符
  81. defer writer.Flush()
  82. // 写入表头到 CSV 文件
  83. if err := writer.Write(headers); err != nil {
  84. return fmt.Errorf("failed to write header to CSV: %v", err)
  85. }
  86. *currentFileSize += estimateRecordSize(headers, delimiter)
  87. } else {
  88. // 如果文件已存在,追加写入
  89. var err error
  90. file, err = os.OpenFile(*currentFilePath, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
  91. if err != nil {
  92. return fmt.Errorf("failed to open file: %v", err)
  93. }
  94. defer file.Close()
  95. writer = csv.NewWriter(file)
  96. writer.Comma = delimiter // 设置 CSV 分隔符
  97. defer writer.Flush()
  98. }
  99. // 遍历记录逐行写入 CSV 文件
  100. for _, record := range records {
  101. // 如果文件大小超过限制,创建新文件
  102. if *currentFileSize > maxFileSize {
  103. writer.Flush()
  104. file.Close()
  105. // 增加文件索引,创建新的文件
  106. *fileIndex++
  107. *currentFilePath = filepath.Join(baseFilePath + "_" + strconv.Itoa(*fileIndex) + ".csv")
  108. var err error
  109. file, err = os.Create(*currentFilePath)
  110. if err != nil {
  111. return fmt.Errorf("failed to create file: %v", err)
  112. }
  113. writer = csv.NewWriter(file)
  114. writer.Comma = delimiter // 设置 CSV 分隔符
  115. defer writer.Flush()
  116. // 写入表头到新文件
  117. if err := writer.Write(headers); err != nil {
  118. return fmt.Errorf("failed to write header to new CSV: %v", err)
  119. }
  120. *currentFileSize = estimateRecordSize(headers, delimiter)
  121. }
  122. // 写入当前记录
  123. if err := writer.Write(record); err != nil {
  124. return fmt.Errorf("failed to write row to CSV: %v", err)
  125. }
  126. // 更新当前文件大小
  127. *currentFileSize += estimateRecordSize(record, delimiter)
  128. }
  129. return nil
  130. }
  131. // WriteRecordsAppendToCSVFile 将数据写入 CSV 文件
  132. func WriteRecordsAppendToCSVFile(filePath string, headers []string, records [][]string, delimiter rune, appendToFile bool) error {
  133. var file *os.File
  134. var err error
  135. if appendToFile {
  136. // 追加写入
  137. file, err = os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
  138. if err != nil {
  139. return fmt.Errorf("failed to open file for appending: %v", err)
  140. }
  141. } else {
  142. // 新建文件
  143. file, err = os.Create(filePath)
  144. if err != nil {
  145. return fmt.Errorf("failed to create file: %v", err)
  146. }
  147. // 写入表头
  148. writer := csv.NewWriter(file)
  149. writer.Comma = delimiter
  150. defer writer.Flush()
  151. if err := writer.Write(headers); err != nil {
  152. return fmt.Errorf("failed to write header to CSV: %v", err)
  153. }
  154. }
  155. defer file.Close()
  156. writer := csv.NewWriter(file)
  157. writer.Comma = delimiter
  158. defer writer.Flush()
  159. // 写入记录
  160. if err := writer.WriteAll(records); err != nil {
  161. return fmt.Errorf("failed to write records to CSV: %v", err)
  162. }
  163. return nil
  164. }
  165. // EstimateRecordSize 估算一条记录在 CSV 中的大小
  166. // 参数:
  167. // - record: 要估算的 CSV 记录
  168. // - delimiter: CSV 的分隔符
  169. // 返回值:
  170. // - 记录的估算大小,以字节为单位
  171. func estimateRecordSize(record []string, delimiter rune) int64 {
  172. size := int64(0)
  173. for i, field := range record {
  174. size += int64(len(field))
  175. // 每个字段之间的分隔符也会占用空间
  176. if i < len(record)-1 {
  177. size++
  178. }
  179. }
  180. // 加上换行符
  181. size++
  182. return size
  183. }