export: add CsvNullValue Cli usage (pingcap#58)

* add CsvNullValue Cli usage

* Update cmd/dumpling/main.go

Co-Authored-By: kennytm <[email protected]>

Co-authored-by: kennytm <[email protected]>

Author: AndrewDi

cmd/dumpling/main.go

@@ -45,6 +45,7 @@ var (
 	noHeader      bool
 	noSchemas     bool
 	noData        bool
+	csvNullValue  string
 
 	escapeBackslash bool
 
@@ -86,6 +87,7 @@ func init() {
 	rootCmd.PersistentFlags().BoolVar(&noHeader, "no-header", false, "whether not to dump CSV table header")
 	rootCmd.PersistentFlags().BoolVarP(&noSchemas, "no-schemas", "m", false, "Do not dump table schemas with the data")
 	rootCmd.PersistentFlags().BoolVarP(&noData, "no-data", "d", false, "Do not dump table data")
+	rootCmd.PersistentFlags().StringVar(&csvNullValue, "csv-null-value", "\\N", "The null value used when export to csv")
 }
 
 func run() {
@@ -112,6 +114,7 @@ func run() {
 	conf.NoHeader = noHeader
 	conf.NoSchemas = noSchemas
 	conf.NoData = noData
+	conf.CsvNullValue = csvNullValue
 
 	err := export.Dump(conf)
 	if err != nil {

v4/export/config.go

@@ -34,6 +34,7 @@ type Config struct {
 	NoHeader      bool
 	NoSchemas     bool
 	NoData        bool
+	CsvNullValue  string
 
 	BlackWhiteList  BWListConf
 	Rows            uint64
@@ -67,6 +68,7 @@ func DefaultConfig() *Config {
 		NoHeader:      false,
 		NoSchemas:     false,
 		NoData:        false,
+		CsvNullValue:  "\\N",
 	}
 }
 

v4/export/ir.go

@@ -38,7 +38,7 @@ type RowReceiverStringer interface {
 
 type Stringer interface {
 	WriteToBuffer(*bytes.Buffer, bool)
-	WriteToBufferInCsv(*bytes.Buffer, bool)
+	WriteToBufferInCsv(*bytes.Buffer, bool, string)
 }
 
 type RowReceiver interface {

v4/export/sql_type.go

@@ -9,7 +9,6 @@ import (
 var colTypeRowReceiverMap = map[string]func() RowReceiverStringer{}
 
 var nullValue = "NULL"
-var csvNullValue = "\\N"
 var quotationMark byte = '\''
 var doubleQuotationMark byte = '"'
 var quotationMarkNotQuote = []byte{quotationMark}
@@ -148,9 +147,9 @@ func (r RowReceiverArr) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	bf.WriteByte(')')
 }
 
-func (r RowReceiverArr) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool) {
+func (r RowReceiverArr) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, csvNullValue string) {
 	for i, receiver := range r {
-		receiver.WriteToBufferInCsv(bf, escapeBackslash)
+		receiver.WriteToBufferInCsv(bf, escapeBackslash, csvNullValue)
 		if i != len(r)-1 {
 			bf.WriteByte(',')
 		}
@@ -169,7 +168,7 @@ func (s SQLTypeNumber) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	}
 }
 
-func (s SQLTypeNumber) WriteToBufferInCsv(bf *bytes.Buffer, _ bool) {
+func (s SQLTypeNumber) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, csvNullValue string) {
 	if s.RawBytes != nil {
 		bf.Write(s.RawBytes)
 	} else {
@@ -201,7 +200,7 @@ func (s *SQLTypeString) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	}
 }
 
-func (s *SQLTypeString) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool) {
+func (s *SQLTypeString) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, csvNullValue string) {
 	if s.RawBytes != nil {
 		bf.WriteByte(doubleQuotationMark)
 		escape(s.RawBytes, bf, escapeBackslash)
@@ -226,7 +225,7 @@ func (s *SQLTypeBytes) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	fmt.Fprintf(bf, "x'%x'", s.RawBytes)
 }
 
-func (s *SQLTypeBytes) WriteToBufferInCsv(bf *bytes.Buffer, _ bool) {
+func (s *SQLTypeBytes) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, csvNullValue string) {
 	if s.RawBytes != nil {
 		bf.WriteByte(doubleQuotationMark)
 		bf.Write(s.RawBytes)

v4/export/writer.go

@@ -123,7 +123,7 @@ func (f *CsvWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
 	for {
 		filePath := path.Join(f.cfg.OutputDirPath, fileName)
 		fileWriter, tearDown := buildInterceptFileWriter(filePath)
-		err := WriteInsertInCsv(chunksIter, fileWriter, f.cfg.NoHeader)
+		err := WriteInsertInCsv(chunksIter, fileWriter, f.cfg.NoHeader, f.cfg.CsvNullValue)
 		tearDown()
 		if err != nil {
 			return err

v4/export/writer_util.go

@@ -182,7 +182,7 @@ func WriteInsert(tblIR TableDataIR, w io.Writer) error {
 	return wp.Error()
 }
 
-func WriteInsertInCsv(tblIR TableDataIR, w io.Writer, noHeader bool) error {
+func WriteInsertInCsv(tblIR TableDataIR, w io.Writer, noHeader bool, csvNullValue string) error {
 	fileRowIter := tblIR.Rows()
 	if !fileRowIter.HasNext() {
 		return nil
@@ -234,7 +234,7 @@ func WriteInsertInCsv(tblIR TableDataIR, w io.Writer, noHeader bool) error {
 				return err
 			}
 
-			row.WriteToBufferInCsv(bf, escapeBackSlash)
+			row.WriteToBufferInCsv(bf, escapeBackSlash, csvNullValue)
 			counter += 1
 
 			if bf.Len() >= lengthLimit {

v4/export/writer_util_test.go

@@ -82,7 +82,7 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	tableIR := newMockTableIR("test", "employee", data, nil, colTypes)
 	bf := &bytes.Buffer{}
 
-	err := WriteInsertInCsv(tableIR, bf, true)
+	err := WriteInsertInCsv(tableIR, bf, true, "\\N")
 	c.Assert(err, IsNil)
 	expected := "1,\"male\",\"[email protected]\",\"020-1234\",\\N\n" +
 		"2,\"female\",\"[email protected]\",\"020-1253\",\"healthy\"\n" +