Repository: pingcap/dumpling
Pull request #26 (2020-01-10 13:44:15)

*: add integration test for dumping in given file size
Changes:

- add integration tests: SQL files are split into size close to the given configuration.
- fix the one-offset bug in `tableDataChunks.Rows()`: calling it multiple times should not push forward its inner `*sql.Rows`.
- add test for `sizedRowIter`.
- add test for `prepareDumpingDatabases`.
- rewrite `SimpleWriter`.

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


```diff
@@            Coverage Diff             @@
##           master      #26      +/-   ##
==========================================
+ Coverage   55.78%   58.57%   +2.79%     
==========================================
  Files          13       13              
  Lines         683      729      +46     
==========================================
+ Hits          381      427      +46     
+ Misses        264      262       -2     
- Partials       38       40       +2
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index bd703ee..1173ccc 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -100,7 +100,7 @@ func main() {
 
 	err = export.Dump(conf)
 	if err != nil {
-		log.Zap().Error("dump failed", zap.String("error", err.Error()))
+		log.Zap().Error("dump failed", zap.Error(err))
 		os.Exit(1)
 	}
 	return
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index 5712831..ad6068d 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -48,7 +48,7 @@ func (s *sizedRowIter) Next(row RowReceiver) error {
 }
 
 func (s *sizedRowIter) HasNext() bool {
-	if s.currentSize > s.sizeLimit {
+	if s.currentSize >= s.sizeLimit {
 		return false
 	}
 	return s.rowIter.HasNext()
@@ -117,12 +117,16 @@ func (td *tableData) SpecialComments() StringIter {
 
 type tableDataChunks struct {
 	TableDataIR
+	rows      SQLRowIter
 	sizeLimit uint64
 }
 
 func (t *tableDataChunks) Rows() SQLRowIter {
+	if t.rows == nil {
+		t.rows = t.TableDataIR.Rows()
+	}
 	return &sizedRowIter{
-		rowIter:   t.Rows(),
+		rowIter:   t.rows,
 		sizeLimit: t.sizeLimit,
 	}
 }
diff --git a/v4/export/writer.go b/v4/export/writer.go
index a633df4..b33272c 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -3,10 +3,11 @@ package export
 import (
 	"context"
 	"fmt"
-	"io"
 	"os"
 	"path"
-	"sync"
+
+	"github.com/pingcap/dumpling/v4/log"
+	"go.uber.org/zap"
 )
 
 type Writer interface {
@@ -25,91 +26,58 @@ func NewSimpleWriter(config *Config) (Writer, error) {
 }
 
 func (f *SimpleWriter) WriteDatabaseMeta(ctx context.Context, db, createSQL string) error {
-	fileName := path.Join(f.cfg.OutputDirPath, fmt.Sprintf("%s-schema-create.sql", db))
-	fsStringWriter := NewFileSystemStringWriter(fileName, false)
-	meta := &metaData{
-		target:  db,
-		metaSQL: createSQL,
-	}
-	return WriteMeta(meta, fsStringWriter)
+	fileName := fmt.Sprintf("%s-schema-create.sql", db)
+	filePath := path.Join(f.cfg.OutputDirPath, fileName)
+	return writeMetaToFile(db, createSQL, filePath)
 }
 
 func (f *SimpleWriter) WriteTableMeta(ctx context.Context, db, table, createSQL string) error {
-	fileName := path.Join(f.cfg.OutputDirPath, fmt.Sprintf("%s.%s-schema.sql", db, table))
-	fsStringWriter := NewFileSystemStringWriter(fileName, false)
-	meta := &metaData{
-		target:  table,
-		metaSQL: createSQL,
-	}
-	return WriteMeta(meta, fsStringWriter)
+	fileName := fmt.Sprintf("%s.%s-schema.sql", db, table)
+	filePath := path.Join(f.cfg.OutputDirPath, fileName)
+	return writeMetaToFile(db, createSQL, filePath)
 }
 
 func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
+	log.Zap().Debug("start dumping table...", zap.String("table", ir.TableName()))
 	if f.cfg.FileSize == UnspecifiedSize {
 		fileName := path.Join(f.cfg.OutputDirPath, fmt.Sprintf("%s.%s.sql", ir.DatabaseName(), ir.TableName()))
-		fsStringWriter := NewFileSystemStringWriter(fileName, true)
-		return WriteInsert(ir, fsStringWriter)
+		fileWriter, tearDown := buildLazyFileWriter(fileName)
+		defer tearDown()
+		return WriteInsert(ir, fileWriter)
 	}
 
-	ir = splitTableDataIntoChunks(ir, f.cfg.FileSize)
-	for chunkCount := 0; ; /* loop */ chunkCount += 1 {
-		fileName := path.Join(f.cfg.OutputDirPath, fmt.Sprintf("%s.%s.%3d.sql", ir.DatabaseName(), ir.TableName(), chunkCount))
-		fsStringWriter := newInterceptStringWriter(NewFileSystemStringWriter(fileName, true))
-		err := WriteInsert(ir, fsStringWriter)
+	chunks := splitTableDataIntoChunks(ir, f.cfg.FileSize)
+	chunkCount := 0
+	for {
+		fileName := fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), chunkCount)
+		filePath := path.Join(f.cfg.OutputDirPath, fileName)
+		fileWriter, tearDown := buildLazyFileWriter(filePath)
+		intWriter := &InterceptStringWriter{StringWriter: fileWriter}
+		err := WriteInsert(chunks, intWriter)
+		tearDown()
 		if err != nil {
 			return err
 		}
-		if fsStringWriter.writeNothingYet {
+
+		if !intWriter.SomethingIsWritten {
 			break
 		}
+		chunkCount += 1
 	}
+	log.Zap().Debug("dumping table successfully",
+		zap.String("table", ir.TableName()))
 	return nil
 }
 
-type FileSystemStringWriter struct {
-	path string
-
-	file *os.File
-	once sync.Once
-	err  error
-}
-
-func (w *FileSystemStringWriter) initFileHandle() {
-	w.file, w.err = os.OpenFile(w.path, os.O_CREATE|os.O_WRONLY, 0755)
-}
-
-func (w *FileSystemStringWriter) WriteString(str string) (int, error) {
-	if w.err != nil {
-		return 0, w.err
+func writeMetaToFile(target, metaSQL, path string) error {
+	fileWriter, tearDown, err := buildFileWriter(path)
+	if err != nil {
+		return err
 	}
-	w.once.Do(w.initFileHandle)
-	return w.file.WriteString(str)
-}
+	defer tearDown()
 
-func NewFileSystemStringWriter(path string, lazyHandleCreation bool) *FileSystemStringWriter {
-	w := &FileSystemStringWriter{path: path}
-	if !lazyHandleCreation {
-		w.once.Do(w.initFileHandle)
-	}
-	return w
-}
-
-type interceptStringWriter struct {
-	sw              io.StringWriter
-	writeNothingYet bool
-}
-
-func (i *interceptStringWriter) WriteString(str string) (int, error) {
-	writtenBytes, err := i.sw.WriteString(str)
-	if writtenBytes != 0 {
-		i.writeNothingYet = false
-	}
-	return writtenBytes, err
-}
-
-func newInterceptStringWriter(sw io.StringWriter) *interceptStringWriter {
-	return &interceptStringWriter{
-		sw:              sw,
-		writeNothingYet: true,
-	}
+	return WriteMeta(&metaData{
+		target:  target,
+		metaSQL: metaSQL,
+	}, fileWriter)
 }
diff --git a/v4/export/writer_util.go b/v4/export/writer_util.go
index b838be8..6cf1227 100644
--- a/v4/export/writer_util.go
+++ b/v4/export/writer_util.go
@@ -1,9 +1,12 @@
 package export
 
 import (
+	"bufio"
 	"fmt"
 	"io"
+	"os"
 	"strings"
+	"sync"
 
 	"github.com/pingcap/dumpling/v4/log"
 	"go.uber.org/zap"
@@ -33,7 +36,6 @@ func WriteInsert(tblIR TableDataIR, w io.StringWriter) error {
 		return nil
 	}
 
-	log.Zap().Debug("start dumping for table", zap.String("table", tblIR.TableName()))
 	specCmtIter := tblIR.SpecialComments()
 	for specCmtIter.HasNext() {
 		if err := write(w, fmt.Sprintf("%s\n", specCmtIter.Next())); err != nil {
@@ -46,16 +48,18 @@ func WriteInsert(tblIR TableDataIR, w io.StringWriter) error {
 		return err
 	}
 
+	counter := 0
 	for rowIter.HasNext() {
 		row := MakeRowReceiver(tblIR.ColumnTypes())
 		if err := rowIter.Next(row); err != nil {
-			log.Zap().Error("scanning from sql.Row failed", zap.String("error", err.Error()))
+			log.Zap().Error("scanning from sql.Row failed", zap.Error(err))
 			return err
 		}
 
 		if err := write(w, row.ToString()); err != nil {
 			return err
 		}
+		counter += 1
 
 		var splitter string
 		if rowIter.HasNext() {
@@ -67,7 +71,9 @@ func WriteInsert(tblIR TableDataIR, w io.StringWriter) error {
 			return err
 		}
 	}
-	log.Zap().Debug("finish dumping for table", zap.String("table", tblIR.TableName()))
+	log.Zap().Debug("dumping table",
+		zap.String("table", tblIR.TableName()),
+		zap.Int("record counts", counter))
 	return nil
 }
 
@@ -76,11 +82,97 @@ func write(writer io.StringWriter, str string) error {
 	if err != nil {
 		log.Zap().Error("writing failed",
 			zap.String("string", str),
-			zap.String("error", err.Error()))
+			zap.Error(err))
 	}
 	return err
 }
 
+func buildFileWriter(path string) (io.StringWriter, func(), error) {
+	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
+	if err != nil {
+		log.Zap().Error("open file failed",
+			zap.String("path", path),
+			zap.Error(err))
+		return nil, nil, err
+	}
+	log.Zap().Debug("opened file", zap.String("path", path))
+	buf := bufio.NewWriter(file)
+	tearDownRoutine := func() {
+		_ = buf.Flush()
+		err := file.Close()
+		if err == nil {
+			return
+		}
+		log.Zap().Error("close file failed",
+			zap.String("path", path),
+			zap.Error(err))
+	}
+	return buf, tearDownRoutine, nil
+}
+
+func buildLazyFileWriter(path string) (io.StringWriter, func()) {
+	var file *os.File
+	var buf *bufio.Writer
+	lazyStringWriter := &LazyStringWriter{}
+	initRoutine := func() error {
+		f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
+		file = f
+		if err != nil {
+			log.Zap().Error("open file failed",
+				zap.String("path", path),
+				zap.Error(err))
+		}
+		log.Zap().Debug("opened file", zap.String("path", path))
+		buf = bufio.NewWriter(file)
+		lazyStringWriter.StringWriter = buf
+		return err
+	}
+	lazyStringWriter.initRoutine = initRoutine
+
+	tearDownRoutine := func() {
+		if file == nil {
+			return
+		}
+		log.Zap().Debug("tear down lazy file writer...")
+		_ = buf.Flush()
+		err := file.Close()
+		if err == nil {
+			return
+		}
+		log.Zap().Error("close file failed", zap.String("path", path))
+	}
+	return lazyStringWriter, tearDownRoutine
+}
+
+type LazyStringWriter struct {
+	initRoutine func() error
+	sync.Once
+	io.StringWriter
+	err error
+}
+
+func (l *LazyStringWriter) WriteString(str string) (int, error) {
+	l.Do(func() { l.err = l.initRoutine() })
+	if l.err != nil {
+		return 0, fmt.Errorf("open file error: %s", l.err.Error())
+	}
+	return l.StringWriter.WriteString(str)
+}
+
+// InterceptStringWriter is an interceptor of io.StringWriter,
+// tracking whether a StringWriter has written something.
+type InterceptStringWriter struct {
+	io.StringWriter
+	SomethingIsWritten bool
+}
+
+func (w *InterceptStringWriter) WriteString(str string) (int, error) {
+	if len(str) > 0 {
+		w.SomethingIsWritten = true
+	}
+	return w.StringWriter.WriteString(str)
+}
+
 func wrapBackTicks(identifier string) string {
 	if !strings.HasPrefix(identifier, "`") && !strings.HasSuffix(identifier, "`") {
 		return wrapStringWith(identifier, "`")


