Repository: pingcap/dumpling
Pull request #122 (2020-07-14 03:13:10)

support specify ouput filename format
<!--
Thank you for contributing to Dumpling! Please read the [CONTRIBUTING](https://github.com/pingcap/dumpling/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

reopened PR for https://github.com/pingcap/dumpling/pull/120

### What problem does this PR solve? <!--add issue link with summary if exists-->
implement https://github.com/pingcap/dumpling/issues/114

### What is changed and how it works?
Add `--output-filename-format` argument to support specifying the format of output filename.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test

Related changes

 - Need to cherry-pick to the release branch
 - Need to update the documentation
 
### Release note

<!-- bugfixes or new feature need a release note, must in the form of a list, such as

- support -T/--tables-list argument

or if no need to be included in the release note, just add the following line

- No release note
-->
- support specify output filename format

codecov[bot]: # [Codecov](https://codecov.io/gh/pingcap/dumpling/pull/122?src=pr&el=h1) Report
> Merging [#122](https://codecov.io/gh/pingcap/dumpling/pull/122?src=pr&el=desc) into [master](https://codecov.io/gh/pingcap/dumpling/commit/dc5af4ed250b4dfd2c30d086772abc5f42e6bfba&el=desc) will **decrease** coverage by `0.26%`.
> The diff coverage is `68.85%`.


```diff
@@            Coverage Diff             @@
##           master     #122      +/-   ##
==========================================
- Coverage   51.94%   51.67%   -0.27%     
==========================================
  Files          17       17              
  Lines        1829     1850      +21     
==========================================
+ Hits          950      956       +6     
- Misses        813      826      +13     
- Partials       66       68       +2     
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index 5fd8392..52816fe 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -21,6 +21,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/docker/go-units"
@@ -66,9 +67,10 @@ var (
 	csvSeparator  string
 	csvDelimiter  string
 
-	dumpEmptyDatabase bool
-	escapeBackslash   bool
-	tidbMemQuotaQuery uint64
+	dumpEmptyDatabase    bool
+	escapeBackslash      bool
+	tidbMemQuotaQuery    uint64
+	outputFilenameFormat string
 )
 
 var defaultOutputDir = timestampDirName()
@@ -119,6 +121,7 @@ func main() {
 	pflag.StringVar(&keyPath, "key", "", "The path name to the client private key file for TLS connection")
 	pflag.StringVar(&csvSeparator, "csv-separator", ",", "The separator for csv files, default ','")
 	pflag.StringVar(&csvDelimiter, "csv-delimiter", "\"", "The delimiter for values in csv files, default '\"'")
+	pflag.StringVar(&outputFilenameFormat, "output-filename-template", "", "The output filename template (without file extension), default '{{.DB}}.{{.Table}}.{{.Index}}'")
 
 	printVersion := pflag.BoolP("version", "V", false, "Print Dumpling version")
 
@@ -152,6 +155,19 @@ func main() {
 		os.Exit(2)
 	}
 
+	if outputFilenameFormat == "" {
+		if sql != "" {
+			outputFilenameFormat = "result.{{.Index}}"
+		} else {
+			outputFilenameFormat = "{{.DB}}.{{.Table}}.{{.Index}}"
+		}
+	}
+	tmpl, err := template.New("filename").Parse(outputFilenameFormat)
+	if err != nil {
+		fmt.Printf("failed to parse output filename template (--output-filename-template '%s')\n", outputFilenameFormat)
+		os.Exit(2)
+	}
+
 	if threads <= 0 {
 		fmt.Printf("--threads is set to %d. It should be greater than 0\n", threads)
 		os.Exit(2)
@@ -191,6 +207,7 @@ func main() {
 	conf.SessionParams["tidb_mem_quota_query"] = tidbMemQuotaQuery
 	conf.CsvSeparator = csvSeparator
 	conf.CsvDelimiter = csvDelimiter
+	conf.OutputFileTemplate = tmpl
 
 	err = export.Dump(context.Background(), conf)
 	if err != nil {
diff --git a/v4/export/config.go b/v4/export/config.go
index 85ffd93..d1bf824 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"regexp"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/coreos/go-semver/semver"
@@ -49,46 +50,49 @@ type Config struct {
 	CsvSeparator  string
 	CsvDelimiter  string
 
-	TableFilter       filter.Filter
-	Rows              uint64
-	Where             string
-	FileType          string
-	EscapeBackslash   bool
-	DumpEmptyDatabase bool
-	SessionParams     map[string]interface{}
+	TableFilter        filter.Filter
+	Rows               uint64
+	Where              string
+	FileType           string
+	EscapeBackslash    bool
+	DumpEmptyDatabase  bool
+	OutputFileTemplate *template.Template
+	SessionParams      map[string]interface{}
 }
 
 func DefaultConfig() *Config {
 	allFilter, _ := filter.Parse([]string{"*.*"})
+	tmpl := template.Must(template.New("filename").Parse("{{.DB}}.{{.Table}}.{{.Index}}"))
 	return &Config{
-		Databases:         nil,
-		Host:              "127.0.0.1",
-		User:              "root",
-		Port:              3306,
-		Password:          "",
-		Threads:           4,
-		Logger:            nil,
-		StatusAddr:        ":8281",
-		FileSize:          UnspecifiedSize,
-		StatementSize:     UnspecifiedSize,
-		OutputDirPath:     ".",
-		ServerInfo:        ServerInfoUnknown,
-		SortByPk:          true,
-		Tables:            nil,
-		Snapshot:          "",
-		Consistency:       "auto",
-		NoViews:           true,
-		Rows:              UnspecifiedSize,
-		Where:             "",
-		FileType:          "SQL",
-		NoHeader:          false,
-		NoSchemas:         false,
-		NoData:            false,
-		CsvNullValue:      "\\N",
-		Sql:               "",
-		TableFilter:       allFilter,
-		DumpEmptyDatabase: true,
-		SessionParams:     make(map[string]interface{}),
+		Databases:          nil,
+		Host:               "127.0.0.1",
+		User:               "root",
+		Port:               3306,
+		Password:           "",
+		Threads:            4,
+		Logger:             nil,
+		StatusAddr:         ":8281",
+		FileSize:           UnspecifiedSize,
+		StatementSize:      UnspecifiedSize,
+		OutputDirPath:      ".",
+		ServerInfo:         ServerInfoUnknown,
+		SortByPk:           true,
+		Tables:             nil,
+		Snapshot:           "",
+		Consistency:        "auto",
+		NoViews:            true,
+		Rows:               UnspecifiedSize,
+		Where:              "",
+		FileType:           "SQL",
+		NoHeader:           false,
+		NoSchemas:          false,
+		NoData:             false,
+		CsvNullValue:       "\\N",
+		Sql:                "",
+		TableFilter:        allFilter,
+		DumpEmptyDatabase:  true,
+		SessionParams:      make(map[string]interface{}),
+		OutputFileTemplate: tmpl,
 	}
 }
 
diff --git a/v4/export/prepare.go b/v4/export/prepare.go
index 81d1516..7bbc5ac 100644
--- a/v4/export/prepare.go
+++ b/v4/export/prepare.go
@@ -3,6 +3,7 @@ package export
 import (
 	"database/sql"
 	"strings"
+	"text/template"
 
 	"github.com/go-sql-driver/mysql"
 	"github.com/pingcap/dumpling/v4/log"
@@ -44,6 +45,13 @@ func adjustConfig(conf *Config) error {
 	if conf.SessionParams == nil {
 		conf.SessionParams = make(map[string]interface{})
 	}
+	if conf.OutputFileTemplate == nil {
+		var err error
+		conf.OutputFileTemplate, err = template.New("filename").Parse("{{.DB}}.{{.Table}}.{{.Index}}")
+		if err != nil {
+			return err
+		}
+	}
 	resolveAutoConsistency(conf)
 
 	return nil
diff --git a/v4/export/writer.go b/v4/export/writer.go
index d28b232..126aa4a 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -1,10 +1,12 @@
 package export
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"os"
 	"path"
+	"text/template"
 
 	"go.uber.org/zap"
 
@@ -49,9 +51,12 @@ func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error
 			fileName = fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), 0)
 		}
 	}*/
-
 	namer := newOutputFileNamer(ir)
-	fileName := fmt.Sprintf("%s.sql", namer.NextName())
+	fileName, err := namer.NextName(f.cfg.OutputFileTemplate)
+	if err != nil {
+		return err
+	}
+	fileName += ".sql"
 	chunksIter := buildChunksIter(ir, f.cfg.FileSize, f.cfg.StatementSize)
 	defer chunksIter.Rows().Close()
 
@@ -71,7 +76,11 @@ func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error
 		if f.cfg.FileSize == UnspecifiedSize {
 			break
 		}
-		fileName = fmt.Sprintf("%s.sql", namer.NextName())
+		fileName, err = namer.NextName(f.cfg.OutputFileTemplate)
+		if err != nil {
+			return err
+		}
+		fileName += ".sql"
 	}
 	log.Debug("dumping table successfully",
 		zap.String("table", ir.TableName()))
@@ -113,9 +122,9 @@ func (f *CsvWriter) WriteTableMeta(ctx context.Context, db, table, createSQL str
 }
 
 type outputFileNamer struct {
-	chunkIndex int
-	dbName     string
-	tableName  string
+	Index int
+	DB    string
+	Table string
 }
 
 type csvOption struct {
@@ -126,25 +135,31 @@ type csvOption struct {
 
 func newOutputFileNamer(ir TableDataIR) *outputFileNamer {
 	return &outputFileNamer{
-		chunkIndex: ir.ChunkIndex(),
-		dbName:     ir.DatabaseName(),
-		tableName:  ir.TableName(),
+		Index: ir.ChunkIndex(),
+		DB:    ir.DatabaseName(),
+		Table: ir.TableName(),
 	}
 }
 
-func (namer *outputFileNamer) NextName() string {
-	defer func() { namer.chunkIndex++ }()
-	if namer.dbName == "" || namer.tableName == "" {
-		return fmt.Sprintf("result.%d", namer.chunkIndex)
+func (namer *outputFileNamer) NextName(tmpl *template.Template) (string, error) {
+	defer func() { namer.Index++ }()
+	bf := bytes.NewBufferString("")
+	err := tmpl.Execute(bf, namer)
+	if err != nil {
+		return "", err
 	}
-	return fmt.Sprintf("%s.%s.%d", namer.dbName, namer.tableName, namer.chunkIndex)
+	return bf.String(), nil
 }
 
 func (f *CsvWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
 	log.Debug("start dumping table in csv format...", zap.String("table", ir.TableName()))
 
 	namer := newOutputFileNamer(ir)
-	fileName := fmt.Sprintf("%s.csv", namer.NextName())
+	fileName, err := namer.NextName(f.cfg.OutputFileTemplate)
+	if err != nil {
+		return err
+	}
+	fileName += ".csv"
 	chunksIter := buildChunksIter(ir, f.cfg.FileSize, f.cfg.StatementSize)
 	defer chunksIter.Rows().Close()
 
@@ -170,7 +185,11 @@ func (f *CsvWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
 		if f.cfg.FileSize == UnspecifiedSize {
 			break
 		}
-		fileName = fmt.Sprintf("%s.csv", namer.NextName())
+		fileName, err = namer.NextName(f.cfg.OutputFileTemplate)
+		if err != nil {
+			return err
+		}
+		fileName += ".csv"
 	}
 	log.Debug("dumping table in csv format successfully",
 		zap.String("table", ir.TableName()))


Test patch:
diff --git a/v4/export/writer_test.go b/v4/export/writer_test.go
index 3e7524d..563303d 100644
--- a/v4/export/writer_test.go
+++ b/v4/export/writer_test.go
@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path"
+	"text/template"
 
 	. "github.com/pingcap/check"
 )
@@ -158,6 +159,8 @@ func (s *testDumpSuite) TestWriteTableDataWithStatementSize(c *C) {
 	config := DefaultConfig()
 	config.OutputDirPath = dir
 	config.StatementSize = 50
+	config.OutputFileTemplate, err = template.New("filename").Parse("specified-name")
+	c.Assert(err, IsNil)
 	ctx := context.Background()
 	defer os.RemoveAll(config.OutputDirPath)
 
@@ -181,7 +184,7 @@ func (s *testDumpSuite) TestWriteTableDataWithStatementSize(c *C) {
 
 	// only with statement size
 	cases := map[string]string{
-		"test.employee.0.sql": "/*!40101 SET NAMES binary*/;\n" +
+		"specified-name.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" +
@@ -203,20 +206,23 @@ func (s *testDumpSuite) TestWriteTableDataWithStatementSize(c *C) {
 	// with file size and statement size
 	config.FileSize = 90
 	config.StatementSize = 30
+	// test specifying filename format
+	config.OutputFileTemplate, err = template.New("filename").Parse("{{.Index}}-{{.Table}}-{{.DB}}")
+	c.Assert(err, IsNil)
 	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" +
+		"0-employee-test.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" +
+		"1-employee-test.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",