Repository: pingcap/dumpling
Pull request #103 (2020-06-18 13:23:19)

Refine dumpling user experience
1. add dump-empty-database option 
2. check whether database exists before dump
3. allow filetype sql for --Sql command



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


```diff
@@            Coverage Diff             @@
##           master     #103      +/-   ##
==========================================
+ Coverage   56.30%   56.87%   +0.56%     
==========================================
  Files          16       16              
  Lines        1593     1607      +14     
==========================================
+ Hits          897      914      +17     
+ Misses        623      619       -4     
- Partials       73       74       +1     
```

kennytm: @cyliu0 can the current PR fix https://github.com/pingcap/dumpling/issues/104#issuecomment-646445174?

Patch:
diff --git a/README.md b/README.md
index 45ebe61..f945139 100644
--- a/README.md
+++ b/README.md
@@ -20,7 +20,7 @@ Features
 - [x] Export multiple tables in parallel to speed up execution.
 - [ ] Multiple output formats: SQL, CSV, ...
 - [ ] Write to cloud storage (S3, GCS) natively
-- [ ] Advanced table filtering via black-/white-lists
+- [x] Advanced table filtering
 
 Building
 --------
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index 8408ddd..d54ce16 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -61,7 +61,9 @@ var (
 	filters       []string
 	caseSensitive bool
 
-	escapeBackslash bool
+	dumpEmptyDatabase bool
+	escapeBackslash   bool
+	tidbMemQuotaQuery uint64
 )
 
 var defaultOutputDir = timestampDirName()
