Merge pull request #2978 from actiontech/2250-cp

feat(utils): add CSVBuilder for CSV file generation with Excel compatibility

Author: ColdWaterLW

sqle/api/controller/v1/audit_plan.go

@@ -1,9 +1,7 @@
 package v1
 
 import (
-	"bytes"
 	"context"
-	"encoding/csv"
 	"fmt"
 	"mime"
 	"net"
@@ -1539,7 +1537,6 @@ func spliceAuditResults(auditResults []model.AuditResult) string {
 // @router /v1/projects/{project_name}/audit_plans/{audit_plan_name}/reports/{audit_plan_report_id}/export [get]
 func ExportAuditPlanReportV1(c echo.Context) error {
 	s := model.GetStorage()
-	buff := new(bytes.Buffer)
 	reportIdStr := c.Param("audit_plan_report_id")
 	auditPlanName := c.Param("audit_plan_name")
 	projectName := c.Param("project_name")
@@ -1548,8 +1545,6 @@ func ExportAuditPlanReportV1(c echo.Context) error {
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
-	csvWriter := csv.NewWriter(buff)
-	buff.WriteString("\xEF\xBB\xBF") // 写入UTF-8 BOM
 	reportInfo, exist, err := s.GetReportWithAuditPlanByReportID(reportId)
 	if !exist {
 		return controller.JSONBaseErrorReq(c, fmt.Errorf("not found audit report"))
@@ -1572,18 +1567,19 @@ func ExportAuditPlanReportV1(c echo.Context) error {
 		{"数据库类型", reportInfo.AuditPlan.DBType},
 		{"审核的数据库", reportInfo.AuditPlan.InstanceDatabase},
 	}
-	err = csvWriter.WriteAll(baseInfo)
+	csvBuilder := utils.NewCSVBuilder()
+	err = csvBuilder.WriteRows(baseInfo)
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
 
 	// Add a split line between report information and sql audit information
-	err = csvWriter.Write([]string{})
+	err = csvBuilder.WriteRow([]string{})
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
 
-	err = csvWriter.Write([]string{"编号", "SQL", "审核结果"})
+	err = csvBuilder.WriteRow([]string{"编号", "SQL", "审核结果"})
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
@@ -1593,17 +1589,16 @@ func ExportAuditPlanReportV1(c echo.Context) error {
 		sqlInfo = append(sqlInfo, []string{strconv.Itoa(idx + 1), sql.SQL, spliceAuditResults(sql.AuditResults)})
 	}
 
-	err = csvWriter.WriteAll(sqlInfo)
+	err = csvBuilder.WriteRows(sqlInfo)
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
 
-	csvWriter.Flush()
 
 	fileName := fmt.Sprintf("扫描任务报告_%s_%s.csv", auditPlanName, time.Now().Format("20060102150405"))
 	c.Response().Header().Set(echo.HeaderContentDisposition, mime.FormatMediaType("attachment", map[string]string{
 		"filename": fileName,
 	}))
 
-	return c.Blob(http.StatusOK, "text/csv", buff.Bytes())
+	return c.Blob(http.StatusOK, "text/csv", csvBuilder.FlushAndGetBuffer().Bytes())
 }

sqle/api/controller/v1/task.go

@@ -1,9 +1,7 @@
 package v1
 
 import (
-	"bytes"
 	"context"
-	"encoding/csv"
 	"fmt"
 	"mime"
 	"mime/multipart"
@@ -488,10 +486,8 @@ func DownloadTaskSQLReportFile(c echo.Context) error {
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, err)
 	}
-	buff := &bytes.Buffer{}
-	buff.WriteString("\xEF\xBB\xBF") // 写入UTF-8 BOM
-	cw := csv.NewWriter(buff)
-	err = cw.Write([]string{"序号", "SQL", "SQL审核状态", "SQL审核结果", "SQL执行状态", "SQL执行结果", "SQL对应的回滚语句", "SQL描述"})
+	csvBuilder := utils.NewCSVBuilder()
+	err = csvBuilder.WriteHeader([]string{"序号", "SQL", "SQL审核状态", "SQL审核结果", "SQL执行状态", "SQL执行结果", "SQL对应的回滚语句", "SQL描述"})
 	if err != nil {
 		return controller.JSONBaseErrorReq(c, errors.New(errors.WriteDataToTheFileError, err))
 	}
@@ -501,7 +497,7 @@ func DownloadTaskSQLReportFile(c echo.Context) error {
 			AuditStatus:  td.AuditStatus,
 		}
 		taskSql.ExecStatus = td.ExecStatus
-		err := cw.Write([]string{
+		err := csvBuilder.WriteRow([]string{
 			strconv.FormatUint(uint64(td.Number), 10),
 			td.ExecSQL,
 			taskSql.GetAuditStatusDesc(),
@@ -515,11 +511,10 @@ func DownloadTaskSQLReportFile(c echo.Context) error {
 			return controller.JSONBaseErrorReq(c, errors.New(errors.WriteDataToTheFileError, err))
 		}
 	}
-	cw.Flush()
 	fileName := fmt.Sprintf("SQL审核报告_%v_%v.csv", task.InstanceName(), taskId)
 	c.Response().Header().Set(echo.HeaderContentDisposition,
 		mime.FormatMediaType("attachment", map[string]string{"filename": fileName}))
-	return c.Blob(http.StatusOK, "text/csv", buff.Bytes())
+	return c.Blob(http.StatusOK, "text/csv", csvBuilder.FlushAndGetBuffer().Bytes())
 }
 
 // @Summary 下载指定扫描任务的SQL文件

sqle/utils/csv_builder.go

@@ -0,0 +1,130 @@
+package utils
+
+import (
+	"bytes"
+	"encoding/csv"
+	"fmt"
+)
+
+/*
+	该文件用于维护SQLE中创建CSV文件的构建器
+*/
+
+const (
+	// MaxNumberPerCell Excel 单元格字符数限制
+	MaxNumberPerCell int = 32767
+)
+
+var ErrCSVColumnCountNotMatch = fmt.Errorf("csv column count not match")
+
+// CSVBuilder 用于构建CSV文件，支持Excel兼容性处理
+//
+// 主要功能：
+//   - 自动添加UTF-8 BOM头
+//   - 处理Excel单元格字符数限制
+//   - 提供简单的API进行CSV文件构建
+//
+// 参考：https://support.microsoft.com/search/results?query=excel-specifications-and-limits
+//
+// 使用示例：
+//
+//	builder := NewCSVBuilder()
+//	builder.WriteHeader([]string{"Name", "Age"}) // 若设置了表头，则每一行的长度需要与表头长度匹配
+//	builder.WriteRows([][]string{{"John", "30"}, {"Alice", "25"}})
+//	buffer := builder.FlushAndGetBuffer()
+type CSVBuilder struct {
+	columnCount uint
+	buffer      *bytes.Buffer
+	csvWriter   *csv.Writer
+}
+
+// NewCSVBuilder 创建并返回一个新的CSVBuilder实例
+//
+// 返回：
+//
+//	*CSVBuilder: 新的CSVBuilder实例
+func NewCSVBuilder() *CSVBuilder {
+	buffer := new(bytes.Buffer)
+	// 写入 UTF-8 BOM 有助于Excel正确识别文件的编码格式
+	buffer.WriteString("\xEF\xBB\xBF")
+	csvWriter := csv.NewWriter(buffer)
+	return &CSVBuilder{
+		buffer:    buffer,
+		csvWriter: csvWriter,
+	}
+}
+
+// WriteHeader 写入CSV文件的表头
+//
+// 参数：
+//
+//	header: 表头字符串数组
+//
+// 返回：
+//
+//	error: 如果写入失败则返回错误
+func (b *CSVBuilder) WriteHeader(header []string) error {
+	b.columnCount = uint(len(header))
+	return b.csvWriter.Write(header)
+}
+
+// WriteRows 写入多行数据到CSV文件
+//
+// 参数：
+//
+//	rows: 二维字符串数组，每行代表一条记录
+//
+// 返回：
+//
+//	error: 如果写入失败则返回错误
+func (b *CSVBuilder) WriteRows(rows [][]string) error {
+	for _, row := range rows {
+		err := b.WriteRow(row)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// WriteRow 写入单行数据到CSV文件
+//
+// 参数：
+//
+//	row: 字符串数组，代表一条记录
+//
+// 返回：
+//
+//	error: 如果列数不匹配或写入失败则返回错误
+func (b *CSVBuilder) WriteRow(row []string) error {
+	// 检查列数是否匹配
+	if b.columnCount > 0 && b.columnCount != uint(len(row)) {
+		return ErrCSVColumnCountNotMatch
+	}
+	// 超过最大字符数，则截断并标记 Excel 单元格
+	for idx, cell := range row {
+		if len(cell) > MaxNumberPerCell {
+			row[idx] = TruncateAndMarkForExcelCell(cell)
+		}
+	}
+	return b.csvWriter.Write(row)
+}
+
+// FlushAndGetBuffer 刷新缓冲区并返回包含CSV内容的缓冲区
+//
+// 返回：
+//
+//	*bytes.Buffer: 包含CSV文件内容的缓冲区
+func (b *CSVBuilder) FlushAndGetBuffer() *bytes.Buffer {
+	b.csvWriter.Flush()
+	return b.buffer
+}
+
+// Error 返回CSV写入过程中遇到的任何错误
+//
+// 返回：
+//
+// 	error: 如果存在错误则返回错误，否则返回nil
+func (b *CSVBuilder) Error() error {
+	return b.csvWriter.Error()
+}

sqle/utils/util.go

@@ -353,3 +353,46 @@ func GenerateRandomString(halfLength int) string {
 	rand.Read(bytes)
 	return fmt.Sprintf("%x", bytes)
 }
+
+// TruncateStringByRunes 按字符数截取字符串
+func TruncateStringByRunes(s string, maxRunes uint) string {
+	// 字节数不大于 maxRunes ，那字符数肯定不大于 maxRunes
+	if uint(len(s)) <= maxRunes {
+		return s
+	}
+
+	// UTF-8一个字符的字节数是不确定的，如：s="a一b二c"，汉字为多字节字符，len(s)=9
+	//    s的hexdump结果：
+	//    00000000  61 e4 b8 80 62 e4 ba 8c 63                       |a...b...c|
+	//
+	//    当想截取头两个字符：“a一”，即 maxRunes 为2时，
+	//    直接返回s[:maxRunes]的话得到是：“61 e4”这两个字节组成的字符串，并非“a一”，“a一”是“61 e4 b8 80”这四个字节，此时应取s[:4]
+	//
+	//    为得到s[:4]中4这个索引，可以“range s”：逐个rune遍历s，i为每个rune起始的字节索引
+	//    i依次为　0　　1　　4　　5　　8
+	//          　^ａ　^一　^ｂ　^二　^ｃ
+	//    遍历 maxRunes (2)次后，i为下一个字符(b)的起始索引，即4，此时s[:i]就是要截取的头两个字符“a一”
+	var runesCount uint
+	for i := range s {
+		if runesCount == maxRunes {
+			// 达到截取的字符数了，将字符截取至此时rune的字节索引
+			return s[:i]
+		}
+		// 未达到要截取的字符数，继续获取下一个rune
+		runesCount++
+	}
+	// 字符串字符数不足 maxRunes
+	return s
+}
+
+const excelCellMaxRunes = 32766
+
+// TruncateAndMarkForExcelCell 对超长字符串进行截取，以符合Excel类工具对单元格字符数上限的限制
+func TruncateAndMarkForExcelCell(s string) string {
+	truncated := TruncateStringByRunes(s, excelCellMaxRunes-4)
+	if truncated != s {
+		// 截取了的话，做标记
+		return truncated + " ..."
+	}
+	return s
+}