*: add integration test for dumping in given file size (#26)

* *: add integration test for dumping in given file size

* log: use zap.Error instead of zap.String('err', xx)

* *: fix style problems

Author: tangenta

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

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

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)

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,
 	}
 }

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)
+}

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)
+}

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)
 }