HTTP导出和上传文件处理

需求

• 导出节目单[不保存文件到本地]
• 输出不同的文件类型[csv,xls,xlsx]
• 上传[csv,xls,xlsx]导入节目单

导出节目单

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"net/http"
	"strings"

	"github.com/tealeg/xlsx"
)

//导出节目单的handler
func Export(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	format := strings.ToLower(q.Get("format"))
	// 获取节目单
	ps := dbqueryprograms()

	// 设置文件头
	w.Header().Set("ContentType", "application/octet-stream")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment;filename=节目单.%s", format))
	if format == "csv" {
		b := &bytes.Buffer{}
		wr := csv.NewWriter(b)
		wr.WriteString("xEFxBBxBF")
		wr.Write([]string{"节目名称", "开始时间", "结束时间", "播出类型", "素材ID"})
		for _, p := range ps {
			wr.Write([]string{p.Name, p.StartTime, p.EndTime, p.Type, p.VideoId})
		}

		wr.Flush()
		w.Write(b.Bytes())
		return
	} else if format == "xls" || format == "xlsx" {
		file := xlsx.NewFile()
		sheet, err := file.AddSheet("sheet1")
		if err != nil {
			DoHttpError(w, "文件创建失败")
			return
		}

		title := &rowContent{"节目名称", "开始时间", "结束时间", "播出类型", "素材ID"}
		writeRow(sheet, title)

		for _, p := range ps {
			line := &rowContent{
				Name:        p.Name,
				Start:       p.StartTime,
				End:         p.EndTime,
				PlayingType: p.Type,
				VideoID:     p.VideoId,
			}
			writeRow(sheet, line)
		}
		// 直接输出到response
		file.Write(w)
		return
	}

	DoHttpError(w, "不支持的文件类型")
	return
}

func writeRow(sheet *xlsx.Sheet, c *rowContent) {
	row := sheet.AddRow()
	nameCell := row.AddCell()
	nameCell.Value = c.Name

	startCell := row.AddCell()
	startCell.Value = c.Start

	endCell := row.AddCell()
	endCell.Value = c.End

	typeCell := row.AddCell()
	typeCell.Value = c.PlayingType

	videoCell := row.AddCell()
	videoCell.Value = c.VideoID
}

type rowContent struct {
	Name, Start, End, PlayingType, VideoID string
}

文件上传

import (
	"encoding/csv"
	"io"
	"mime/multipart"
	"net/http"
	"path"
	"strings"

	"github.com/pkg/errors"
	"github.com/tealeg/xlsx"
)

func ImportHandler(w http.ResponseWriter, r *http.Request) {
	file, handler, err := r.FormFile("file")
	if err != nil {
		DoHttpError(w, "获取上传文件失败")
		return
	}
	defer file.Close()

	if handler.Size > 100*1024*1024 {
		DoHttpError(w, "上传不能超过100M")
		return
	}

	fileExt := strings.ToLower(path.Ext(handler.Filename))
	var content [][]string
	if fileExt == ".csv" {
		content, err = parseCsv(file)
		if err != nil {
			// handle err and log
			return
		}
	} else if fileExt == ".xls" || fileExt == ".xlsx" {
		content, err = parseXlsx(file, handler.Size)
		if err != nil {
			// handle err and log
			return
		}
	} else {
		DoHttpError(w, "不支持的类型")
		return
	}

	// db and handle err log
	// err := DbCreatePrograms(content)
	DoHttpSuccess(w, "解析成功", nil)
	return
}

// 解析csv文件
func parseCsv(file multipart.File) ([][]string, error) {
	result := [][]string{}
	r := csv.NewReader(file)
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, errors.Wrap(err, "Read csv failed")
		}

		var st = []string{}
		for i, j := range record {
			// 头信息过滤掉
			if i == 0 {
				continue
			}
			st = append(st, j)
		}
		result = append(result, st)
	}
	return result, nil
}

// 解析xlsx文件
func parseXlsx(file multipart.File, size int64) ([][]string, error) {
	result := [][]string{}
	xlFile, err := xlsx.OpenReaderAt(file, size)
	if err != nil {
		return nil, errors.Wrap(err, "Read file failed;")
	}

	for _, sheet := range xlFile.Sheets {
		// 不读隐藏的工作区
		if sheet.Hidden {
			continue
		}

		for i, row := range sheet.Rows {
			// 过滤头信息
			if i == 0 {
				continue
			}
			var st = []string{}
			for j, cell := range row.Cells {
				c, err := cell.FormattedValue()
				if err != nil {
					return nil, errors.Wrapf(err, "Get (%v, %v) value failed;", i, j)
				}
				st = append(st, c)
			}
			result = append(result, st)
		}
	}
	return result, nil
}