Test patch:
diff --git a/tests/file_size/run.sh b/tests/file_size/run.sh
new file mode 100644
index 0000000..9efe39d
--- /dev/null
+++ b/tests/file_size/run.sh
@@ -0,0 +1,32 @@
+#!/bin/sh
+
+set -eu
+
+run_sql "drop database if exists file_size"
+run_sql "create database file_size"
+export DUMPLING_TEST_DATABASE=file_size
+run_sql "create table t (a varchar(255))"
+
+chars_20="1111_0000_1111_0000_"
+
+# insert 100 records, each occupies 20 bytes
+i=0; while [ $i -lt 100 ]; do
+  run_sql "insert into t values (\"$chars_20\")"
+  i=$(( i + 1 ))
+done
+
+# dumping with file size = 200 bytes
+run_dumpling -filesize 200
+
+# the dumping result is expected to be:
+# 10 files for insertion(each conatins 10 records / 200 bytes)
+file_num=$(find "$DUMPLING_OUTPUT_DIR" -maxdepth 1 -iname "file_size.t.*.sql" | wc -l)
+if [ "$file_num" -ne 10 ]; then
+  echo "obtain file number: $file_num, but expect: 10" && exit 1
+fi
+
+total_lines=$(find "$DUMPLING_OUTPUT_DIR" -maxdepth 1 -iname "file_size.t.*.sql" -print0 \
+ | xargs -0 cat | grep "$chars_20" -c)
+if [ ! "$total_lines" = 100 ]; then
+  echo "obtain record number: $total_lines, but expect: 100" && exit 1
+fi
diff --git a/v4/export/consistency_test.go b/v4/export/consistency_test.go
index c2bc3cf..47c62d2 100644
--- a/v4/export/consistency_test.go
+++ b/v4/export/consistency_test.go
@@ -68,7 +68,7 @@ func (s *testConsistencySuite) TestConsistencyController(c *C) {
 		"db1": {"t1", "t2", "t3"},
 		"db2": {"t4"},
 	}
