MarketingSystemDataExportTool/server/internal/exporter/writer.go

// Package exporter 提供数据导出功能
package exporter

import (
	"bufio"
	"encoding/csv"
	"errors"
	"os"
	"path/filepath"
	"server/internal/constants"
	"time"

	"github.com/xuri/excelize/v2"
)

// ==================== 接口定义 ====================

// RowWriter 行写入器接口
// 所有导出格式必须实现此接口
type RowWriter interface {
	// WriteHeader 写入表头
	WriteHeader(cols []string) error
	// WriteRow 写入数据行
	WriteRow(vals []string) error
	// Close 关闭并返回文件路径、大小
	Close() (path string, size int64, err error)
}

// WriterFactory 写入器工厂函数类型
type WriterFactory func() (RowWriter, error)

// ==================== 工厂方法 ====================

// NewWriter 根据格式创建对应的写入器
func NewWriter(format constants.FileFormat, dir, name string) (RowWriter, error) {
	switch format {
	case constants.FileFormatCSV:
		return NewCSVWriter(dir, name)
	case constants.FileFormatXLSX:
		return NewXLSXWriter(dir, name, "Sheet1")
	default:
		return nil, errors.New("unsupported format: " + string(format))
	}
}

// NewWriterFactory 创建写入器工厂函数
func NewWriterFactory(format constants.FileFormat, dir, name string) WriterFactory {
	return func() (RowWriter, error) {
		return NewWriter(format, dir, name)
	}
}

type CSVWriter struct {
	f     *os.File
	buf   *bufio.Writer
	w     *csv.Writer
	count int64
}

func NewCSVWriter(dir, name string) (*CSVWriter, error) {
	os.MkdirAll(dir, 0755)
	p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".csv")
	f, err := os.Create(p)
	if err != nil {
		return nil, err
	}
	buf := bufio.NewWriterSize(f, 4<<20)
	return &CSVWriter{f: f, buf: buf, w: csv.NewWriter(buf)}, nil
}

func (c *CSVWriter) WriteHeader(cols []string) error {
	if err := c.w.Write(cols); err != nil {
		return err
	}
	c.count++
	return nil
}

func (c *CSVWriter) WriteRow(vals []string) error {
	if err := c.w.Write(vals); err != nil {
		return err
	}
	c.count++
	return nil
}

func (c *CSVWriter) Close() (string, int64, error) {
	c.w.Flush()
	if c.buf != nil {
		_ = c.buf.Flush()
	}
	p := c.f.Name()
	info, _ := c.f.Stat()
	c.f.Close()
	return p, info.Size(), nil
}

type XLSXWriter struct {
	f     *excelize.File
	sheet string
	row   int
	path  string
	sw    *excelize.StreamWriter
}

func NewXLSXWriter(dir, name, sheet string) (*XLSXWriter, error) {
	os.MkdirAll(dir, 0755)
	p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".xlsx")
	f := excelize.NewFile()
	idx, err := f.GetSheetIndex(sheet)
	if err != nil || idx < 0 {
		idx, _ = f.NewSheet(sheet)
		f.SetActiveSheet(idx)
		if sheet != "Sheet1" {
			_ = f.DeleteSheet("Sheet1")
		}
	} else {
		f.SetActiveSheet(idx)
	}
	sw, e := f.NewStreamWriter(sheet)
	if e != nil {
		return nil, e
	}
	return &XLSXWriter{f: f, sheet: sheet, row: 1, path: p, sw: sw}, nil
}

func (x *XLSXWriter) WriteHeader(cols []string) error {
	vals := make([]interface{}, len(cols))
	for i := range cols {
		vals[i] = cols[i]
	}
	axis := "A" + itoa(1)
	if err := x.sw.SetRow(axis, vals); err != nil {
		return err
	}
	x.row = 2
	return nil
}

func (x *XLSXWriter) WriteRow(vals []string) error {
	rowVals := make([]interface{}, len(vals))
	for i := range vals {
		rowVals[i] = vals[i]
	}
	axis := "A" + itoa(x.row)
	if err := x.sw.SetRow(axis, rowVals); err != nil {
		return err
	}
	x.row++
	return nil
}

func (x *XLSXWriter) Close() (string, int64, error) {
	if x.sw != nil {
		_ = x.sw.Flush()
	}
	if err := x.f.SaveAs(x.path); err != nil {
		return "", 0, err
	}
	info, err := os.Stat(x.path)
	if err != nil {
		return x.path, 0, nil
	}
	return x.path, info.Size(), nil
}

func col(n int) string {
	s := ""
	for n > 0 {
		n--
		s = string(rune('A'+(n%26))) + s
		n /= 26
	}
	return s
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	b := make([]byte, 0, 10)
	m := n
	for m > 0 {
		b = append(b, byte('0'+(m%10)))
		m /= 10
	}
	for i, j := 0, len(b)-1; i < j; i, j = i+1, j-1 {
		b[i], b[j] = b[j], b[i]
	}
	return string(b)
}