From 9559e983dca9842b8bf66ccb5b83c7acef1e96f3 Mon Sep 17 00:00:00 2001 From: Chunzhu Li Date: Tue, 14 Jul 2020 15:50:03 +0800 Subject: [PATCH] support specify ouput filename format (#122) * support specify ouput filename format * revise variable name * address comments * try to address comment * address comment --- dumpling/cmd/dumpling/main.go | 23 ++++++++-- dumpling/v4/export/config.go | 74 ++++++++++++++++--------------- dumpling/v4/export/prepare.go | 8 ++++ dumpling/v4/export/writer.go | 51 ++++++++++++++------- dumpling/v4/export/writer_test.go | 12 +++-- 5 files changed, 111 insertions(+), 57 deletions(-) diff --git a/dumpling/cmd/dumpling/main.go b/dumpling/cmd/dumpling/main.go index 5fd83920..52816fe0 100644 --- a/dumpling/cmd/dumpling/main.go +++ b/dumpling/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/dumpling/v4/export/config.go b/dumpling/v4/export/config.go index 85ffd936..d1bf8242 100644 --- a/dumpling/v4/export/config.go +++ b/dumpling/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/dumpling/v4/export/prepare.go b/dumpling/v4/export/prepare.go index 81d1516f..7bbc5ac5 100644 --- a/dumpling/v4/export/prepare.go +++ b/dumpling/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/dumpling/v4/export/writer.go b/dumpling/v4/export/writer.go index d28b2324..126aa4a1 100644 --- a/dumpling/v4/export/writer.go +++ b/dumpling/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())) diff --git a/dumpling/v4/export/writer_test.go b/dumpling/v4/export/writer_test.go index 3e7524dc..563303db 100644 --- a/dumpling/v4/export/writer_test.go +++ b/dumpling/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",