github.com/bosssauce/ponzu@v0.11.1-0.20200102001432-9bc41b703131/system/admin/export.go (about)

     1  package admin
     2  
     3  import (
     4  	"encoding/csv"
     5  	"fmt"
     6  	"io/ioutil"
     7  	"log"
     8  	"net/http"
     9  	"os"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/ponzu-cms/ponzu/management/format"
    14  	"github.com/ponzu-cms/ponzu/system/db"
    15  	"github.com/ponzu-cms/ponzu/system/item"
    16  
    17  	"github.com/tidwall/gjson"
    18  )
    19  
    20  func exportHandler(res http.ResponseWriter, req *http.Request) {
    21  	// /admin/contents/export?type=Blogpost&format=csv
    22  	q := req.URL.Query()
    23  	t := q.Get("type")
    24  	f := strings.ToLower(q.Get("format"))
    25  
    26  	if t == "" || f == "" {
    27  		v, err := Error400()
    28  		if err != nil {
    29  			res.WriteHeader(http.StatusInternalServerError)
    30  			return
    31  		}
    32  
    33  		res.WriteHeader(http.StatusBadRequest)
    34  		_, err = res.Write(v)
    35  		if err != nil {
    36  			res.WriteHeader(http.StatusInternalServerError)
    37  			return
    38  		}
    39  
    40  	}
    41  
    42  	pt, ok := item.Types[t]
    43  	if !ok {
    44  		res.WriteHeader(http.StatusBadRequest)
    45  		return
    46  	}
    47  
    48  	switch f {
    49  	case "csv":
    50  		csv, ok := pt().(format.CSVFormattable)
    51  		if !ok {
    52  			res.WriteHeader(http.StatusBadRequest)
    53  			return
    54  		}
    55  
    56  		fields := csv.FormatCSV()
    57  		exportCSV(res, req, pt, fields)
    58  
    59  	default:
    60  		res.WriteHeader(http.StatusBadRequest)
    61  		return
    62  	}
    63  }
    64  
    65  func exportCSV(res http.ResponseWriter, req *http.Request, pt func() interface{}, fields []string) {
    66  	tmpFile, err := ioutil.TempFile(os.TempDir(), "exportcsv-")
    67  	if err != nil {
    68  		log.Println("Failed to create tmp file for CSV export:", err)
    69  		res.WriteHeader(http.StatusInternalServerError)
    70  		return
    71  	}
    72  
    73  	err = os.Chmod(tmpFile.Name(), 0666)
    74  	if err != nil {
    75  		log.Println("chmod err:", err)
    76  	}
    77  
    78  	csvBuf := csv.NewWriter(tmpFile)
    79  
    80  	t := req.URL.Query().Get("type")
    81  
    82  	// get content data and loop through creating a csv row per result
    83  	bb := db.ContentAll(t)
    84  
    85  	// add field names to first row
    86  	err = csvBuf.Write(fields)
    87  	if err != nil {
    88  		res.WriteHeader(http.StatusInternalServerError)
    89  		log.Println("Failed to write column headers:", fields)
    90  		return
    91  	}
    92  
    93  	for row := range bb {
    94  		// unmarshal data and loop over fields
    95  		rowBuf := []string{}
    96  
    97  		for _, col := range fields {
    98  			// pull out each field as the column value
    99  			result := gjson.GetBytes(bb[row], col)
   100  
   101  			// append it to the buffer
   102  			rowBuf = append(rowBuf, result.String())
   103  		}
   104  
   105  		// write row to csv
   106  		err := csvBuf.Write(rowBuf)
   107  		if err != nil {
   108  			res.WriteHeader(http.StatusInternalServerError)
   109  			log.Println("Failed to write column headers:", fields)
   110  			return
   111  		}
   112  	}
   113  
   114  	csvBuf.Flush()
   115  
   116  	// write the buffer to a content-disposition response
   117  	fi, err := tmpFile.Stat()
   118  	if err != nil {
   119  		log.Println("Failed to read tmp file info for CSV export:", err)
   120  		res.WriteHeader(http.StatusInternalServerError)
   121  		return
   122  	}
   123  
   124  	err = tmpFile.Close()
   125  	if err != nil {
   126  		log.Println("Failed to close tmp file for CSV export:", err)
   127  	}
   128  
   129  	ts := time.Now().Unix()
   130  	disposition := `attachment; filename="export-%s-%d.csv"`
   131  
   132  	res.Header().Set("Content-Type", "text/csv")
   133  	res.Header().Set("Content-Disposition", fmt.Sprintf(disposition, t, ts))
   134  	res.Header().Set("Content-Length", fmt.Sprintf("%d", int(fi.Size())))
   135  
   136  	http.ServeFile(res, req, tmpFile.Name())
   137  
   138  	err = os.Remove(tmpFile.Name())
   139  	if err != nil {
   140  		log.Println("Failed to remove tmp file for CSV export:", err)
   141  	}
   142  }