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("\xEF\xBB\xBF")
        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
}