Skip to content

Commit

Permalink
support specify ouput filename format (pingcap#122)
Browse files Browse the repository at this point in the history 
* support specify ouput filename format

* revise variable name

* address comments

* try to address comment

* address comment
  • Loading branch information
@lichunzhu
lichunzhu authored Jul 14, 2020
1 parent 19ba921 commit 9559e98
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 57 deletions.
23 changes: 20 additions & 3 deletions dumpling/cmd/dumpling/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"os"
"strconv"
"strings"
"text/template"
"time"

"github.com/docker/go-units"
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down
74 changes: 39 additions & 35 deletions dumpling/v4/export/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"regexp"
"strings"
"text/template"
"time"

"github.com/coreos/go-semver/semver"
Expand Down Expand Up @@ -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,
}
}

Expand Down
8 changes: 8 additions & 0 deletions dumpling/v4/export/prepare.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package export
import (
"database/sql"
"strings"
"text/template"

"github.com/go-sql-driver/mysql"
"github.com/pingcap/dumpling/v4/log"
Expand Down Expand Up @@ -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
Expand Down
51 changes: 35 additions & 16 deletions dumpling/v4/export/writer.go
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
package export

import (
"bytes"
"context"
"fmt"
"os"
"path"
"text/template"

"go.uber.org/zap"

Expand Down Expand Up @@ -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()

Expand All @@ -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()))
Expand Down Expand Up @@ -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 {
Expand All @@ -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()

Expand All @@ -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()))
Expand Down
12 changes: 9 additions & 3 deletions dumpling/v4/export/writer_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"io/ioutil"
"os"
"path"
"text/template"

. "github.com/pingcap/check"
)
Expand Down Expand Up @@ -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)

Expand All @@ -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" +
Expand All @@ -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",
Expand Down

0 comments on commit 9559e98

Please sign in to comment.