Repository: pingcap/dumpling
Pull request #39 (2020-03-06 04:13:58)

support `--statement-size`  in dumpling
solve 1,000,000 bytes

support `statement-size` parameters  in `dumpling`, it can be used to control the size of an insert statement

the default value  of `statement-size` in `dumpling` is unspecified, which means it wouldn't take effect
the default value of `statement-size` in `mydumper` is `1,000,000 bytes`




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


```diff
@@            Coverage Diff             @@
##           master      #39      +/-   ##
==========================================
+ Coverage   76.58%   77.28%   +0.69%     
==========================================
  Files          14       14              
  Lines         850      876      +26     
==========================================
+ Hits          651      677      +26     
  Misses        149      149              
  Partials       50       50
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index 4bb50c0..55e8b97 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -27,18 +27,19 @@ import (
 )
 
 var (
-	database    string
-	host        string
-	user        string
-	port        int
-	password    string
-	threads     int
-	outputDir   string
-	fileSize    uint64
-	logLevel    string
-	consistency string
-	snapshot    string
-	noViews     bool
+	database      string
+	host          string
+	user          string
+	port          int
+	password      string
+	threads       int
+	outputDir     string
+	fileSize      uint64
+	statementSize uint64
+	logLevel      string
+	consistency   string
+	snapshot      string
+	noViews       bool
 
 	rootCmd = &cobra.Command{
 		Use:   "dumpling",
@@ -64,6 +65,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&password, "password", "p", "", "User password")
 	rootCmd.PersistentFlags().IntVarP(&threads, "threads", "t", 4, "Number of goroutines to use, default 4")
 	rootCmd.PersistentFlags().Uint64VarP(&fileSize, "filesize", "F", export.UnspecifiedSize, "The approximate size of output file")
+	rootCmd.PersistentFlags().Uint64VarP(&statementSize, "statement-size", "S", export.UnspecifiedSize, "Attempted size of INSERT statement in bytes")
 	rootCmd.PersistentFlags().StringVarP(&outputDir, "output", "o", defaultOutputDir, "Output directory")
 	rootCmd.PersistentFlags().StringVar(&logLevel, "loglevel", "info", "Log level: {debug|info|warn|error|dpanic|panic|fatal}")
 	rootCmd.PersistentFlags().StringVar(&consistency, "consistency", "auto", "Consistency level during dumping: {auto|none|flush|lock|snapshot}")
@@ -88,6 +90,7 @@ func run() {
 	conf.Password = password
 	conf.Threads = threads
 	conf.FileSize = fileSize
+	conf.StatementSize = statementSize
 	conf.OutputDirPath = outputDir
 	conf.NoViews = noViews
 
diff --git a/v4/export/config.go b/v4/export/config.go
index 134e5c4..210b880 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -19,6 +19,7 @@ type Config struct {
 	Threads  int
 
 	FileSize      uint64
+	StatementSize uint64
 	OutputDirPath string
 	ServerInfo    ServerInfo
 	SortByPk      bool
@@ -39,6 +40,7 @@ func DefaultConfig() *Config {
 		Password:      "",
 		Threads:       4,
 		FileSize:      UnspecifiedSize,
+		StatementSize: UnspecifiedSize,
 		OutputDirPath: ".",
 		ServerInfo:    ServerInfoUnknown,
 		SortByPk:      true,
diff --git a/v4/export/ir.go b/v4/export/ir.go
index b50a6e2..8bf9f74 100644
--- a/v4/export/ir.go
+++ b/v4/export/ir.go
@@ -17,6 +17,8 @@ type TableDataIR interface {
 type SQLRowIter interface {
 	Next(RowReceiver) error
 	HasNext() bool
+	HasNextSQLRowIter() bool
+	NextSQLRowIter() SQLRowIter
 }
 
 type RowReceiverStringer interface {
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index ad6068d..01e3482 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -32,26 +32,61 @@ func (iter *rowIter) HasNext() bool {
 	return iter.hasNext
 }
 
-type sizedRowIter struct {
-	rowIter     SQLRowIter
-	sizeLimit   uint64
-	currentSize uint64
+func (iter *rowIter) HasNextSQLRowIter() bool {
+	return iter.hasNext
+}
+
+func (iter *rowIter) NextSQLRowIter() SQLRowIter {
+	return iter
 }
 
-func (s *sizedRowIter) Next(row RowReceiver) error {
-	err := s.rowIter.Next(row)
+type fileRowIter struct {
+	rowIter            SQLRowIter
+	fileSizeLimit      uint64
+	statementSizeLimit uint64
+
+	currentStatementSize uint64
+	currentFileSize      uint64
+}
+
+func (c *fileRowIter) Next(row RowReceiver) error {
+	err := c.rowIter.Next(row)
 	if err != nil {
 		return err
 	}
-	s.currentSize += row.ReportSize()
+
+	size := row.ReportSize()
+	c.currentFileSize += size
+	c.currentStatementSize += size
 	return nil
 }
 
-func (s *sizedRowIter) HasNext() bool {
-	if s.currentSize >= s.sizeLimit {
+func (c *fileRowIter) HasNext() bool {
+	if c.fileSizeLimit != UnspecifiedSize && c.currentFileSize >= c.fileSizeLimit {
+		return false
+	}
+
+	if c.statementSizeLimit != UnspecifiedSize && c.currentStatementSize >= c.statementSizeLimit {
+		return false
+	}
+	return c.rowIter.HasNext()
+}
+
+func (c *fileRowIter) HasNextSQLRowIter() bool {
+	if c.fileSizeLimit != UnspecifiedSize && c.currentFileSize >= c.fileSizeLimit {
 		return false
 	}
-	return s.rowIter.HasNext()
+	return c.rowIter.HasNext()
+}
+
+func (c *fileRowIter) NextSQLRowIter() SQLRowIter {
+	return &fileRowIter{
+		rowIter:              c.rowIter,
+		fileSizeLimit:        c.fileSizeLimit,
+		statementSizeLimit:   c.statementSizeLimit,
+		currentFileSize:      c.currentFileSize,
+		currentStatementSize: 0,
+	}
 }
 
 type stringIter struct {
@@ -117,24 +152,28 @@ func (td *tableData) SpecialComments() StringIter {
 
 type tableDataChunks struct {
 	TableDataIR
-	rows      SQLRowIter
-	sizeLimit uint64
+	rows               SQLRowIter
+	chunkSizeLimit     uint64
+	statementSizeLimit uint64
 }
 
 func (t *tableDataChunks) Rows() SQLRowIter {
 	if t.rows == nil {
 		t.rows = t.TableDataIR.Rows()
 	}
-	return &sizedRowIter{
-		rowIter:   t.rows,
-		sizeLimit: t.sizeLimit,
+
+	return &fileRowIter{
+		rowIter:            t.rows,
+		statementSizeLimit: t.statementSizeLimit,
+		fileSizeLimit:      t.chunkSizeLimit,
 	}
 }
 
-func splitTableDataIntoChunks(td TableDataIR, chunkSize uint64) *tableDataChunks {
+func splitTableDataIntoChunks(td TableDataIR, chunkSize uint64, statementSize uint64) *tableDataChunks {
 	return &tableDataChunks{
-		TableDataIR: td,
-		sizeLimit:   chunkSize,
+		TableDataIR:        td,
+		chunkSizeLimit:     chunkSize,
+		statementSizeLimit: statementSize,
 	}
 }
 
diff --git a/v4/export/writer.go b/v4/export/writer.go
index b33272c..ff2018b 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -39,17 +39,15 @@ func (f *SimpleWriter) WriteTableMeta(ctx context.Context, db, table, createSQL
 
 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()))
-		fileWriter, tearDown := buildLazyFileWriter(fileName)
-		defer tearDown()
-		return WriteInsert(ir, fileWriter)
-	}
 
-	chunks := splitTableDataIntoChunks(ir, f.cfg.FileSize)
+	chunks := splitTableDataIntoChunks(ir, f.cfg.FileSize, f.cfg.StatementSize)
 	chunkCount := 0
+	fileName := fmt.Sprintf("%s.%s.sql", ir.DatabaseName(), ir.TableName())
+	if f.cfg.FileSize != UnspecifiedSize {
+		fileName = fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), chunkCount)
+	}
+
 	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}
@@ -62,7 +60,12 @@ func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error
 		if !intWriter.SomethingIsWritten {
 			break
 		}
+
+		if f.cfg.FileSize == UnspecifiedSize {
+			break
+		}
 		chunkCount += 1
+		fileName = fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), chunkCount)
 	}
 	log.Zap().Debug("dumping table successfully",
 		zap.String("table", ir.TableName()))
diff --git a/v4/export/writer_util.go b/v4/export/writer_util.go
index 6cf1227..2d64ca1 100644
--- a/v4/export/writer_util.go
+++ b/v4/export/writer_util.go
@@ -31,8 +31,8 @@ func WriteMeta(meta MetaIR, w io.StringWriter) error {
 }
 
 func WriteInsert(tblIR TableDataIR, w io.StringWriter) error {
-	rowIter := tblIR.Rows()
-	if !rowIter.HasNext() {
+	fileRowIter := tblIR.Rows()
+	if !fileRowIter.HasNext() {
 		return nil
 	}
 
@@ -43,34 +43,40 @@ func WriteInsert(tblIR TableDataIR, w io.StringWriter) error {
 		}
 	}
 
-	tblName := wrapBackTicks(tblIR.TableName())
-	if err := write(w, fmt.Sprintf("INSERT INTO %s VALUES\n", tblName)); err != nil {
-		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.Error(err))
+	var (
+		insertStatementPrefix = fmt.Sprintf("INSERT INTO %s VALUES\n", wrapBackTicks(tblIR.TableName()))
+		row                   = MakeRowReceiver(tblIR.ColumnTypes())
+		counter               = 0
+	)
+	for fileRowIter.HasNextSQLRowIter() {
+		if err := write(w, insertStatementPrefix); err != nil {
 			return err
 		}
 
-		if err := write(w, row.ToString()); err != nil {
-			return err
-		}
-		counter += 1
-
-		var splitter string
-		if rowIter.HasNext() {
-			splitter = ","
-		} else {
-			splitter = ";"
-		}
-		if err := write(w, fmt.Sprintf("%s\n", splitter)); err != nil {
-			return err
+		fileRowIter = fileRowIter.NextSQLRowIter()
+		for fileRowIter.HasNext() {
+			if err := fileRowIter.Next(row); err != nil {
+				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 fileRowIter.HasNext() {
+				splitter = ","
+			} else {
+				splitter = ";"
+			}
+			if err := write(w, fmt.Sprintf("%s\n", splitter)); err != nil {
+				return err
+			}
 		}
 	}
+
 	log.Zap().Debug("dumping table",
 		zap.String("table", tblIR.TableName()),
 		zap.Int("record counts", counter))


Test patch:
diff --git a/v4/export/ir_impl_test.go b/v4/export/ir_impl_test.go
index 4c8f8d9..9e002e3 100644
--- a/v4/export/ir_impl_test.go
+++ b/v4/export/ir_impl_test.go
@@ -63,7 +63,7 @@ func (s *testIRImplSuite) TestRowIter(c *C) {
 	c.Assert(iter.HasNext(), IsFalse)
 }
 
-func (s *testIRImplSuite) TestSizedRowIter(c *C) {
+func (s *testIRImplSuite) TestChunkRowIter(c *C) {
 	db, mock, err := sqlmock.New()
 	c.Assert(err, IsNil)
 	defer db.Close()
@@ -78,18 +78,41 @@ func (s *testIRImplSuite) TestSizedRowIter(c *C) {
 	rows, err := db.Query("SELECT a, b FROM t")
 	c.Assert(err, IsNil)
 
-	sizedRowIter := sizedRowIter{
-		rowIter:   newRowIter(rows, 2),
-		sizeLimit: 200,
-	}
+	var (
+		testFileSize      uint64 = 200
+		testStatementSize uint64 = 101
+
+		expectedSize = [][]uint64{
+			[]uint64{50, 50},
+			[]uint64{100, 100},
+			[]uint64{150, 150},
+			[]uint64{200, 50},
+		}
+	)
+
+	sqlRowIter := SQLRowIter(&fileRowIter{
+		rowIter:            newRowIter(rows, 2),
+		fileSizeLimit:      testFileSize,
+		statementSizeLimit: testStatementSize,
+	})
+
 	res := newSimpleRowReceiver(2)
-	for i := 0; i < 200/50; i++ {
-		c.Assert(sizedRowIter.HasNext(), IsTrue)
-		err := sizedRowIter.Next(res)
-		c.Assert(err, IsNil)
+
+	var resSize [][]uint64
+	for sqlRowIter.HasNextSQLRowIter() {
+		sqlRowIter = sqlRowIter.NextSQLRowIter()
+		fileRowIter, ok := sqlRowIter.(*fileRowIter)
+		c.Assert(ok, IsTrue)
+
+		for sqlRowIter.HasNext() {
+			c.Assert(sqlRowIter.Next(res), IsNil)
+			resSize = append(resSize, []uint64{fileRowIter.currentFileSize, fileRowIter.currentStatementSize})
+		}
 	}
-	c.Assert(sizedRowIter.HasNext(), IsFalse)
-	c.Assert(sizedRowIter.HasNext(), IsFalse)
+
+	c.Assert(resSize, DeepEquals, expectedSize)
+	c.Assert(sqlRowIter.HasNextSQLRowIter(), IsFalse)
+	c.Assert(sqlRowIter.HasNext(), IsFalse)
 	rows.Close()
-	c.Assert(sizedRowIter.Next(res), NotNil)
+	c.Assert(sqlRowIter.Next(res), NotNil)
 }
diff --git a/v4/export/writer_test.go b/v4/export/writer_test.go
index 5844474..52465bd 100644
--- a/v4/export/writer_test.go
+++ b/v4/export/writer_test.go
@@ -3,6 +3,7 @@ package export
 import (
 	"context"
 	"database/sql/driver"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path"
@@ -149,3 +150,87 @@ func (s *testDumpSuite) TestWriteTableDataWithFileSize(c *C) {
 		c.Assert(string(bytes), Equals, expected)
 	}
 }
+
+func (s *testDumpSuite) TestWriteTableDataWithStatementSize(c *C) {
+	dir, err := ioutil.TempDir("", "dumpling")
+	c.Assert(err, IsNil)
+
+	config := DefaultConfig()
+	config.OutputDirPath = dir
+	config.StatementSize = 50
+	ctx := context.Background()
+	defer os.RemoveAll(config.OutputDirPath)
+
+	writer, err := NewSimpleWriter(config)
+	c.Assert(err, IsNil)
+
+	data := [][]driver.Value{
+		{"1", "male", "bob@mail.com", "020-1234", nil},
+		{"2", "female", "sarah@mail.com", "020-1253", "healthy"},
+		{"3", "male", "john@mail.com", "020-1256", "healthy"},
+		{"4", "female", "sarah@mail.com", "020-1235", "healthy"},
+	}
+	colTypes := []string{"INT", "SET", "VARCHAR", "VARCHAR", "TEXT"}
+	specCmts := []string{
+		"/*!40101 SET NAMES binary*/;",
+		"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;",
+	}
+	tableIR := newMockTableIR("test", "employee", data, specCmts, colTypes)
+	err = writer.WriteTableData(ctx, tableIR)
+	c.Assert(err, IsNil)
+
+	// only with statement size
+	cases := map[string]string{
+		"test.employee.sql": "/*!40101 SET NAMES binary*/;\n" +
+			"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n" +
+			"INSERT INTO `employee` VALUES\n" +
+			"(1, 'male', 'bob@mail.com', '020-1234', NULL),\n" +
+			"(2, 'female', 'sarah@mail.com', '020-1253', 'healthy');\n" +
+			"INSERT INTO `employee` VALUES\n" +
+			"(3, 'male', 'john@mail.com', '020-1256', 'healthy'),\n" +
+			"(4, 'female', 'sarah@mail.com', '020-1235', 'healthy');\n",
+	}
+
+	for p, expected := range cases {
+		p := path.Join(config.OutputDirPath, p)
+		_, err = os.Stat(p)
+		c.Assert(err, IsNil)
+		bytes, err := ioutil.ReadFile(p)
+		c.Assert(err, IsNil)
+		c.Assert(string(bytes), Equals, expected)
+	}
+
+	// with file size and statement size
+	config.FileSize = 90
+	config.StatementSize = 30
+	os.RemoveAll(config.OutputDirPath)
+	config.OutputDirPath, err = ioutil.TempDir("", "dumpling")
+	fmt.Println(config.OutputDirPath)
+	c.Assert(err, IsNil)
+
+	cases = map[string]string{
+		"test.employee.0.sql": "/*!40101 SET NAMES binary*/;\n" +
+			"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n" +
+			"INSERT INTO `employee` VALUES\n" +
+			"(1, 'male', 'bob@mail.com', '020-1234', NULL),\n" +
+			"(2, 'female', 'sarah@mail.com', '020-1253', 'healthy');\n" +
+			"INSERT INTO `employee` VALUES\n" +
+			"(3, 'male', 'john@mail.com', '020-1256', 'healthy');\n",
+		"test.employee.1.sql": "/*!40101 SET NAMES binary*/;\n" +
+			"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n" +
+			"INSERT INTO `employee` VALUES\n" +
+			"(4, 'female', 'sarah@mail.com', '020-1235', 'healthy');\n",
+	}
+
+	tableIR = newMockTableIR("test", "employee", data, specCmts, colTypes)
+	c.Assert(writer.WriteTableData(ctx, tableIR), IsNil)
+	c.Assert(err, IsNil)
+	for p, expected := range cases {
+		p := path.Join(config.OutputDirPath, p)
+		_, err = os.Stat(p)
+		c.Assert(err, IsNil)
+		bytes, err := ioutil.ReadFile(p)
+		c.Assert(err, IsNil)
+		c.Assert(string(bytes), Equals, expected)
+	}
+}