Repository: pingcap/dumpling
Pull request #55 (2020-04-16 05:05:24)

writer: Implement csv writer
Implement basic csv writer
issue #6 

codecov-io: # [Codecov](https://codecov.io/gh/pingcap/dumpling/pull/55?src=pr&el=h1) Report
> Merging [#55](https://codecov.io/gh/pingcap/dumpling/pull/55?src=pr&el=desc) into [master](https://codecov.io/gh/pingcap/dumpling/commit/a4687ce515b12d06f6484fee2758238253bc0286&el=desc) will **decrease** coverage by `4.79%`.
> The diff coverage is `1.61%`.


```diff
@@            Coverage Diff             @@
##           master      #55      +/-   ##
==========================================
- Coverage   59.97%   55.17%   -4.80%     
==========================================
  Files          16       16              
  Lines        1374     1497     +123     
==========================================
+ Hits          824      826       +2     
- Misses        488      609     +121     
  Partials       62       62              
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index 7f22101..0473e2c 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -41,6 +41,8 @@ var (
 	statusAddr    string
 	rows          uint64
 	where         string
+	fileType      string
+	noHeader      bool
 
 	escapeBackslash bool
 
@@ -78,6 +80,8 @@ func init() {
 	rootCmd.PersistentFlags().Uint64VarP(&rows, "rows", "r", export.UnspecifiedSize, "Split table into chunks of this many rows, default unlimited")
 	rootCmd.PersistentFlags().StringVar(&where, "where", "", "Dump only selected records")
 	rootCmd.PersistentFlags().BoolVar(&escapeBackslash, "escape-backslash", true, "use backslash to escape quotation marks")
+	rootCmd.PersistentFlags().StringVar(&fileType, "filetype", "sql", "The type of export file (sql/csv)")
+	rootCmd.PersistentFlags().BoolVar(&noHeader, "no-header", false, "whether not to dump CSV table header")
 }
 
 func run() {
@@ -100,6 +104,8 @@ func run() {
 	conf.Where = where
 	conf.EscapeBackslash = escapeBackslash
 	conf.LogLevel = logLevel
+	conf.FileType = fileType
+	conf.NoHeader = noHeader
 
 	err := export.Dump(conf)
 	if err != nil {
diff --git a/v4/export/config.go b/v4/export/config.go
index 0206a17..0d723e9 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -31,10 +31,12 @@ type Config struct {
 	Snapshot      string
 	Consistency   string
 	NoViews       bool
+	NoHeader      bool
 
 	BlackWhiteList  BWListConf
 	Rows            uint64
 	Where           string
+	FileType        string
 	EscapeBackslash bool
 }
 
@@ -59,6 +61,8 @@ func DefaultConfig() *Config {
 		NoViews:       true,
 		Rows:          UnspecifiedSize,
 		Where:         "",
+		FileType:      "SQL",
+		NoHeader:      false,
 	}
 }
 
diff --git a/v4/export/dump.go b/v4/export/dump.go
index dddee86..4cd3d64 100644
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -3,6 +3,7 @@ package export
 import (
 	"context"
 	"database/sql"
+	"strings"
 	"time"
 
 	"github.com/pingcap/dumpling/v4/log"
@@ -76,7 +77,13 @@ func Dump(conf *Config) (err error) {
 		log.Info("get global metadata failed", zap.Error(err))
 	}
 
-	writer, err := NewSimpleWriter(conf)
+	var writer Writer
+	switch strings.ToLower(conf.FileType) {
+	case "sql":
+		writer, err = NewSimpleWriter(conf)
+	case "csv":
+		writer, err = NewCsvWriter(conf)
+	}
 	if err != nil {
 		return err
 	}
diff --git a/v4/export/ir.go b/v4/export/ir.go
index db7a0c4..928eff8 100644
--- a/v4/export/ir.go
+++ b/v4/export/ir.go
@@ -12,6 +12,7 @@ type TableDataIR interface {
 	ChunkIndex() int
 	ColumnCount() uint
 	ColumnTypes() []string
+	ColumnNames() []string
 	SelectedField() string
 	EscapeBackSlash() bool
 
@@ -37,6 +38,7 @@ type RowReceiverStringer interface {
 
 type Stringer interface {
 	WriteToBuffer(*bytes.Buffer, bool)
+	WriteToBufferInCsv(*bytes.Buffer, bool)
 }
 
 type RowReceiver interface {
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index bd09804..e2b21fb 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -154,6 +154,14 @@ func (td *tableData) ColumnTypes() []string {
 	return colTypes
 }
 
+func (td *tableData) ColumnNames() []string {
+	colNames := make([]string, len(td.colTypes))
+	for i, ct := range td.colTypes {
+		colNames[i] = ct.Name()
+	}
+	return colNames
+}
+
 func (td *tableData) DatabaseName() string {
 	return td.database
 }
diff --git a/v4/export/sql.go b/v4/export/sql.go
index b98086f..998f3da 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -401,14 +401,14 @@ func estimateCount(dbName, tableName string, db *sql.DB, field string, conf *Con
 	var estRows int
 	if conf.ServerInfo.ServerType == ServerTypeTiDB {
 		/* tidb results field name is estrows
-		+-----------------------+----------+-----------+-----------------------------------------+
-		| id                    | estrows  | task      | operator info                           |
-		+-----------------------+----------+-----------+-----------------------------------------+
-		| tablereader_5         | 10000.00 | root      | data:tablefullscan_4                    |
-		| └─tablefullscan_4     | 10000.00 | cop[tikv] | table:a, keep order:false, stats:pseudo |
-		+-----------------------+----------+-----------+----------------------------------------
+		+-----------------------+----------+-----------+---------------------------------------------------------+
+		| id                    | estrows  | task      | access object | operator info                           |
+		+-----------------------+----------+-----------+---------------------------------------------------------+
+		| tablereader_5         | 10000.00 | root      |               | data:tablefullscan_4                    |
+		| └─tablefullscan_4     | 10000.00 | cop[tikv] | table:a       | table:a, keep order:false, stats:pseudo |
+		+-----------------------+----------+-----------+----------------------------------------------------------
 		*/
-		estRows = detectEstimateRows(db, query, "tidb", 1, 4)
+		estRows = detectEstimateRows(db, query, "tidb", "estrows")
 	} else if conf.ServerInfo.ServerType == ServerTypeMariaDB {
 		/* mariadb result field name is rows
 		+------+-------------+---------+-------+---------------+------+---------+------+----------+-------------+
@@ -417,7 +417,7 @@ func estimateCount(dbName, tableName string, db *sql.DB, field string, conf *Con
 		|    1 | SIMPLE      | sbtest1 | index | NULL          | k_1  | 4       | NULL | 15000049 | Using index |
 		+------+-------------+---------+-------+---------------+------+---------+------+----------+-------------+
 		*/
-		estRows = detectEstimateRows(db, query, "mariadb", 8, 10)
+		estRows = detectEstimateRows(db, query, "mariadb", "rows")
 	} else if conf.ServerInfo.ServerType == ServerTypeMySQL {
 		/* mysql result field name is rows
 		+----+-------------+-------+------------+-------+---------------+-----------+---------+------+------+----------+-------------+
@@ -426,7 +426,7 @@ func estimateCount(dbName, tableName string, db *sql.DB, field string, conf *Con
 		|  1 | SIMPLE      | t1    | NULL       | index | NULL          | multi_col | 10      | NULL |    5 |   100.00 | Using index |
 		+----+-------------+-------+------------+-------+---------------+-----------+---------+------+------+----------+-------------+
 		*/
-		estRows = detectEstimateRows(db, query, "mysql", 9, 12)
+		estRows = detectEstimateRows(db, query, "mysql", "rows")
 	}
 	if estRows > 0 {
 		return uint64(estRows)
@@ -434,20 +434,32 @@ func estimateCount(dbName, tableName string, db *sql.DB, field string, conf *Con
 	return 0
 }
 
-func detectEstimateRows(db *sql.DB, query string, dbType string, fieldIndex int, fieldLength int) int {
-	oneRow := make([]sql.NullString, fieldLength)
-	addr := make([]interface{}, fieldLength)
+func detectEstimateRows(db *sql.DB, query string, dbType string, fieldName string) int {
+	row, err := db.Query(query)
+	if err != nil {
+		log.Warn("can't execute query from db",
+			zap.String("query", query), zap.Error(err))
+		return 0
+	}
+	row.Next()
+	columns, _ := row.Columns()
+	addr := make([]interface{}, len(columns))
+	oneRow := make([]sql.NullString, len(columns))
+	var fieldIndex = -1
 	for i := range oneRow {
 		addr[i] = &oneRow[i]
+		if strings.ToLower(columns[i]) == fieldName {
+			fieldIndex = i
+		}
 	}
-	row := db.QueryRow(query)
-	err := row.Scan(addr...)
-	if err != nil {
+	err = row.Scan(addr...)
+	if err != nil || fieldIndex < 0 {
 		log.Warn("can't get estimate count from db",
 			zap.String("db", dbType),
 			zap.String("query", query), zap.Error(err))
 		return 0
 	}
+
 	estRows, err := strconv.ParseFloat(oneRow[fieldIndex].String, 64)
 	if err != nil {
 		log.Warn("can't get parse rows from db",
diff --git a/v4/export/sql_type.go b/v4/export/sql_type.go
index 2afa2df..2a459e4 100644
--- a/v4/export/sql_type.go
+++ b/v4/export/sql_type.go
@@ -9,7 +9,9 @@ import (
 var colTypeRowReceiverMap = map[string]func() RowReceiverStringer{}
 
 var nullValue = "NULL"
+var csvNullValue = "\\N"
 var quotationMark byte = '\''
+var doubleQuotationMark byte = '"'
 var quotationMarkNotQuote = []byte{quotationMark}
 var quotationMarkQuote = []byte{quotationMark, quotationMark}
 
@@ -146,6 +148,15 @@ func (r RowReceiverArr) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	bf.WriteByte(')')
 }
 
+func (r RowReceiverArr) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool) {
+	for i, receiver := range r {
+		receiver.WriteToBufferInCsv(bf, escapeBackslash)
+		if i != len(r)-1 {
+			bf.WriteByte(',')
+		}
+	}
+}
+
 type SQLTypeNumber struct {
 	SQLTypeString
 }
@@ -158,6 +169,14 @@ func (s SQLTypeNumber) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	}
 }
 
+func (s SQLTypeNumber) WriteToBufferInCsv(bf *bytes.Buffer, _ bool) {
+	if s.RawBytes != nil {
+		bf.Write(s.RawBytes)
+	} else {
+		bf.WriteString(csvNullValue)
+	}
+}
+
 type SQLTypeString struct {
 	sql.RawBytes
 }
@@ -182,6 +201,16 @@ func (s *SQLTypeString) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	}
 }
 
+func (s *SQLTypeString) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool) {
+	if s.RawBytes != nil {
+		bf.WriteByte(doubleQuotationMark)
+		escape(s.RawBytes, bf, escapeBackslash)
+		bf.WriteByte(doubleQuotationMark)
+	} else {
+		bf.WriteString(csvNullValue)
+	}
+}
+
 type SQLTypeBytes struct {
 	sql.RawBytes
 }
@@ -196,3 +225,13 @@ func (s *SQLTypeBytes) ReportSize() uint64 {
 func (s *SQLTypeBytes) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	fmt.Fprintf(bf, "x'%x'", s.RawBytes)
 }
+
+func (s *SQLTypeBytes) WriteToBufferInCsv(bf *bytes.Buffer, _ bool) {
+	if s.RawBytes != nil {
+		bf.WriteByte(doubleQuotationMark)
+		bf.Write(s.RawBytes)
+		bf.WriteByte(doubleQuotationMark)
+	} else {
+		bf.WriteString(csvNullValue)
+	}
+}
diff --git a/v4/export/writer.go b/v4/export/writer.go
index 3b5684c..c20afd0 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -90,3 +90,56 @@ func writeMetaToFile(target, metaSQL, path string) error {
 		metaSQL: metaSQL,
 	}, fileWriter)
 }
+
+type CsvWriter struct {
+	cfg *Config
+}
+
+func NewCsvWriter(config *Config) (Writer, error) {
+	sw := &CsvWriter{cfg: config}
+	return sw, os.MkdirAll(config.OutputDirPath, 0755)
+}
+
+func (f *CsvWriter) WriteDatabaseMeta(ctx context.Context, db, createSQL string) error {
+	fileName := fmt.Sprintf("%s-schema-create.sql", db)
+	filePath := path.Join(f.cfg.OutputDirPath, fileName)
+	return writeMetaToFile(db, createSQL, filePath)
+}
+
+func (f *CsvWriter) WriteTableMeta(ctx context.Context, db, table, createSQL string) error {
+	fileName := fmt.Sprintf("%s.%s-schema.sql", db, table)
+	filePath := path.Join(f.cfg.OutputDirPath, fileName)
+	return writeMetaToFile(db, createSQL, filePath)
+}
+
+func (f *CsvWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
+	log.Debug("start dumping table in csv format...", zap.String("table", ir.TableName()))
+
+	chunkIndex := ir.ChunkIndex()
+	fileName := fmt.Sprintf("%s.%s.%d.csv", ir.DatabaseName(), ir.TableName(), ir.ChunkIndex())
+	chunksIter := buildChunksIter(ir, f.cfg.FileSize, f.cfg.StatementSize)
+	defer chunksIter.Rows().Close()
+
+	for {
+		filePath := path.Join(f.cfg.OutputDirPath, fileName)
+		fileWriter, tearDown := buildInterceptFileWriter(filePath)
+		err := WriteInsertInCsv(chunksIter, fileWriter, f.cfg.NoHeader)
+		tearDown()
+		if err != nil {
+			return err
+		}
+
+		if w, ok := fileWriter.(*InterceptFileWriter); ok && !w.SomethingIsWritten {
+			break
+		}
+
+		if f.cfg.FileSize == UnspecifiedSize {
+			break
+		}
+		chunkIndex += 1
+		fileName = fmt.Sprintf("%s.%s.%d.csv", ir.DatabaseName(), ir.TableName(), chunkIndex)
+	}
+	log.Debug("dumping table in csv format successfully",
+		zap.String("table", ir.TableName()))
+	return nil
+}
diff --git a/v4/export/writer_util.go b/v4/export/writer_util.go
index 67fe874..8d1c3a7 100644
--- a/v4/export/writer_util.go
+++ b/v4/export/writer_util.go
@@ -182,6 +182,88 @@ func WriteInsert(tblIR TableDataIR, w io.Writer) error {
 	return wp.Error()
 }
 
+func WriteInsertInCsv(tblIR TableDataIR, w io.Writer, noHeader bool) error {
+	fileRowIter := tblIR.Rows()
+	if !fileRowIter.HasNext() {
+		return nil
+	}
+
+	bf := pool.Get().(*bytes.Buffer)
+	if bfCap := bf.Cap(); bfCap < lengthLimit {
+		bf.Grow(lengthLimit - bfCap)
+	}
+
+	wp := newWriterPipe(w)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		wp.Run(ctx)
+		wg.Done()
+	}()
+	defer func() {
+		cancel()
+		wg.Wait()
+	}()
+
+	var (
+		row             = MakeRowReceiver(tblIR.ColumnTypes())
+		counter         = 0
+		escapeBackSlash = tblIR.EscapeBackSlash()
+		err             error
+	)
+
+	if !noHeader && len(tblIR.ColumnNames()) != 0 {
+		for i, col := range tblIR.ColumnNames() {
+			bf.WriteByte(doubleQuotationMark)
+			escape([]byte(col), bf, escapeBackSlash)
+			bf.WriteByte(doubleQuotationMark)
+			if i != len(tblIR.ColumnTypes())-1 {
+				bf.WriteByte(',')
+			}
+		}
+		bf.WriteByte('\n')
+	}
+
+	for fileRowIter.HasNextSQLRowIter() {
+		fileRowIter = fileRowIter.NextSQLRowIter()
+		for fileRowIter.HasNext() {
+			if err = fileRowIter.Decode(row); err != nil {
+				log.Error("scanning from sql.Row failed", zap.Error(err))
+				return err
+			}
+
+			row.WriteToBufferInCsv(bf, escapeBackSlash)
+			counter += 1
+
+			if bf.Len() >= lengthLimit {
+				wp.input <- bf
+				bf = pool.Get().(*bytes.Buffer)
+				if bfCap := bf.Cap(); bfCap < lengthLimit {
+					bf.Grow(lengthLimit - bfCap)
+				}
+			}
+
+			fileRowIter.Next()
+			bf.WriteByte('\n')
+			if err = wp.Error(); err != nil {
+				return err
+			}
+		}
+	}
+
+	log.Debug("dumping table",
+		zap.String("table", tblIR.TableName()),
+		zap.Int("record counts", counter))
+	if bf.Len() > 0 {
+		wp.input <- bf
+	}
+	close(wp.input)
+	<-wp.closed
+	return wp.Error()
+}
+
 func write(writer io.StringWriter, str string) error {
 	_, err := writer.WriteString(str)
 	if err != nil {


Test patch:
diff --git a/v4/export/test_util.go b/v4/export/test_util.go
index 5212cd2..ee95e1f 100644
--- a/v4/export/test_util.go
+++ b/v4/export/test_util.go
@@ -82,6 +82,7 @@ type mockTableIR struct {
 	selectedField   string
 	specCmt         []string
 	colTypes        []string
+	colNames        []string
 	escapeBackSlash bool
 }
 
@@ -105,6 +106,10 @@ func (m *mockTableIR) ColumnTypes() []string {
 	return m.colTypes
 }
 
+func (m *mockTableIR) ColumnNames() []string {
+	return m.colNames
+}
+
 func (m *mockTableIR) SelectedField() string {
 	if m.selectedField == "*" {
 		return ""
diff --git a/v4/export/writer_util_test.go b/v4/export/writer_util_test.go
index 2c0b04e..e85fb24 100644
--- a/v4/export/writer_util_test.go
+++ b/v4/export/writer_util_test.go
@@ -71,6 +71,26 @@ func (s *testUtilSuite) TestWriteInsert(c *C) {
 	c.Assert(bf.String(), Equals, expected)
 }
 
+func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
+	data := [][]driver.Value{
+		{"1", "male", "bob@mail.com", "020-1234", nil},
+		{"2", "female", "sarah@mail.com", "020-1253", "healthy"},
+		{"3", "male", "john@mail.com", "020-1256", "healthy"},
+		{"4", "female", "sarah@mail.com", "020-1235", "healthy"},
+	}
+	colTypes := []string{"INT", "SET", "VARCHAR", "VARCHAR", "TEXT"}
+	tableIR := newMockTableIR("test", "employee", data, nil, colTypes)
+	bf := &bytes.Buffer{}
+
+	err := WriteInsertInCsv(tableIR, bf, true)
+	c.Assert(err, IsNil)
+	expected := "1,\"male\",\"bob@mail.com\",\"020-1234\",\\N\n" +
+		"2,\"female\",\"sarah@mail.com\",\"020-1253\",\"healthy\"\n" +
+		"3,\"male\",\"john@mail.com\",\"020-1256\",\"healthy\"\n" +
+		"4,\"female\",\"sarah@mail.com\",\"020-1235\",\"healthy\"\n"
+	c.Assert(bf.String(), Equals, expected)
+}
+
 func (s *testUtilSuite) TestSQLDataTypes(c *C) {
 	data := [][]driver.Value{
 		{"CHAR", "char1", `'char1'`},