@@ -105,6 +107,8 @@ func main() {
 	pflag.StringVarP(&sql, "sql", "S", "", "Dump data with given sql. This argument doesn't support concurrent dump")
 	pflag.StringArrayVarP(&filters, "filter", "f", []string{"*.*"}, "filter to select which tables to dump")
 	pflag.BoolVar(&caseSensitive, "case-sensitive", false, "whether the filter should be case-sensitive")
+	pflag.BoolVar(&dumpEmptyDatabase, "dump-empty-database", true, "whether to dump empty database")
+	pflag.Uint64Var(&tidbMemQuotaQuery, "tidb-mem-quota-query", export.DefaultTiDBMemQuotaQuery, "The maximum memory limit for a single SQL statement, in bytes. Default: 32GB")
 
 	printVersion := pflag.BoolP("version", "V", false, "Print Dumpling version")
 
@@ -159,6 +163,7 @@ func main() {
 	conf.Rows = rows
 	conf.Where = where
 	conf.EscapeBackslash = escapeBackslash
+	conf.DumpEmptyDatabase = dumpEmptyDatabase
 	conf.LogLevel = logLevel
 	conf.LogFile = logFile
 	conf.LogFormat = logFormat
@@ -170,6 +175,7 @@ func main() {
 	conf.CsvNullValue = csvNullValue
 	conf.Sql = sql
 	conf.TableFilter = tableFilter
+	conf.TiDBMemQuotaQuery = tidbMemQuotaQuery
 
 	err = export.Dump(context.Background(), conf)
 	if err != nil {
diff --git a/v4/export/block_allow_list.go b/v4/export/block_allow_list.go
index f85652f..0068025 100644
--- a/v4/export/block_allow_list.go
+++ b/v4/export/block_allow_list.go
@@ -37,8 +37,10 @@ func filterTables(conf *Config) {
 		}
 		// 1. this dbName doesn't match block allow list, don't add
 		// 2. this dbName matches block allow list, but there is no table in this database, add
-		if _, ok := dbTables[dbName]; !ok && conf.TableFilter.MatchSchema(dbName) {
-			dbTables[dbName] = make([]*TableInfo, 0)
+		if conf.DumpEmptyDatabase {
+			if _, ok := dbTables[dbName]; !ok && conf.TableFilter.MatchSchema(dbName) {
+				dbTables[dbName] = make([]*TableInfo, 0)
+			}
 		}
 	}
 
diff --git a/v4/export/config.go b/v4/export/config.go
index 32f5ce2..f524e3c 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -41,51 +41,60 @@ type Config struct {
 	CsvNullValue  string
 	Sql           string
 
-	TableFilter     filter.Filter
-	Rows            uint64
-	Where           string
-	FileType        string
-	EscapeBackslash bool
+	TableFilter       filter.Filter
+	Rows              uint64
+	Where             string
+	FileType          string
+	EscapeBackslash   bool
+	DumpEmptyDatabase bool
+	TiDBMemQuotaQuery uint64
 }
 
 func DefaultConfig() *Config {
 	allFilter, _ := filter.Parse([]string{"*.*"})
 	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,
+		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,
+		TiDBMemQuotaQuery: UnspecifiedSize,
 	}
 }
 
 func (conf *Config) getDSN(db string) string {
-	return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4", conf.User, conf.Password, conf.Host, conf.Port, db)
+	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4", conf.User, conf.Password, conf.Host, conf.Port, db)
+	if conf.TiDBMemQuotaQuery != UnspecifiedSize {
+		dsn += fmt.Sprintf("&tidb_mem_quota_query=%v", conf.TiDBMemQuotaQuery)
+	}
+	return dsn
 }
 
 const (
 	UnspecifiedSize            = 0
+	DefaultTiDBMemQuotaQuery   = 32 * (1 << 30)
 	defaultDumpThreads         = 128
 	defaultDumpGCSafePointTTL  = 5 * 60
 	dumplingServiceSafePointID = "dumpling"
diff --git a/v4/export/dump.go b/v4/export/dump.go
index e0f8dd1..2777659 100644
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -36,7 +36,19 @@ func Dump(pCtx context.Context, conf *Config) (err error) {
 
 	conf.ServerInfo, err = detectServerInfo(pool)
 	if err != nil {
-		return err
+		if strings.Contains(err.Error(), "tidb_mem_quota_query") {
+			conf.TiDBMemQuotaQuery = UnspecifiedSize
+			pool, err = sql.Open("mysql", conf.getDSN(""))
+			if err != nil {
+				return withStack(err)
+			}
+			conf.ServerInfo, err = detectServerInfo(pool)
+			if err != nil {
+				return withStack(err)
+			}
+		} else {
+			return withStack(err)
+		}
 	}
 
 	databases, err := prepareDumpingDatabases(conf, pool)
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index 8940110..cc50494 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -289,6 +289,7 @@ func splitTableDataIntoChunks(
 	}
 
 	count := estimateCount(dbName, tableName, db, field, conf)
+	log.Info("get estimated rows count", zap.Uint64("estimateCount", count))
 	if count < conf.Rows {
 		// skip chunk logic if estimates are low
 		log.Debug("skip concurrent dump due to estimate count < rows",
diff --git a/v4/export/prepare.go b/v4/export/prepare.go
index ae216fe..4947cee 100644
--- a/v4/export/prepare.go
+++ b/v4/export/prepare.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/pingcap/dumpling/v4/log"
+	"github.com/pingcap/errors"
 )
 
 func adjustConfig(conf *Config) error {
@@ -39,9 +40,23 @@ func detectServerInfo(db *sql.DB) (ServerInfo, error) {
 }
 
 func prepareDumpingDatabases(conf *Config, db *sql.DB) ([]string, error) {
+	databases, err := ShowDatabases(db)
 	if len(conf.Databases) == 0 {
-		return ShowDatabases(db)
+		return databases, err
 	} else {
+		dbMap := make(map[string]interface{}, len(databases))
+		for _, database := range databases {
+			dbMap[database] = struct{}{}
+		}
+		var notExistsDatabases []string
+		for _, database := range conf.Databases {
+			if _, ok := dbMap[database]; !ok {
+				notExistsDatabases = append(notExistsDatabases, database)
+			}
+		}
+		if len(notExistsDatabases) > 0 {
+			return nil, errors.Errorf("Unknown databases [%s]", strings.Join(notExistsDatabases, ","))
+		}
 		return conf.Databases, nil
 	}
 }
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 98fd217..2f258e5 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -142,12 +142,19 @@ func SelectFromSql(conf *Config, db *sql.DB) (TableDataIR, error) {
 	if err != nil {
 		return nil, withStack(errors.WithMessage(err, conf.Sql))
 	}
+	cols, err := rows.Columns()
+	if err != nil {
+		return nil, withStack(errors.WithMessage(err, conf.Sql))
+	}
+	for i := range cols {
+		cols[i] = wrapBackTicks(cols[i])
+	}
 	return &tableData{
 		database:        "",
 		table:           "",
 		rows:            rows,
 		colTypes:        colTypes,
-		selectedField:   "",
+		selectedField:   strings.Join(cols, ","),
 		escapeBackslash: conf.EscapeBackslash,
 		specCmts: []string{
 			"/*!40101 SET NAMES binary*/;",
@@ -502,12 +509,12 @@ func estimateCount(dbName, tableName string, db *sql.DB, field string, conf *Con
 		+----+-------------+-------+------------+-------+---------------+-----------+---------+------+------+----------+-------------+
 	*/
 	if estRows > 0 {
-		return uint64(estRows)
+		return estRows
 	}
 	return 0
 }
 
-func detectEstimateRows(db *sql.DB, query string, fieldNames []string) int {
+func detectEstimateRows(db *sql.DB, query string, fieldNames []string) uint64 {
 	row, err := db.Query(query)
 	if err != nil {
 		log.Warn("can't execute query from db",
@@ -544,7 +551,7 @@ found:
 			zap.String("query", query), zap.Error(err))
 		return 0
 	}
-	return int(estRows)
+	return uint64(estRows)
 }
 
 func buildWhereCondition(conf *Config, where string) string {
diff --git a/v4/export/writer.go b/v4/export/writer.go
index 3d70e0f..234bbbd 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -6,8 +6,6 @@ import (
 	"os"
 	"path"
 
-	"github.com/pingcap/errors"
-
 	"go.uber.org/zap"
 
 	"github.com/pingcap/dumpling/v4/log"
@@ -24,10 +22,6 @@ type SimpleWriter struct {
 }
 
 func NewSimpleWriter(config *Config) (Writer, error) {
-	if config.Sql != "" {
-		log.Error("unsupported dump data in sql format", zap.String("sql", config.Sql))
-		return nil, errors.New("unsupported dump data in sql format when specific sql")
-	}
 	sw := &SimpleWriter{cfg: config}
 	return sw, os.MkdirAll(config.OutputDirPath, 0755)
 }
@@ -47,8 +41,6 @@ func (f *SimpleWriter) WriteTableMeta(ctx context.Context, db, table, createSQL
 func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error {
 	log.Debug("start dumping table...", zap.String("table", ir.TableName()))
 
-	chunkIndex := ir.ChunkIndex()
-	fileName := fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), ir.ChunkIndex())
 	// just let `database.table.sql` be `database.table.0.sql`
 	/*if fileName == "" {
 		// set initial file name
@@ -57,6 +49,9 @@ 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())
 	chunksIter := buildChunksIter(ir, f.cfg.FileSize, f.cfg.StatementSize)
 	defer chunksIter.Rows().Close()
 
@@ -76,8 +71,7 @@ func (f *SimpleWriter) WriteTableData(ctx context.Context, ir TableDataIR) error
 		if f.cfg.FileSize == UnspecifiedSize {
 			break
 		}
-		chunkIndex += 1
-		fileName = fmt.Sprintf("%s.%s.%d.sql", ir.DatabaseName(), ir.TableName(), chunkIndex)
+		fileName = fmt.Sprintf("%s.sql", namer.NextName())
 	}
 	log.Debug("dumping table successfully",
 		zap.String("table", ir.TableName()))


Test patch:
diff --git a/tests/no_table_and_db_name/run.sh b/tests/no_table_and_db_name/run.sh
index 240ee39..0f5c4bf 100644
--- a/tests/no_table_and_db_name/run.sh
+++ b/tests/no_table_and_db_name/run.sh
@@ -28,4 +28,13 @@ assert [ $( ls -lh $DUMPLING_OUTPUT_DIR | grep -e ".csv$" | wc -l ) -eq 10 ]
 # 10 files with header.
 assert [ $( cat $DUMPLING_OUTPUT_DIR/*.csv | wc -l ) -eq $(( 100 + 10 )) ]
 
+
+# dumping with file size = 200 bytes
+run_dumpling -F 200B --filetype sql --sql "select * from $TEST_NAME.t"
+
+assert [ $( ls -lh $DUMPLING_OUTPUT_DIR | grep -e ".sql$" | wc -l ) -eq 10 ]
+
+# 10 files with header.
+assert [ $( cat $DUMPLING_OUTPUT_DIR/*.sql | wc -l ) -eq $(( 100 + 10 * 2 )) ]
+
 echo "TEST: [$TEST_NAME] passed."
\ No newline at end of file
diff --git a/v4/export/block_allow_list_test.go b/v4/export/block_allow_list_test.go
index eb32f15..19c72d6 100644
--- a/v4/export/block_allow_list_test.go
+++ b/v4/export/block_allow_list_test.go
@@ -47,8 +47,9 @@ func (s *testBWListSuite) TestFilterDatabaseWithNoTable(c *C) {
 		ServerInfo: ServerInfo{
 			ServerType: ServerTypeTiDB,
 		},
-		Tables:      dbTables,
-		TableFilter: tf.NewSchemasFilter("yyy"),
+		Tables:            dbTables,
+		TableFilter:       tf.NewSchemasFilter("yyy"),
+		DumpEmptyDatabase: true,
 	}
 	filterTables(conf)
 	c.Assert(conf.Tables, HasLen, 0)
@@ -60,4 +61,12 @@ func (s *testBWListSuite) TestFilterDatabaseWithNoTable(c *C) {
 	filterTables(conf)
 	c.Assert(conf.Tables, HasLen, 1)
 	c.Assert(conf.Tables, DeepEquals, expectedDBTables)
+
+	dbTables["xxx"] = []*TableInfo{}
+	expectedDBTables = DatabaseTables{}
+	conf.Tables = dbTables
+	conf.DumpEmptyDatabase = false
+	filterTables(conf)
+	c.Assert(conf.Tables, HasLen, 0)
+	c.Assert(conf.Tables, DeepEquals, expectedDBTables)
 }
diff --git a/v4/export/prepare_test.go b/v4/export/prepare_test.go
index 6056a1c..395de9f 100644
--- a/v4/export/prepare_test.go
+++ b/v4/export/prepare_test.go
@@ -16,6 +16,12 @@ func (s *testPrepareSuite) TestPrepareDumpingDatabases(c *C) {
 	c.Assert(err, IsNil)
 	defer db.Close()
 
+	rows := sqlmock.NewRows([]string{"Database"}).
+		AddRow("db1").
+		AddRow("db2").
+		AddRow("db3").
+		AddRow("db5")
+	mock.ExpectQuery("SHOW DATABASES").WillReturnRows(rows)
 	conf := DefaultConfig()
 	conf.Databases = []string{"db1", "db2", "db3"}
 	result, err := prepareDumpingDatabases(conf, db)
@@ -23,7 +29,7 @@ func (s *testPrepareSuite) TestPrepareDumpingDatabases(c *C) {
 	c.Assert(result, DeepEquals, []string{"db1", "db2", "db3"})
 
 	conf.Databases = nil
-	rows := sqlmock.NewRows([]string{"Database"}).
+	rows = sqlmock.NewRows([]string{"Database"}).
 		AddRow("db1").
 		AddRow("db2")
 	mock.ExpectQuery("SHOW DATABASES").WillReturnRows(rows)
@@ -34,6 +40,16 @@ func (s *testPrepareSuite) TestPrepareDumpingDatabases(c *C) {
 	mock.ExpectQuery("SHOW DATABASES").WillReturnError(fmt.Errorf("err"))
 	_, err = prepareDumpingDatabases(conf, db)
 	c.Assert(err, NotNil)
+
+	rows = sqlmock.NewRows([]string{"Database"}).
+		AddRow("db1").
+		AddRow("db2").
+		AddRow("db3").
+		AddRow("db5")
+	mock.ExpectQuery("SHOW DATABASES").WillReturnRows(rows)
+	conf.Databases = []string{"db1", "db2", "db4", "db6"}
+	_, err = prepareDumpingDatabases(conf, db)
+	c.Assert(err, ErrorMatches, `Unknown databases \[db4,db6\]`)
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 }