export statements

pkg/apiserver/statement/statement.go

@@ -14,11 +14,24 @@
 package statement
 
 import (
+	"bufio"
+	"bytes"
+	"encoding/csv"
+	"errors"
+	"fmt"
+	"io/ioutil"
 	"net/http"
+	"os"
+	"reflect"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/gtank/cryptopasta"
+	"github.com/pingcap/log"
 	"go.uber.org/fx"
+	"go.uber.org/zap"
 
 	"github.com/pingcap-incubator/tidb-dashboard/pkg/apiserver/user"
 	"github.com/pingcap-incubator/tidb-dashboard/pkg/apiserver/utils"
@@ -40,15 +53,23 @@ func NewService(p ServiceParams) *Service {
 
 func Register(r *gin.RouterGroup, auth *user.AuthService, s *Service) {
 	endpoint := r.Group("/statements")
-	endpoint.Use(auth.MWAuthRequired())
-	endpoint.Use(utils.MWConnectTiDB(s.params.TiDBClient))
-	endpoint.GET("/config", s.configHandler)
-	endpoint.POST("/config", s.modifyConfigHandler)
-	endpoint.GET("/time_ranges", s.timeRangesHandler)
-	endpoint.GET("/stmt_types", s.stmtTypesHandler)
-	endpoint.GET("/overviews", s.overviewsHandler)
-	endpoint.GET("/plans", s.getPlansHandler)
-	endpoint.GET("/plan/detail", s.getPlanDetailHandler)
+	{
+		endpoint.GET("/download", s.downloadHandler)
+
+		endpoint.Use(auth.MWAuthRequired())
+		endpoint.Use(utils.MWConnectTiDB(s.params.TiDBClient))
+		{
+			endpoint.GET("/config", s.configHandler)
+			endpoint.POST("/config", s.modifyConfigHandler)
+			endpoint.GET("/time_ranges", s.timeRangesHandler)
+			endpoint.GET("/stmt_types", s.stmtTypesHandler)
+			endpoint.GET("/overviews", s.overviewsHandler)
+			endpoint.GET("/plans", s.getPlansHandler)
+			endpoint.GET("/plan/detail", s.getPlanDetailHandler)
+
+			endpoint.POST("/download/token", s.downloadTokenHandler)
+		}
+	}
 }
 
 // @Summary Get statement configurations
@@ -210,3 +231,136 @@ func (s *Service) getPlanDetailHandler(c *gin.Context) {
 	}
 	c.JSON(http.StatusOK, result)
 }
+
+// @Router /statements/download/token [post]
+// @Summary Generate a download token for exported statements
+// @Produce plain
+// @Param request body GetStatementsRequest true "Request body"
+// @Success 200 {string} string "xxx"
+// @Security JwtAuth
+// @Failure 401 {object} utils.APIError "Unauthorized failure"
+func (s *Service) downloadTokenHandler(c *gin.Context) {
+	var req GetStatementsRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		utils.MakeInvalidRequestErrorFromError(c, err)
+		return
+	}
+	db := utils.GetTiDBConnection(c)
+	fields := []string{}
+	if strings.TrimSpace(req.Fields) != "" {
+		fields = strings.Split(req.Fields, ",")
+	}
+	overviews, err := QueryStatementsOverview(
+		db,
+		req.BeginTime, req.EndTime,
+		req.Schemas,
+		req.StmtTypes,
+		req.Text,
+		fields)
+	if err != nil {
+		_ = c.Error(err)
+		return
+	}
+	if len(overviews) == 0 {
+		utils.MakeInvalidRequestErrorFromError(c, errors.New("no data to export"))
+		return
+	}
+
+	// convert data
+	fieldsMap := make(map[string]string)
+	t := reflect.TypeOf(overviews[0])
+	fieldsNum := t.NumField()
+	for i := 0; i < fieldsNum; i++ {
+		field := t.Field(i)
+		fieldsMap[strings.ToLower(field.Tag.Get("json"))] = field.Name
+	}
+
+	csvData := [][]string{fields}
+	for _, overview := range overviews {
+		row := []string{}
+		v := reflect.ValueOf(overview)
+		for _, field := range fields {
+			filedName := fieldsMap[field]
+			s := v.FieldByName(filedName)
+			var val string
+			switch s.Interface().(type) {
+			case int:
+				val = strconv.FormatInt(s.Int(), 10)
+			default:
+				val = s.String()
+			}
+			row = append(row, val)
+		}
+		csvData = append(csvData, row)
+	}
+
+	// generate temp file that persist encrypted data
+	timeLayout := "01021504"
+	beginTime := time.Unix(int64(req.BeginTime), 0).Format(timeLayout)
+	endTime := time.Unix(int64(req.EndTime), 0).Format(timeLayout)
+	csvFile, err := ioutil.TempFile("", fmt.Sprintf("statements_%s_%s_*.csv", beginTime, endTime))
+	if err != nil {
+		_ = c.Error(err)
+		return
+	}
+	defer csvFile.Close()
+	// generate buf writer for tmpfile
+	csvBuf := bufio.NewWriter(csvFile)
+	// generate encrypted key
+	secretKey := cryptopasta.NewEncryptionKey()
+
+	rowBuf := bytes.NewBuffer(nil)
+	csvwriter := csv.NewWriter(rowBuf)
+	for _, csvRow := range csvData {
+		_ = csvwriter.Write(csvRow)
+		csvwriter.Flush()
+		row := make([]byte, len(rowBuf.String()))
+		_, _ = rowBuf.Read(row)
+
+		encrypted, err := cryptopasta.Encrypt(row, secretKey)
+		if err != nil {
+			_ = c.Error(err)
+			return
+		}
+		csvBuf.Write(encrypted)
+	}
+	csvBuf.Flush()
+
+	// zip file and encrypt
+
+	// generate token by filepath
+	token, err := utils.NewJWTString("statements/download", csvFile.Name())
+	if err != nil {
+		_ = c.Error(err)
+		return
+	}
+	c.String(http.StatusOK, token)
+}
+
+// @Router /statements/download [get]
+// @Summary Download statements
+// @Produce application/zip
+// @Param token query string true "download token"
+// @Failure 400 {object} utils.APIError
+// @Failure 401 {object} utils.APIError "Unauthorized failure"
+func (s *Service) downloadHandler(c *gin.Context) {
+	token := c.Query("token")
+	fielPath, err := utils.ParseJWTString("statements/download", token)
+	if err != nil {
+		utils.MakeInvalidRequestErrorFromError(c, err)
+		return
+	}
+
+	_, err = os.Stat(fielPath)
+	if err != nil {
+		_ = c.Error(err)
+		return
+	}
+
+	c.Writer.Header().Set("Content-type", "application/octet-stream")
+	c.Writer.Header().Set("Content-Disposition", "attachment; filename=\"statements.zip\"")
+	err = utils.StreamZipPack(c.Writer, []string{fielPath}, true)
+	if err != nil {
+		log.Error("Stream zip pack failed", zap.Error(err))
+	}
+}

