address comments

pingcap · Jul 7, 2020 · 7eb9742 · 7eb9742
1 parent 273c82f
commit 7eb9742
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 27 deletions.
diff --git a/v4/export/ir.go b/v4/export/ir.go
@@ -39,7 +39,7 @@ type RowReceiverStringer interface {
 
 type Stringer interface {
 	WriteToBuffer(*bytes.Buffer, bool)
-	WriteToBufferInCsv(*bytes.Buffer, bool, csvOption)
+	WriteToBufferInCsv(*bytes.Buffer, bool, *csvOption)
 }
 
 type RowReceiver interface {

diff --git a/v4/export/sql_type.go b/v4/export/sql_type.go
@@ -153,7 +153,7 @@ func (r RowReceiverArr) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	bf.WriteByte(')')
 }
 
-func (r RowReceiverArr) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, opt csvOption) {
+func (r RowReceiverArr) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, opt *csvOption) {
 	for i, receiver := range r {
 		receiver.WriteToBufferInCsv(bf, escapeBackslash, opt)
 		if i != len(r)-1 {
@@ -174,7 +174,7 @@ func (s SQLTypeNumber) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	}
 }
 
-func (s SQLTypeNumber) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, opt csvOption) {
+func (s SQLTypeNumber) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, opt *csvOption) {
 	if s.RawBytes != nil {
 		bf.Write(s.RawBytes)
 	} else {
@@ -206,10 +206,10 @@ func (s *SQLTypeString) WriteToBuffer(bf *bytes.Buffer, escapeBackslash bool) {
 	}
 }
 
-func (s *SQLTypeString) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, opt csvOption) {
+func (s *SQLTypeString) WriteToBufferInCsv(bf *bytes.Buffer, escapeBackslash bool, opt *csvOption) {
 	if s.RawBytes != nil {
 		bf.Write(opt.delimiter)
-		escape(s.RawBytes, bf, getEscapeQuotation(escapeBackslash, opt.separator))
+		escape(s.RawBytes, bf, getEscapeQuotation(escapeBackslash, opt.delimiter))
 		bf.Write(opt.delimiter)
 	} else {
 		bf.WriteString(opt.nullValue)
@@ -231,7 +231,7 @@ func (s *SQLTypeBytes) WriteToBuffer(bf *bytes.Buffer, _ bool) {
 	fmt.Fprintf(bf, "x'%x'", s.RawBytes)
 }
 
-func (s *SQLTypeBytes) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, opt csvOption) {
+func (s *SQLTypeBytes) WriteToBufferInCsv(bf *bytes.Buffer, _ bool, opt *csvOption) {
 	if s.RawBytes != nil {
 		bf.Write(opt.delimiter)
 		bf.Write(s.RawBytes)

diff --git a/v4/export/test_util.go b/v4/export/test_util.go
@@ -148,7 +148,7 @@ func (m *mockTableIR) EscapeBackSlash() bool {
 	return m.escapeBackSlash
 }
 
-func newMockTableIR(databaseName, tableName string, data [][]driver.Value, specialComments, colTypes []string) TableDataIR {
+func newMockTableIR(databaseName, tableName string, data [][]driver.Value, specialComments, colTypes []string) *mockTableIR {
 	return &mockTableIR{
 		dbName:        databaseName,
 		tblName:       tableName,
@@ -158,15 +158,3 @@ func newMockTableIR(databaseName, tableName string, data [][]driver.Value, speci
 		colTypes:      colTypes,
 	}
 }
-
-func newMockTableIRWithError(databaseName, tableName string, data [][]driver.Value, specialComments, colTypes []string, err error) TableDataIR {
-	return &mockTableIR{
-		dbName:        databaseName,
-		tblName:       tableName,
-		data:          data,
-		specCmt:       specialComments,
-		selectedField: "*",
-		colTypes:      colTypes,
-		rowErr:        err,
-	}
-}
diff --git a/v4/export/writer.go b/v4/export/writer.go
@@ -148,7 +148,7 @@ func (f *CsvWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
 	chunksIter := buildChunksIter(ir, f.cfg.FileSize, f.cfg.StatementSize)
 	defer chunksIter.Rows().Close()
 
-	opt := csvOption{
+	opt := &csvOption{
 		nullValue: f.cfg.CsvNullValue,
 		separator: []byte(f.cfg.CsvSeparator),
 		delimiter: []byte(f.cfg.CsvDelimiter),

diff --git a/v4/export/writer_util.go b/v4/export/writer_util.go
@@ -189,7 +189,7 @@ func WriteInsert(pCtx context.Context, tblIR TableDataIR, w io.Writer) error {
 	return wp.Error()
 }
 
-func WriteInsertInCsv(pCtx context.Context, tblIR TableDataIR, w io.Writer, noHeader bool, opt csvOption) error {
+func WriteInsertInCsv(pCtx context.Context, tblIR TableDataIR, w io.Writer, noHeader bool, opt *csvOption) error {
 	fileRowIter := tblIR.Rows()
 	if !fileRowIter.HasNext() {
 		return nil
@@ -223,11 +223,11 @@ func WriteInsertInCsv(pCtx context.Context, tblIR TableDataIR, w io.Writer, noHe
 
 	if !noHeader && len(tblIR.ColumnNames()) != 0 {
 		for i, col := range tblIR.ColumnNames() {
-			bf.Write(opt.separator)
-			escape([]byte(col), bf, getEscapeQuotation(escapeBackSlash, opt.separator))
-			bf.Write(opt.separator)
+			bf.Write(opt.delimiter)
+			escape([]byte(col), bf, getEscapeQuotation(escapeBackSlash, opt.delimiter))
+			bf.Write(opt.delimiter)
 			if i != len(tblIR.ColumnTypes())-1 {
-				bf.WriteByte(',')
+				bf.Write(opt.separator)
 			}
 		}
 		bf.WriteByte('\n')

diff --git a/v4/export/writer_util_test.go b/v4/export/writer_util_test.go
@@ -87,7 +87,8 @@ func (s *testUtilSuite) TestWriteInsertReturnsError(c *C) {
 	}
 	// row errors at last line
 	rowErr := errors.New("mock row error")
-	tableIR := newMockTableIRWithError("test", "employee", data, specCmts, colTypes, rowErr)
+	tableIR := newMockTableIR("test", "employee", data, specCmts, colTypes)
+	tableIR.rowErr = rowErr
 	bf := &bytes.Buffer{}
 
 	err := WriteInsert(context.Background(), tableIR, bf)
@@ -112,7 +113,8 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	tableIR := newMockTableIR("test", "employee", data, nil, colTypes)
 	bf := &bytes.Buffer{}
 
-	opt := csvOption{separator: []byte(","), delimiter: doubleQuotationMark, nullValue: "\\N"}
+	// test nullValue
+	opt := &csvOption{separator: []byte(","), delimiter: doubleQuotationMark, nullValue: "\\N"}
 	err := WriteInsertInCsv(context.Background(), tableIR, bf, true, opt)
 	c.Assert(err, IsNil)
 	expected := "1,\"male\",\"[email protected]\",\"020-1234\",\\N\n" +
@@ -121,6 +123,7 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 		"4,\"female\",\"[email protected]\",\"020-1235\",\"healthy\"\n"
 	c.Assert(bf.String(), Equals, expected)
 
+	// test delimiter
 	bf.Reset()
 	opt.delimiter = quotationMark
 	err = WriteInsertInCsv(context.Background(), tableIR, bf, true, opt)
@@ -131,6 +134,7 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 		"4,'female','[email protected]','020-1235','healthy'\n"
 	c.Assert(bf.String(), Equals, expected)
 
+	// test separator
 	bf.Reset()
 	opt.separator = []byte(";")
 	err = WriteInsertInCsv(context.Background(), tableIR, bf, true, opt)
@@ -140,6 +144,20 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 		"3;'male';'[email protected]';'020-1256';'healthy'\n" +
 		"4;'female';'[email protected]';'020-1235';'healthy'\n"
 	c.Assert(bf.String(), Equals, expected)
+
+	// test delimiter that included in values
+	bf.Reset()
+	opt.separator = []byte(",")
+	opt.delimiter = []byte("ma")
+	tableIR.colNames = []string{"id", "gender", "email", "phone_number", "status"}
+	err = WriteInsertInCsv(context.Background(), tableIR, bf, false, opt)
+	c.Assert(err, IsNil)
+	expected = "maidma,magenderma,maemamailma,maphone_numberma,mastatusma\n" +
+		"1,mamamalema,[email protected],ma020-1234ma,\\N\n" +
+		"2,mafemamalema,[email protected],ma020-1253ma,mahealthyma\n" +
+		"3,mamamalema,[email protected],ma020-1256ma,mahealthyma\n" +
+		"4,mafemamalema,[email protected],ma020-1235ma,mahealthyma\n"
+	c.Assert(bf.String(), Equals, expected)
 }
 
 func (s *testUtilSuite) TestSQLDataTypes(c *C) {