Repository: pingcap/tidb-lightning
Pull request #364 (2020-08-04 14:01:48)

parser: fix csv parse header with empty line
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Enable csv parser to skip empty line. See: https://asktug.com/t/topic/36359/24

### What is changed and how it works?
Csv parse will treat the first non-empty line as header instead of the first line. And will also skip empty lines.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Side effects

Related changes

# Release notes

Empty lines in CSV are ignored.

kennytm: The unit test entered an infinite loop.
glorv: > The unit test entered an infinite loop.

fixed and add a unit test

Patch:
diff --git a/lightning/mydump/csv_parser.go b/lightning/mydump/csv_parser.go
index 30288d0..1c52a9b 100644
--- a/lightning/mydump/csv_parser.go
+++ b/lightning/mydump/csv_parser.go
@@ -14,8 +14,10 @@
 package mydump
 
 import (
+	"bytes"
 	"io"
 	"strings"
+	"unicode"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb-lightning/lightning/config"
@@ -51,6 +53,9 @@ type CSVParser struct {
 	fieldIndexes []int
 
 	lastRecord []string
+
+	// if set to true, csv parser will treat the first non-empty line as header line
+	shouldParseHeader bool
 }
 
 func NewCSVParser(
@@ -58,6 +63,7 @@ func NewCSVParser(
 	reader ReadSeekCloser,
 	blockBufSize int64,
 	ioWorkers *worker.Pool,
+	shouldParseHeader bool,
 ) *CSVParser {
 	quote := byte(0)
 	if len(cfg.Delimiter) > 0 {
@@ -78,13 +84,14 @@ func NewCSVParser(
 	}
 
 	return &CSVParser{
-		blockParser:      makeBlockParser(reader, blockBufSize, ioWorkers),
-		cfg:              cfg,
-		comma:            cfg.Separator[0],
-		quote:            quote,
-		escFlavor:        escFlavor,
-		quoteIndexFunc:   makeBytesIndexFunc(quoteStopSet),
-		unquoteIndexFunc: makeBytesIndexFunc(unquoteStopSet),
+		blockParser:       makeBlockParser(reader, blockBufSize, ioWorkers),
+		cfg:               cfg,
+		comma:             cfg.Separator[0],
+		quote:             quote,
+		escFlavor:         escFlavor,
+		quoteIndexFunc:    makeBytesIndexFunc(quoteStopSet),
+		unquoteIndexFunc:  makeBytesIndexFunc(unquoteStopSet),
+		shouldParseHeader: shouldParseHeader,
 	}
 }
 
@@ -177,6 +184,7 @@ func (parser *CSVParser) readRecord(dst []string) ([]string, error) {
 	parser.fieldIndexes = parser.fieldIndexes[:0]
 
 	isEmptyLine := true
+	whitespaceLine := true
 outside:
 	for {
 		firstByte, err := parser.readByte()
@@ -191,17 +199,23 @@ outside:
 		switch firstByte {
 		case parser.comma:
 			parser.fieldIndexes = append(parser.fieldIndexes, len(parser.recordBuffer))
-
+			whitespaceLine = false
 		case parser.quote:
 			if err := parser.readQuotedField(); err != nil {
 				return nil, err
 			}
+			whitespaceLine = false
 
 		case '\r', '\n':
 			// new line = end of record (ignore empty lines)
 			if isEmptyLine {
 				continue
 			}
+			// skip lines only contain whitespaces
+			if err == nil && whitespaceLine && len(bytes.TrimFunc(parser.recordBuffer, unicode.IsSpace)) == 0 {
+				parser.recordBuffer = parser.recordBuffer[:0]
+				continue
+			}
 			parser.fieldIndexes = append(parser.fieldIndexes, len(parser.recordBuffer))
 			break outside
 
@@ -327,7 +341,7 @@ func (parser *CSVParser) ReadRow() error {
 	row.RowID++
 
 	// skip the header first
-	if parser.pos == 0 && parser.cfg.Header {
+	if parser.shouldParseHeader {
 		columns, err := parser.readRecord(nil)
 		if err != nil {
 			return errors.Trace(err)
@@ -337,6 +351,7 @@ func (parser *CSVParser) ReadRow() error {
 			colName, _ = parser.unescapeString(colName)
 			parser.columns = append(parser.columns, strings.ToLower(colName))
 		}
+		parser.shouldParseHeader = false
 	}
 
 	records, err := parser.readRecord(parser.lastRecord)
diff --git a/lightning/mydump/region.go b/lightning/mydump/region.go
index 2680946..cdbd4b4 100644
--- a/lightning/mydump/region.go
+++ b/lightning/mydump/region.go
@@ -234,7 +234,7 @@ func SplitLargeFile(
 			if err != nil {
 				return 0, nil, nil, err
 			}
-			parser := NewCSVParser(&cfg.Mydumper.CSV, reader, cfg.Mydumper.ReadBlockSize, ioWorker)
+			parser := NewCSVParser(&cfg.Mydumper.CSV, reader, cfg.Mydumper.ReadBlockSize, ioWorker, false)
 			parser.SetPos(endOffset, prevRowIdMax)
 			pos, err := parser.ReadUntilTokNewLine()
 			if err != nil {
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index 38bfd7e..cb08579 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -1258,7 +1258,8 @@ func newChunkRestore(
 	var parser mydump.Parser
 	switch path.Ext(strings.ToLower(chunk.Key.Path)) {
 	case ".csv":
-		parser = mydump.NewCSVParser(&cfg.Mydumper.CSV, reader, blockBufSize, ioWorkers)
+		hasHeader := cfg.Mydumper.CSV.Header && chunk.Chunk.Offset == 0
+		parser = mydump.NewCSVParser(&cfg.Mydumper.CSV, reader, blockBufSize, ioWorkers, hasHeader)
 	default:
 		parser = mydump.NewChunkParser(cfg.TiDB.SQLMode, reader, blockBufSize, ioWorkers)
 	}


Test patch:
diff --git a/lightning/mydump/csv_parser_test.go b/lightning/mydump/csv_parser_test.go
index 13ec9cf..69dff19 100644
--- a/lightning/mydump/csv_parser_test.go
+++ b/lightning/mydump/csv_parser_test.go
@@ -55,7 +55,7 @@ type testCase struct {
 
 func (s *testMydumpCSVParserSuite) runTestCases(c *C, cfg *config.CSVConfig, blockBufSize int64, cases []testCase) {
 	for _, tc := range cases {
-		parser := mydump.NewCSVParser(cfg, mydump.NewStringReader(tc.input), blockBufSize, s.ioWorkers)
+		parser := mydump.NewCSVParser(cfg, mydump.NewStringReader(tc.input), blockBufSize, s.ioWorkers, false)
 		for i, row := range tc.expected {
 			comment := Commentf("input = %q, row = %d", tc.input, i+1)
 			e := parser.ReadRow()
@@ -68,7 +68,7 @@ func (s *testMydumpCSVParserSuite) runTestCases(c *C, cfg *config.CSVConfig, blo
 
 func (s *testMydumpCSVParserSuite) runFailingTestCases(c *C, cfg *config.CSVConfig, blockBufSize int64, cases []string) {
 	for _, tc := range cases {
-		parser := mydump.NewCSVParser(cfg, mydump.NewStringReader(tc), blockBufSize, s.ioWorkers)
+		parser := mydump.NewCSVParser(cfg, mydump.NewStringReader(tc), blockBufSize, s.ioWorkers, false)
 		e := parser.ReadRow()
 		c.Assert(e, ErrorMatches, "syntax error.*", Commentf("input = %q / %s", tc, errors.ErrorStack(e)))
 	}
@@ -87,7 +87,7 @@ func (s *testMydumpCSVParserSuite) TestTPCH(c *C) {
 		TrimLastSep: true,
 	}
 
-	parser := mydump.NewCSVParser(&cfg, reader, config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, reader, config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -151,7 +151,7 @@ func (s *testMydumpCSVParserSuite) TestRFC4180(c *C) {
 
 	// example 1, trailing new lines
 
-	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader("aaa,bbb,ccc\nzzz,yyy,xxx\n"), config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader("aaa,bbb,ccc\nzzz,yyy,xxx\n"), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -179,7 +179,7 @@ func (s *testMydumpCSVParserSuite) TestRFC4180(c *C) {
 
 	// example 2, no trailing new lines
 
-	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader("aaa,bbb,ccc\nzzz,yyy,xxx"), config.ReadBlockSize, s.ioWorkers)
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader("aaa,bbb,ccc\nzzz,yyy,xxx"), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -207,7 +207,7 @@ func (s *testMydumpCSVParserSuite) TestRFC4180(c *C) {
 
 	// example 5, quoted fields
 
-	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"aaa","bbb","ccc"`+"\nzzz,yyy,xxx"), config.ReadBlockSize, s.ioWorkers)
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"aaa","bbb","ccc"`+"\nzzz,yyy,xxx"), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -237,7 +237,7 @@ func (s *testMydumpCSVParserSuite) TestRFC4180(c *C) {
 
 	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"aaa","b
 bb","ccc"
-zzz,yyy,xxx`), config.ReadBlockSize, s.ioWorkers)
+zzz,yyy,xxx`), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -265,7 +265,7 @@ zzz,yyy,xxx`), config.ReadBlockSize, s.ioWorkers)
 
 	// example 7, quote escaping
 
-	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"aaa","b""bb","ccc"`), config.ReadBlockSize, s.ioWorkers)
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"aaa","b""bb","ccc"`), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -292,7 +292,7 @@ func (s *testMydumpCSVParserSuite) TestMySQL(c *C) {
 
 	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(`"\"","\\","\?"
 "\
-",\N,\\N`), config.ReadBlockSize, s.ioWorkers)
+",\N,\\N`), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -357,7 +357,7 @@ func (s *testMydumpCSVParserSuite) TestTSV(c *C) {
 	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(`a	b	c	d	e	f
 0				foo	0000-00-00
 0				foo	0000-00-00
-0	abc	def	ghi	bar	1999-12-31`), config.ReadBlockSize, s.ioWorkers)
+0	abc	def	ghi	bar	1999-12-31`), config.ReadBlockSize, s.ioWorkers, true)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -405,23 +405,70 @@ func (s *testMydumpCSVParserSuite) TestTSV(c *C) {
 	c.Assert(errors.Cause(parser.ReadRow()), Equals, io.EOF)
 }
 
+func (s *testMydumpCSVParserSuite) TestCsvWithWhiteSpaceLine(c *C) {
+	cfg := config.CSVConfig{
+		Separator: ",",
+		Delimiter: `"`,
+	}
+	data := " \r\n\r\n0,,abc\r\n \r\n123,1999-12-31,test\r\n"
+	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(data), config.ReadBlockSize, s.ioWorkers, false)
+	c.Assert(parser.ReadRow(), IsNil)
+	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
+		RowID: 1,
+		Row: []types.Datum{
+			types.NewStringDatum("0"),
+			nullDatum,
+			types.NewStringDatum("abc"),
+		},
+	})
+
+	c.Assert(parser, posEq, 12, 1)
+	c.Assert(parser.ReadRow(), IsNil)
+	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
+		RowID: 2,
+		Row: []types.Datum{
+			types.NewStringDatum("123"),
+			types.NewStringDatum("1999-12-31"),
+			types.NewStringDatum("test"),
+		},
+	})
+	c.Assert(parser.Close(), IsNil)
+
+	cfg.Header = true
+	data = " \r\na,b,c\r\n0,,abc\r\n"
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(data), config.ReadBlockSize, s.ioWorkers, true)
+	c.Assert(parser.ReadRow(), IsNil)
+	c.Assert(parser.Columns(), DeepEquals, []string{"a", "b", "c"})
+	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
+		RowID: 1,
+		Row: []types.Datum{
+			types.NewStringDatum("0"),
+			nullDatum,
+			types.NewStringDatum("abc"),
+		},
+	})
+
+	c.Assert(parser, posEq, 17, 1)
+	c.Assert(parser.Close(), IsNil)
+}
+
 func (s *testMydumpCSVParserSuite) TestEmpty(c *C) {
 	cfg := config.CSVConfig{
 		Separator: ",",
 		Delimiter: `"`,
 	}
 
-	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(""), config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(""), config.ReadBlockSize, s.ioWorkers, false)
 	c.Assert(errors.Cause(parser.ReadRow()), Equals, io.EOF)
 
 	// Try again with headers.
 
 	cfg.Header = true
 
-	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(""), config.ReadBlockSize, s.ioWorkers)
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader(""), config.ReadBlockSize, s.ioWorkers, true)
 	c.Assert(errors.Cause(parser.ReadRow()), Equals, io.EOF)
 
-	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader("h\n"), config.ReadBlockSize, s.ioWorkers)
+	parser = mydump.NewCSVParser(&cfg, mydump.NewStringReader("h\n"), config.ReadBlockSize, s.ioWorkers, true)
 	c.Assert(errors.Cause(parser.ReadRow()), Equals, io.EOF)
 }
 
@@ -430,7 +477,7 @@ func (s *testMydumpCSVParserSuite) TestCRLF(c *C) {
 		Separator: ",",
 		Delimiter: `"`,
 	}
-	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader("a\rb\r\nc\n\n\n\nd"), config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader("a\rb\r\nc\n\n\n\nd"), config.ReadBlockSize, s.ioWorkers, false)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -465,7 +512,7 @@ func (s *testMydumpCSVParserSuite) TestQuotedSeparator(c *C) {
 		Delimiter: `"`,
 	}
 
-	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(`",",','`), config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, mydump.NewStringReader(`",",','`), config.ReadBlockSize, s.ioWorkers, false)
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
 		RowID: 1,
@@ -636,7 +683,7 @@ func (s *testMydumpCSVParserSuite) TestReadError(c *C) {
 		Delimiter: `"`,
 	}
 
-	parser := mydump.NewCSVParser(&cfg, &errorReader{}, config.ReadBlockSize, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, &errorReader{}, config.ReadBlockSize, s.ioWorkers, false)
 	c.Assert(parser.ReadRow(), ErrorMatches, "fake read error")
 }
 
@@ -648,7 +695,7 @@ func (s *testMydumpCSVParserSuite) TestSyntaxErrorLog(c *C) {
 	}
 
 	tc := mydump.NewStringReader("x'" + strings.Repeat("y", 50000))
-	parser := mydump.NewCSVParser(&cfg, tc, 50000, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, tc, 50000, s.ioWorkers, false)
 	logger, buffer := log.MakeTestLogger()
 	parser.SetLogger(logger)
 	c.Assert(parser.ReadRow(), ErrorMatches, "syntax error.*")
@@ -697,7 +744,7 @@ func (s *benchCSVParserSuite) BenchmarkReadRowUsingMydumpCSVParser(c *C) {
 	}()
 
 	cfg := config.CSVConfig{Separator: ","}
-	parser := mydump.NewCSVParser(&cfg, file, 65536, s.ioWorkers)
+	parser := mydump.NewCSVParser(&cfg, file, 65536, s.ioWorkers, false)
 	parser.SetLogger(log.Logger{Logger: zap.NewNop()})
 
 	rowsCount := 0