ui/lib/apps/Statement/pages/List/index.tsx

@@ -17,6 +17,7 @@ import TimeRangeSelector from './TimeRangeSelector'
 import useStatementTableController, {
   DEF_STMT_COLUMN_KEYS,
 } from '../../utils/useStatementTableController'
+import client from '@lib/client'
 
 const { Search } = Input
 
@@ -47,8 +48,19 @@ export default function StatementsOverview() {
     allStmtTypes,
     loadingStatements,
     tableColumns,
+
+    getDownloadToken,
+    downloading,
   } = controller
 
+  async function exportCSV() {
+    const token = await getDownloadToken()
+    if (token) {
+      const url = `${client.getBasePath()}/statements/download?token=${token}`
+      window.open(url)
+    }
+  }
+
   return (
     <div style={{ height: '100vh', display: 'flex', flexDirection: 'column' }}>
       <Card>
@@ -106,6 +118,9 @@ export default function StatementsOverview() {
                 setQueryOptions({ ...queryOptions, searchText })
               }
             />
+            <Button type="primary" onClick={exportCSV} loading={downloading}>
+              {t('statement.pages.overview.toolbar.export')}
+            </Button>
           </Space>
 
           <Space>

ui/lib/apps/Statement/translations/en.yaml

@@ -37,6 +37,7 @@ statement:
           recent: Recent
           usual_time_ranges: Common
           custom_time_ranges: Custom
+        export: export
   settings:
     title: Settings
     disabled_result:

ui/lib/apps/Statement/translations/zh.yaml

@@ -37,6 +37,7 @@ statement:
           recent: 最近
           usual_time_ranges: 常用时间范围
           custom_time_ranges: 自定义时间范围
+        export: 导出
   settings:
     title: 设置
     disabled_result:

ui/lib/apps/Statement/utils/useStatementTableController.ts

@@ -67,6 +67,9 @@ export interface IStatementTableController {
 
   tableColumns: IColumn[]
   visibleColumnKeys: IColumnKeys
+
+  getDownloadToken: () => Promise<string>
+  downloading: boolean
 }
 
 export default function useStatementTableController(
@@ -218,6 +221,26 @@ export default function useStatementTableController(
     queryStatementList()
   }, [queryOptions, allTimeRanges, validTimeRange, selectedFields])
 
+  const [downloading, setDownloading] = useState(false)
+  async function getDownloadToken() {
+    let token = ''
+    try {
+      setDownloading(true)
+      const res = await client.getInstance().statementsDownloadTokenPost({
+        begin_time: validTimeRange.begin_time,
+        end_time: validTimeRange.end_time,
+        fields: selectedFields,
+        schemas: queryOptions.schemas,
+        stmt_types: queryOptions.stmtTypes,
+        text: queryOptions.searchText,
+      })
+      token = res.data
+    } finally {
+      setDownloading(false)
+    }
+    return token
+  }
+
   return {
     queryOptions,
     setQueryOptions,
@@ -237,5 +260,8 @@ export default function useStatementTableController(
 
     tableColumns,
     visibleColumnKeys,
+
+    getDownloadToken,
+    downloading,
   }
 }