Repository: pingcap/dumpling
Pull request #16 (2019-12-27 16:56:48)

support sorting output by primary key
Support sorting output by primary key.
Use [DATA-DOG/go-sqlmock](https://github.com/DATA-DOG/go-sqlmock) to mock `database/sql` related types.

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


```diff
@@             Coverage Diff             @@
##           master      #16       +/-   ##
===========================================
+ Coverage   27.61%   38.05%   +10.44%     
===========================================
  Files           9        9               
  Lines         431      515       +84     
===========================================
+ Hits          119      196       +77     
+ Misses        302      300        -2     
- Partials       10       19        +9
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index f44b5d2..559d3b9 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -18,6 +18,7 @@ import (
 	"fmt"
 	_ "net/http/pprof"
 	"os"
+	"time"
 
 	"github.com/pingcap/dumpling/v4/cli"
 	"github.com/pingcap/dumpling/v4/export"
@@ -56,8 +57,14 @@ func init() {
 	flag.Uint64Var(&fileSize, "F", export.UnspecifiedSize, "The approximate size of output file")
 	flag.Uint64Var(&fileSize, "filesize", export.UnspecifiedSize, "The approximate size of output file")
 
-	flag.StringVar(&outputDir, "output", ".", "Output directory")
-	flag.StringVar(&outputDir, "o", ".", "Output directory")
+	flag.StringVar(&outputDir, "output", defaultOutputDir, "Output directory")
+	flag.StringVar(&outputDir, "o", defaultOutputDir, "Output directory")
+}
+
+var defaultOutputDir = timestampDirName()
+
+func timestampDirName() string {
+	return fmt.Sprintf("./export-%s", time.Now().Format(time.RFC3339))
 }
 
 func main() {
diff --git a/go.mod b/go.mod
index 0972af2..203a2dc 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,8 @@
 module github.com/pingcap/dumpling
 
 require (
+	github.com/DATA-DOG/go-sqlmock v1.3.3
+	github.com/coreos/go-semver v0.3.0
 	github.com/go-sql-driver/mysql v1.4.1
 	github.com/pingcap/check v0.0.0-20191107115940-caf2b9e6ccf4
 )
diff --git a/go.sum b/go.sum
index 7b4543e..974d81a 100644
--- a/go.sum
+++ b/go.sum
@@ -1,10 +1,14 @@
-github.com/go-sql-driver/mysql v1.4.1 h1:g24URVg0OFbNUTx9qqY1IRZ9D9z3iPyi5zKhQZpNwpA=
-github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/DATA-DOG/go-sqlmock v1.3.3 h1:CWUqKXe0s8A2z6qCgkP4Kru7wC11YoAnoupUKFDnH08=
+github.com/DATA-DOG/go-sqlmock v1.3.3/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
+github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM=
+github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-sql-driver/mysql v1.4.1 h1:g24URVg0OFbNUTx9qqY1IRZ9D9z3iPyi5zKhQZpNwpA=
+github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
diff --git a/v4/export/config.go b/v4/export/config.go
index bc29ce7..40787b9 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -2,6 +2,10 @@ package export
 
 import (
 	"fmt"
+	"regexp"
+	"strings"
+
+	"github.com/coreos/go-semver/semver"
 )
 
 type Config struct {
@@ -15,6 +19,8 @@ type Config struct {
 	Logger        Logger
 	FileSize      uint64
 	OutputDirPath string
+	ServerInfo    ServerInfo
+	SortByPk      bool
 }
 
 func DefaultConfig() *Config {
@@ -28,6 +34,8 @@ func DefaultConfig() *Config {
 		Logger:        &DummyLogger{},
 		FileSize:      UnspecifiedSize,
 		OutputDirPath: ".",
+		ServerInfo:    ServerInfoUnknown,
+		SortByPk:      false,
 	}
 }
 
@@ -35,21 +43,56 @@ func (conf *Config) getDSN(db string) string {
 	return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s", conf.User, conf.Password, conf.Host, conf.Port, db)
 }
 
-func extractOutputConfig(conf *Config) *Config {
-	return &Config{
-		Logger:        conf.Logger,
-		FileSize:      conf.FileSize,
-		OutputDirPath: conf.OutputDirPath,
-	}
+const UnspecifiedSize = 0
+
+type ServerInfo struct {
+	ServerType    ServerType
+	ServerVersion *semver.Version
 }
 
-type WriteConfig struct {
-	// Logger is used to log the export routine.
-	Logger Logger
-	// Output size limit in bytes.
-	OutputSize int
-	// OutputDirPath is the directory to output.
-	OutputDirPath string
+var ServerInfoUnknown = ServerInfo{
+	ServerType:    ServerTypeUnknown,
+	ServerVersion: nil,
 }
 
-const UnspecifiedSize = 0
+var versionRegex = regexp.MustCompile(`^\d+\.\d+\.\d+([0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?`)
+var tidbVersionRegex = regexp.MustCompile(`v\d+\.\d+\.\d+([0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?`)
+
+func ParseServerInfo(src string) ServerInfo {
+	lowerCase := strings.ToLower(src)
+	serverInfo := ServerInfo{}
+	if strings.Contains(lowerCase, "tidb") {
+		serverInfo.ServerType = ServerTypeTiDB
+	} else if strings.Contains(lowerCase, "mariadb") {
+		serverInfo.ServerType = ServerTypeMariaDB
+	} else if versionRegex.MatchString(lowerCase) {
+		serverInfo.ServerType = ServerTypeMySQL
+	} else {
+		serverInfo.ServerType = ServerTypeUnknown
+	}
+
+	var versionStr string
+	if serverInfo.ServerType == ServerTypeTiDB {
+		versionStr = tidbVersionRegex.FindString(src)[1:]
+	} else {
+		versionStr = versionRegex.FindString(src)
+	}
+
+	var err error
+	serverInfo.ServerVersion, err = semver.NewVersion(versionStr)
+	if err != nil {
+		return serverInfo
+	}
+	return serverInfo
+}
+
+type ServerType int8
+
+const (
+	ServerTypeUnknown = iota
+	ServerTypeMySQL
+	ServerTypeMariaDB
+	ServerTypeTiDB
+
+	ServerTypeAll
+)
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 9b43011..067628e 100644
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -11,26 +11,14 @@ import (
 )
 
 func Dump(conf *Config) error {
-	pool, err := sql.Open("mysql", conf.getDSN(""))
+	databases, serverInfo, err := prepareMeta(conf)
 	if err != nil {
-		return withStack(err)
-	}
-
-	var databases []string
-	if conf.Database == "" {
-		var err error
-		databases, err = showDatabases(pool)
-		if err != nil {
-			pool.Close()
-			return err
-		}
-	} else {
-		databases = strings.Split(conf.Database, ",")
+		return err
 	}
-	pool.Close()
 
+	conf.ServerInfo = serverInfo
 	for _, database := range databases {
-		fsWriter, err := NewSimpleWriter(extractOutputConfig(conf))
+		fsWriter, err := NewSimpleWriter(conf)
 		if err != nil {
 			return err
 		}
@@ -41,6 +29,42 @@ func Dump(conf *Config) error {
 	return nil
 }
 
+func prepareMeta(conf *Config) (databases []string, info ServerInfo, err error) {
+	pool, err := sql.Open("mysql", conf.getDSN(""))
+	if err != nil {
+		return nil, ServerInfoUnknown, withStack(err)
+	}
+	defer pool.Close()
+	databases, err = getDumpingDatabaseNames(pool, conf)
+	if err != nil {
+		return nil, ServerInfoUnknown, withStack(err)
+	}
+	info, err = detectServerInfo(pool)
+	if err != nil {
+		return databases, ServerInfoUnknown, withStack(err)
+	}
+	return
+}
+
+func getDumpingDatabaseNames(pool *sql.DB, conf *Config) ([]string, error) {
+	if conf.Database == "" {
+		return showDatabases(pool)
+	}
+	return strings.Split(conf.Database, ","), nil
+}
+
+func detectServerInfo(db *sql.DB) (ServerInfo, error) {
+	var versionInfo string
+	handleOneRow := func(rows *sql.Rows) error {
+		return rows.Scan(&versionInfo)
+	}
+	err := simpleQuery(db, "SELECT version()", handleOneRow)
+	if err != nil {
+		return ServerInfoUnknown, withStack(err)
+	}
+	return ParseServerInfo(versionInfo), nil
+}
+
 func simpleQuery(db *sql.DB, sql string, handleOneRow func(*sql.Rows) error) error {
 	rows, err := db.Query(sql)
 	if err != nil {
@@ -151,7 +175,7 @@ func dumpDatabase(ctx context.Context, conf *Config, dbName string, writer Write
 			}
 
 			rateLimit.getToken()
-			tableIR, err := dumpTable(ctx, db, dbName, table)
+			tableIR, err := dumpTable(conf, db, dbName, table)
 			defer rateLimit.putToken()
 			if err != nil {
 				res[ith] = err
@@ -188,13 +212,17 @@ type tableDumper interface {
 	finishTable(ctx context.Context)
 }
 
-func dumpTable(ctx context.Context, db *sql.DB, database, table string) (TableDataIR, error) {
+func dumpTable(conf *Config, db *sql.DB, database, table string) (TableDataIR, error) {
 	colTypes, err := getColumnTypes(db, table)
 	if err != nil {
 		return nil, err
 	}
 
-	rows, err := db.Query(fmt.Sprintf("SELECT * from %s", table))
+	query, err := buildSelectAllQuery(conf, db, database, table)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.Query(query)
 	if err != nil {
 		return nil, withStack(err)
 	}
@@ -206,3 +234,82 @@ func dumpTable(ctx context.Context, db *sql.DB, database, table string) (TableDa
 		colTypes: colTypes,
 	}, nil
 }
+
+func buildSelectAllQuery(conf *Config, db *sql.DB, database, table string) (string, error) {
+	var query strings.Builder
+	query.WriteString("SELECT * FROM ")
+	query.WriteString(database)
+	query.WriteString(".")
+	query.WriteString(table)
+	if conf.SortByPk {
+		orderByClause, err := buildOrderByClause(conf, db, database, table)
+		if err != nil {
+			return "", err
+		}
+		if orderByClause != "" {
+			query.WriteString(" ")
+			query.WriteString(orderByClause)
+		}
+	}
+	return query.String(), nil
+}
+
+const errBadFieldCode = 1054
+
+func buildOrderByClause(conf *Config, db *sql.DB, database, table string) (string, error) {
+	if conf.ServerInfo.ServerType == ServerTypeTiDB {
+		tiDBRowIDQuery := fmt.Sprintf("SELECT _tidb_rowid from %s.%s LIMIT 0", database, table)
+		isBadField, otherErr := queryReturnExpectedErrorCode(db, tiDBRowIDQuery, errBadFieldCode)
+		if otherErr != nil {
+			return "", otherErr
+		}
+		if isBadField {
+			return "", nil
+		}
+		return "ORDER BY _tidb_rowid", nil
+	}
+	pkName, err := getPrimaryKeyName(db, database, table)
+	if err != nil {
+		return "", err
+	}
+	tableContainsPriKey := pkName != ""
+	if tableContainsPriKey {
+		return fmt.Sprintf("ORDER BY %s", pkName), nil
+	}
+	return "", nil
+}
+
+func queryReturnExpectedErrorCode(db *sql.DB, sql string, errCode int) (bool, error) {
+	_, err := db.Exec(sql)
+	if err != nil {
+		errMsg := strings.ToLower(err.Error())
+		if strings.Contains(errMsg, fmt.Sprintf("%d", errCode)) {
+			return true, nil
+		}
+		return false, withStack(err)
+	}
+	return false, nil
+}
+
+func getPrimaryKeyName(db *sql.DB, database, table string) (string, error) {
+	priKeyQuery := `SELECT column_name FROM information_schema.columns
+		WHERE table_schema = ? AND table_name = ? AND column_key = 'PRI';`
+	stmt, err := db.Prepare(priKeyQuery)
+	if err != nil {
+		return "", err
+	}
+	defer stmt.Close()
+	rows, err := stmt.Query(database, table)
+	if err != nil {
+		return "", withStack(err)
+	}
+
+	var colName string
+	for rows.Next() {
+		if err := rows.Scan(&colName); err != nil {
+			rows.Close()
+			return "", withStack(err)
+		}
+	}
+	return colName, nil
+}


Test patch:
diff --git a/v4/export/dump_test.go b/v4/export/dump_test.go
new file mode 100644
index 0000000..a196ce5
--- /dev/null
+++ b/v4/export/dump_test.go
@@ -0,0 +1,131 @@
+package export
+
+import (
+	"errors"
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/coreos/go-semver/semver"
+	. "github.com/pingcap/check"
+)
+
+var _ = Suite(&testDumpSuite{})
+
+type testDumpSuite struct{}
+
+func (s *testDumpSuite) TestDetectServerInfo(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	mkVer := makeVersion
+	data := [][]interface{}{
+		{1, "8.0.18", ServerTypeMySQL, mkVer(8, 0, 18, "")},
+		{2, "10.4.10-MariaDB-1:10.4.10+maria~bionic", ServerTypeMariaDB, mkVer(10, 4, 10, "MariaDB-1")},
+		{3, "5.7.25-TiDB-v4.0.0-alpha-1263-g635f2e1af", ServerTypeTiDB, mkVer(4, 0, 0, "alpha-1263-g635f2e1af")},
+		{4, "5.7.25-TiDB-v3.0.7-58-g6adce2367", ServerTypeTiDB, mkVer(3, 0, 7, "58-g6adce2367")},
+		{5, "invalid version", ServerTypeUnknown, (*semver.Version)(nil)},
+	}
+	dec := func(d []interface{}) (tag int, verStr string, tp ServerType, v *semver.Version) {
+		return d[0].(int), d[1].(string), ServerType(d[2].(int)), d[3].(*semver.Version)
+	}
+
+	for _, datum := range data {
+		tag, r, serverTp, expectVer := dec(datum)
+		cmt := Commentf("test case number: %d", tag)
+
+		rows := sqlmock.NewRows([]string{"version"}).AddRow(r)
+		mock.ExpectQuery("SELECT version()").WillReturnRows(rows)
+
+		info, err := detectServerInfo(db)
+		c.Assert(err, IsNil, cmt)
+		c.Assert(info.ServerType, Equals, serverTp, cmt)
+		c.Assert(info.ServerVersion == nil, Equals, expectVer == nil, cmt)
+		if info.ServerVersion == nil {
+			c.Assert(expectVer, IsNil, cmt)
+		} else {
+			c.Assert(info.ServerVersion.Equal(*expectVer), IsTrue)
+		}
+		c.Assert(mock.ExpectationsWereMet(), IsNil, cmt)
+	}
+}
+
+func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	mockConf := DefaultConfig()
+	mockConf.SortByPk = true
+
+	// Test when the server is TiDB.
+	mockConf.ServerInfo.ServerType = ServerTypeTiDB
+
+	// _tidb_rowid is available.
+	mock.ExpectExec("SELECT _tidb_rowid from test.t").
+		WillReturnResult(sqlmock.NewResult(0, 0))
+	q, err := buildSelectAllQuery(mockConf, db, "test", "t")
+	c.Assert(q, Equals, "SELECT * FROM test.t ORDER BY _tidb_rowid")
+
+	// _tidb_rowid is unavailable, or PKIsHandle.
+	mock.ExpectExec("SELECT _tidb_rowid from test.t").
+		WillReturnError(errors.New(`1054, "Unknown column '_tidb_rowid' in 'field list'"`))
+	q, err = buildSelectAllQuery(mockConf, db, "test", "t")
+	c.Assert(q, Equals, "SELECT * FROM test.t")
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+
+	// Test other server.
+	otherServers := []ServerType{ServerTypeUnknown, ServerTypeMySQL, ServerTypeMariaDB}
+
+	// Test table with primary key.
+	for _, serverTp := range otherServers {
+		mockConf.ServerInfo.ServerType = serverTp
+		cmt := Commentf("server type: %s", serverTp)
+		mock.ExpectPrepare("SELECT column_name FROM information_schema.columns").
+			ExpectQuery().WithArgs("test", "t").
+			WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id"))
+
+		q, err := buildSelectAllQuery(mockConf, db, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(q, Equals, "SELECT * FROM test.t ORDER BY id", cmt)
+		err = mock.ExpectationsWereMet()
+		c.Assert(err, IsNil, cmt)
+		c.Assert(mock.ExpectationsWereMet(), IsNil, cmt)
+	}
+
+	// Test table without primary key.
+	for _, serverTp := range otherServers {
+		mockConf.ServerInfo.ServerType = serverTp
+		cmt := Commentf("server type: %s", serverTp)
+		mock.ExpectPrepare("SELECT column_name FROM information_schema.columns").
+			ExpectQuery().WithArgs("test", "t").
+			WillReturnRows(sqlmock.NewRows([]string{"column_name"}))
+
+		q, err := buildSelectAllQuery(mockConf, db, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(q, Equals, "SELECT * FROM test.t", cmt)
+		err = mock.ExpectationsWereMet()
+		c.Assert(err, IsNil, cmt)
+		c.Assert(mock.ExpectationsWereMet(), IsNil)
+	}
+
+	// Test when config.SortByPk is disabled.
+	mockConf.SortByPk = false
+	for tp := ServerTypeUnknown; tp < ServerTypeAll; tp += 1 {
+		mockConf.ServerInfo.ServerType = ServerType(tp)
+		cmt := Commentf("current server type: ", tp)
+
+		q, err := buildSelectAllQuery(mockConf, db, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(q, Equals, "SELECT * FROM test.t", cmt)
+		c.Assert(mock.ExpectationsWereMet(), IsNil, cmt)
+	}
+}
+
+func makeVersion(major, minor, patch int64, preRelease string) *semver.Version {
+	return &semver.Version{
+		Major:      major,
+		Minor:      minor,
+		Patch:      patch,
+		PreRelease: semver.PreRelease(preRelease),
+		Metadata:   "",
+	}
+}