-	for i := 0; i < 4; i += 1 {
+	for i := 0; i < 4; i++ {
 		mock.ExpectExec("LOCK TABLES").WillReturnResult(resultOk)
 	}
 	mock.ExpectExec("UNLOCK TABLES").WillReturnResult(resultOk)
diff --git a/v4/export/ir_impl_test.go b/v4/export/ir_impl_test.go
index 4a2901a..4c8f8d9 100644
--- a/v4/export/ir_impl_test.go
+++ b/v4/export/ir_impl_test.go
@@ -1,6 +1,8 @@
 package export
 
 import (
+	"strings"
+
 	"github.com/DATA-DOG/go-sqlmock"
 	. "github.com/pingcap/check"
 )
@@ -10,15 +12,25 @@ var _ = Suite(&testIRImplSuite{})
 type testIRImplSuite struct{}
 
 type simpleRowReceiver struct {
-	data string
+	data []string
+}
+
+func newSimpleRowReceiver(length int) *simpleRowReceiver {
+	return &simpleRowReceiver{data: make([]string, length)}
 }
 
-func (s *simpleRowReceiver) BindAddress(arg []interface{}) {
-	arg[0] = &s.data
+func (s *simpleRowReceiver) BindAddress(args []interface{}) {
+	for i := range args {
+		args[i] = &s.data[i]
+	}
 }
 
 func (s *simpleRowReceiver) ReportSize() uint64 {
-	panic("not implement")
+	var sum uint64
+	for _, datum := range s.data {
+		sum += uint64(len(datum))
+	}
+	return sum
 }
 
 func (s *testIRImplSuite) TestRowIter(c *C) {
@@ -35,18 +47,49 @@ func (s *testIRImplSuite) TestRowIter(c *C) {
 	c.Assert(err, IsNil)
 
 	iter := newRowIter(rows, 1)
-	for i := 0; i < 100; i += 1 {
+	for i := 0; i < 100; i++ {
 		c.Assert(iter.HasNext(), IsTrue)
 	}
-	res := &simpleRowReceiver{}
+	res := newSimpleRowReceiver(1)
 	c.Assert(iter.Next(res), IsNil)
-	c.Assert(res.data, Equals, "1")
+	c.Assert(res.data, DeepEquals, []string{"1"})
 	c.Assert(iter.HasNext(), IsTrue)
 	c.Assert(iter.HasNext(), IsTrue)
 	c.Assert(iter.Next(res), IsNil)
-	c.Assert(res.data, Equals, "2")
+	c.Assert(res.data, DeepEquals, []string{"2"})
 	c.Assert(iter.HasNext(), IsTrue)
 	c.Assert(iter.Next(res), IsNil)
-	c.Assert(res.data, Equals, "3")
+	c.Assert(res.data, DeepEquals, []string{"3"})
 	c.Assert(iter.HasNext(), IsFalse)
 }
+
+func (s *testIRImplSuite) TestSizedRowIter(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	twentyBytes := strings.Repeat("x", 20)
+	thirtyBytes := strings.Repeat("x", 30)
+	expectedRows := mock.NewRows([]string{"a", "b"})
+	for i := 0; i < 10; i++ {
+		expectedRows.AddRow(twentyBytes, thirtyBytes)
+	}
+	mock.ExpectQuery("SELECT a, b FROM t").WillReturnRows(expectedRows)
+	rows, err := db.Query("SELECT a, b FROM t")
+	c.Assert(err, IsNil)
+
+	sizedRowIter := sizedRowIter{
+		rowIter:   newRowIter(rows, 2),
+		sizeLimit: 200,
+	}
+	res := newSimpleRowReceiver(2)
+	for i := 0; i < 200/50; i++ {
+		c.Assert(sizedRowIter.HasNext(), IsTrue)
+		err := sizedRowIter.Next(res)
+		c.Assert(err, IsNil)
+	}
+	c.Assert(sizedRowIter.HasNext(), IsFalse)
+	c.Assert(sizedRowIter.HasNext(), IsFalse)
+	rows.Close()
+	c.Assert(sizedRowIter.Next(res), NotNil)
+}
diff --git a/v4/export/prepare_test.go b/v4/export/prepare_test.go
new file mode 100644
index 0000000..a8b510c
--- /dev/null
+++ b/v4/export/prepare_test.go
@@ -0,0 +1,71 @@
+package export
+
+import (
+	"fmt"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	. "github.com/pingcap/check"
+)
+
+var _ = Suite(&testPrepareSuite{})
+
+type testPrepareSuite struct{}
+
+func (s *testPrepareSuite) TestPrepareDumpingDatabases(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	conf := DefaultConfig()
+	conf.Database = "db1,db2,db3"
+	result, err := prepareDumpingDatabases(conf, db)
+	c.Assert(err, IsNil)
+	c.Assert(result, DeepEquals, []string{"db1", "db2", "db3"})
+
+	conf.Database = ""
+	rows := sqlmock.NewRows([]string{"Database"}).
+		AddRow("db1").
+		AddRow("db2")
+	mock.ExpectQuery("SHOW DATABASES").WillReturnRows(rows)
+	result, err = prepareDumpingDatabases(conf, db)
+	c.Assert(err, IsNil)
+	c.Assert(result, DeepEquals, []string{"db1", "db2"})
+
+	mock.ExpectQuery("SHOW DATABASES").WillReturnError(fmt.Errorf("err"))
+	_, err = prepareDumpingDatabases(conf, db)
+	c.Assert(err, NotNil)
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+}
+
+func (s *testPrepareSuite) TestListAllTables(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	resultOk := sqlmock.NewResult(0, 0)
+	data := map[databaseName][]tableName{
+		"db1": {"t1", "t2"},
+		"db2": {"t3", "t4", "t5"},
+	}
+
+	var dbNames []databaseName
+	for dbName, tbNames := range data {
+		dbNames = append(dbNames, dbName)
+		mock.ExpectExec("USE .").WillReturnResult(resultOk)
+		rows := sqlmock.NewRows([]string{"Tables_in_xx"})
+		for _, name := range tbNames {
+			rows.AddRow(name)
+		}
+		mock.ExpectQuery("SHOW TABLES").WillReturnRows(rows)
+	}
+
+	tables, err := listAllTables(db, dbNames)
+	c.Assert(err, IsNil)
+
+	for d, t := range tables {
+		expectedTbs, ok := data[d]
+		c.Assert(ok, IsTrue)
+		c.Assert(t, DeepEquals, expectedTbs)
+	}
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+}