Repository: pingcap/dumpling
Pull request #24 (2020-01-09 02:51:00)

consistency: move code to a new file and add unit tests
Changes:
- Add command line option `consistency` and `snapshot`
- Move consistency-related code to a new file `consistency.go`
- Add unit tests for `consistency.go`

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


```diff
@@             Coverage Diff             @@
##           master      #24       +/-   ##
===========================================
+ Coverage   42.05%   53.52%   +11.46%     
===========================================
  Files          12       13        +1     
  Lines         642      667       +25     
===========================================
+ Hits          270      357       +87     
+ Misses        339      272       -67     
- Partials       33       38        +5
```


Patch:
diff --git a/cmd/dumpling/main.go b/cmd/dumpling/main.go
index eaace85..bd703ee 100644
--- a/cmd/dumpling/main.go
+++ b/cmd/dumpling/main.go
@@ -27,15 +27,17 @@ import (
 )
 
 var (
-	database  string
-	host      string
-	user      string
-	port      int
-	password  string
-	threads   int
-	outputDir string
-	fileSize  uint64
-	logLevel  string
+	database    string
+	host        string
+	user        string
+	port        int
+	password    string
+	threads     int
+	outputDir   string
+	fileSize    uint64
+	logLevel    string
+	consistency string
+	snapshot    string
 )
 
 func init() {
@@ -63,7 +65,11 @@ func init() {
 	flag.StringVar(&outputDir, "output", defaultOutputDir, "Output directory")
 	flag.StringVar(&outputDir, "o", defaultOutputDir, "Output directory")
 
-	flag.StringVar(&logLevel, "loglevel", "info", "Log level {debug|info|warn|error|dpanic|panic|fatal}")
+	flag.StringVar(&logLevel, "loglevel", "info", "Log level: {debug|info|warn|error|dpanic|panic|fatal}")
+
+	flag.StringVar(&consistency, "consistency", "auto", "Consistency level during dumping: {auto|none|flush|lock|snapshot}")
+
+	flag.StringVar(&snapshot, "snapshot", "", "Snapshot position. Valid only when consistency=snapshot")
 }
 
 var defaultOutputDir = timestampDirName()
diff --git a/v4/export/config.go b/v4/export/config.go
index 2ea2ed2..c9a6edf 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -24,6 +24,7 @@ type Config struct {
 	SortByPk      bool
 	Tables        DatabaseTables
 	Snapshot      string
+	Consistency   string
 }
 
 func DefaultConfig() *Config {
@@ -40,6 +41,7 @@ func DefaultConfig() *Config {
 		SortByPk:      false,
 		Tables:        nil,
 		Snapshot:      "",
+		Consistency:   "auto",
 	}
 }
 
diff --git a/v4/export/consistency.go b/v4/export/consistency.go
new file mode 100644
index 0000000..5318ce5
--- /dev/null
+++ b/v4/export/consistency.go
@@ -0,0 +1,134 @@
+package export
+
+import (
+	"database/sql"
+	"errors"
+	"fmt"
+)
+
+func NewConsistencyController(conf *Config, session *sql.DB) (ConsistencyController, error) {
+	resolveAutoConsistency(conf)
+	switch conf.Consistency {
+	case "flush":
+		return &ConsistencyFlushTableWithReadLock{
+			serverType: conf.ServerInfo.ServerType,
+			db:         session,
+		}, nil
+	case "lock":
+		return &ConsistencyLockDumpingTables{
+			db:        session,
+			allTables: conf.Tables,
+		}, nil
+	case "snapshot":
+		return &ConsistencySnapshot{
+			serverType: conf.ServerInfo.ServerType,
+			snapshot:   conf.Snapshot,
+			db:         session,
+		}, nil
+	case "none":
+		return &ConsistencyNone{}, nil
+	default:
+		return nil, withStack(fmt.Errorf("invalid consistency option %s", conf.Consistency))
+	}
+}
+
+type ConsistencyController interface {
+	Setup() error
+	TearDown() error
+}
+
+type ConsistencyNone struct{}
+
+func (c *ConsistencyNone) Setup() error {
+	return nil
+}
+
+func (c *ConsistencyNone) TearDown() error {
+	return nil
+}
+
+type ConsistencyFlushTableWithReadLock struct {
+	serverType ServerType
+	db         *sql.DB
+}
+
+func (c *ConsistencyFlushTableWithReadLock) Setup() error {
+	if c.serverType == ServerTypeTiDB {
+		return withStack(errors.New("'flush table with read lock' cannot be used to ensure the consistency in TiDB"))
+	}
+	return FlushTableWithReadLock(c.db)
+}
+
+func (c *ConsistencyFlushTableWithReadLock) TearDown() error {
+	err := c.db.Ping()
+	if err != nil {
+		return withStack(errors.New("ConsistencyFlushTableWithReadLock lost database connection"))
+	}
+	return UnlockTables(c.db)
+}
+
+type ConsistencyLockDumpingTables struct {
+	db        *sql.DB
+	allTables DatabaseTables
+}
+
+func (c *ConsistencyLockDumpingTables) Setup() error {
+	for dbName, tables := range c.allTables {
+		for _, table := range tables {
+			err := LockTables(c.db, dbName, table)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (c *ConsistencyLockDumpingTables) TearDown() error {
+	err := c.db.Ping()
+	if err != nil {
+		return withStack(errors.New("ConsistencyLockDumpingTables lost database connection"))
+	}
+	return UnlockTables(c.db)
+}
+
+type ConsistencySnapshot struct {
+	serverType ServerType
+	snapshot   string
+	db         *sql.DB
+}
+
+const showMasterStatusFieldNum = 5
+const snapshotFieldIndex = 1
+
+func (c *ConsistencySnapshot) Setup() error {
+	if c.serverType != ServerTypeTiDB {
+		return withStack(errors.New("snapshot consistency is not supported for this server"))
+	}
+	if c.snapshot == "" {
+		str, err := ShowMasterStatus(c.db, showMasterStatusFieldNum)
+		if err != nil {
+			return err
+		}
+		c.snapshot = str[snapshotFieldIndex]
+	}
+	return SetTiDBSnapshot(c.db, c.snapshot)
+}
+
+func (c *ConsistencySnapshot) TearDown() error {
+	return nil
+}
+
+func resolveAutoConsistency(conf *Config) {
+	if conf.Consistency != "auto" {
+		return
+	}
+	switch conf.ServerInfo.ServerType {
+	case ServerTypeTiDB:
+		conf.Consistency = "snapshot"
+	case ServerTypeMySQL, ServerTypeMariaDB:
+		conf.Consistency = "flush"
+	default:
+		conf.Consistency = "none"
+	}
+}
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 8dcba4a..7d71cc2 100644
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -30,7 +30,10 @@ func Dump(conf *Config) (err error) {
 		return err
 	}
 
-	var conCtrl ConsistencyController = &ConsistencyNone{}
+	conCtrl, err := NewConsistencyController(conf, pool)
+	if err != nil {
+		return err
+	}
 	if err = conCtrl.Setup(); err != nil {
 		return err
 	}
diff --git a/v4/export/prepare.go b/v4/export/prepare.go
index 4f97bb4..2349cc3 100644
--- a/v4/export/prepare.go
+++ b/v4/export/prepare.go
@@ -2,7 +2,6 @@ package export
 
 import (
 	"database/sql"
-	"errors"
 	"strings"
 
 	"github.com/pingcap/dumpling/v4/log"
@@ -40,90 +39,3 @@ func listAllTables(db *sql.DB, databaseNames []string) (DatabaseTables, error) {
 	}
 	return dbTables, nil
 }
-
-type ConsistencyController interface {
-	Setup() error
-	TearDown() error
-}
-
-type ConsistencyNone struct{}
-
-func (c *ConsistencyNone) Setup() error {
-	return nil
-}
-
-func (c *ConsistencyNone) TearDown() error {
-	return nil
-}
-
-type ConsistencyFlushTableWithReadLock struct {
-	serverType ServerType
-	db         *sql.DB
-}
-
-func (c *ConsistencyFlushTableWithReadLock) Setup() error {
-	if c.serverType == ServerTypeTiDB {
-		return withStack(errors.New("'flush table with read lock' cannot be used to ensure the consistency in TiDB"))
-	}
-	return FlushTableWithReadLock(c.db)
-}
-
-func (c *ConsistencyFlushTableWithReadLock) TearDown() error {
-	err := c.db.Ping()
-	if err != nil {
-		return withStack(errors.New("ConsistencyFlushTableWithReadLock lost database connection"))
-	}
-	return UnlockTables(c.db)
-}
-
-type ConsistencyLockDumpingTables struct {
-	db        *sql.DB
-	allTables DatabaseTables
-}
-
-func (c *ConsistencyLockDumpingTables) Setup() error {
-	for dbName, tables := range c.allTables {
-		for _, table := range tables {
-			err := LockTables(c.db, dbName, table)
-			if err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}
-
-func (c *ConsistencyLockDumpingTables) TearDown() error {
-	err := c.db.Ping()
-	if err != nil {
-		return withStack(errors.New("ConsistencyLockDumpingTables lost database connection"))
-	}
-	return UnlockTables(c.db)
-}
-
-type ConsistencySnapshot struct {
-	serverType ServerType
-	snapshot   string
-	db         *sql.DB
-}
-
-const showMasterStatusFieldNum = 5
-const snapshotFieldIndex = 1
-
-func (c *ConsistencySnapshot) Setup() error {
-	if c.serverType != ServerTypeTiDB {
-		return withStack(errors.New("snapshot consistency is not supported for this server"))
-	}
-	if c.snapshot == "" {
-		str, err := ShowMasterStatus(c.db, showMasterStatusFieldNum)
-		if err != nil {
-			return err
-		}
-		c.snapshot = str[snapshotFieldIndex]
-	}
-	return SetTiDBSnapshot(c.db, c.snapshot)
-}
-
-func (c *ConsistencySnapshot) TearDown() error {
-	return nil
-}
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 41b9635..0ac9de5 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -178,7 +178,7 @@ func FlushTableWithReadLock(db *sql.DB) error {
 }
 
 func LockTables(db *sql.DB, database, table string) error {
-	_, err := db.Exec(fmt.Sprintf("`%s`.`%s`", database, table))
+	_, err := db.Exec(fmt.Sprintf("LOCK TABLES `%s`.`%s` READ", database, table))
 	return withStack(err)
 }
 


Test patch:
diff --git a/v4/export/consistency_test.go b/v4/export/consistency_test.go
new file mode 100644
index 0000000..c2bc3cf
--- /dev/null
+++ b/v4/export/consistency_test.go
@@ -0,0 +1,137 @@
+package export
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	. "github.com/pingcap/check"
+)
+
+var _ = Suite(&testConsistencySuite{})
+
+type testConsistencySuite struct{}
+
+func (s *testConsistencySuite) assertNil(err error, c *C) {
+	if err != nil {
+		c.Fatalf(err.Error())
+	}
+}
+
+func (s *testConsistencySuite) assertLifetimeErrNil(ctrl ConsistencyController, c *C) {
+	s.assertNil(ctrl.Setup(), c)
+	s.assertNil(ctrl.TearDown(), c)
+}
+
+func (s *testConsistencySuite) TestConsistencyController(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+	conf := DefaultConfig()
+	resultOk := sqlmock.NewResult(0, 1)
+
+	conf.Consistency = "none"
+	ctrl, _ := NewConsistencyController(conf, db)
+	_, ok := ctrl.(*ConsistencyNone)
+	c.Assert(ok, IsTrue)
+	s.assertLifetimeErrNil(ctrl, c)
+
+	conf.Consistency = "flush"
+	mock.ExpectExec("FLUSH TABLES WITH READ LOCK").WillReturnResult(resultOk)
+	mock.ExpectExec("UNLOCK TABLES").WillReturnResult(resultOk)
+	ctrl, _ = NewConsistencyController(conf, db)
+	_, ok = ctrl.(*ConsistencyFlushTableWithReadLock)
+	c.Assert(ok, IsTrue)
+	s.assertLifetimeErrNil(ctrl, c)
+	if err = mock.ExpectationsWereMet(); err != nil {
+		c.Fatalf(err.Error())
+	}
+
+	conf.Consistency = "snapshot"
+	conf.ServerInfo.ServerType = ServerTypeTiDB
+	conf.Snapshot = "" // let dumpling detect the TSO
+	rows := sqlmock.NewRows([]string{"File", "Position", "Binlog_Do_DB", "Binlog_Ignore_DB", "Executed_Gtid_Set"})
+	rows.AddRow("tidb-binlog", "413802961528946688", "", "", "")
+	mock.ExpectQuery("SHOW MASTER STATUS").WillReturnRows(rows)
+	mock.ExpectExec("SET SESSION tidb_snapshot").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+	ctrl, _ = NewConsistencyController(conf, db)
+	_, ok = ctrl.(*ConsistencySnapshot)
+	c.Assert(ok, IsTrue)
+	s.assertLifetimeErrNil(ctrl, c)
+	if err = mock.ExpectationsWereMet(); err != nil {
+		c.Fatalf(err.Error())
+	}
+
+	conf.Consistency = "lock"
+	conf.Tables = map[databaseName][]tableName{
+		"db1": {"t1", "t2", "t3"},
+		"db2": {"t4"},
+	}
+	for i := 0; i < 4; i += 1 {
+		mock.ExpectExec("LOCK TABLES").WillReturnResult(resultOk)
+	}
+	mock.ExpectExec("UNLOCK TABLES").WillReturnResult(resultOk)
+	ctrl, _ = NewConsistencyController(conf, db)
+	_, ok = ctrl.(*ConsistencyLockDumpingTables)
+	c.Assert(ok, IsTrue)
+	s.assertLifetimeErrNil(ctrl, c)
+	if err = mock.ExpectationsWereMet(); err != nil {
+		c.Fatalf(err.Error())
+	}
+}
+
+func (s *testConsistencySuite) TestResolveAutoConsistency(c *C) {
+	conf := DefaultConfig()
+	cases := []struct {
+		serverTp            ServerType
+		resolvedConsistency string
+	}{
+		{ServerTypeTiDB, "snapshot"},
+		{ServerTypeMySQL, "flush"},
+		{ServerTypeMariaDB, "flush"},
+		{ServerTypeUnknown, "none"},
+	}
+
+	for _, x := range cases {
+		conf.Consistency = "auto"
+		conf.ServerInfo.ServerType = x.serverTp
+		resolveAutoConsistency(conf)
+		cmt := Commentf("server type %s", x.serverTp.String())
+		c.Assert(conf.Consistency, Equals, x.resolvedConsistency, cmt)
+	}
+}
+
+func (s *testConsistencySuite) TestConsistencyControllerError(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+	conf := DefaultConfig()
+
+	conf.Consistency = "invalid_str"
+	_, err = NewConsistencyController(conf, db)
+	c.Assert(err, NotNil)
+	c.Assert(strings.Contains(err.Error(), "invalid consistency option"), IsTrue)
+
+	// snapshot consistency is only available in TiDB
+	conf.Consistency = "snapshot"
+	conf.ServerInfo.ServerType = ServerTypeUnknown
+	ctrl, _ := NewConsistencyController(conf, db)
+	err = ctrl.Setup()
+	c.Assert(err, NotNil)
+
+	// flush consistency is unavailable in TiDB
+	conf.Consistency = "flush"
+	conf.ServerInfo.ServerType = ServerTypeTiDB
+	ctrl, _ = NewConsistencyController(conf, db)
+	err = ctrl.Setup()
+	c.Assert(err, NotNil)
+
+	// lock table fail
+	conf.Consistency = "lock"
+	conf.Tables = map[databaseName][]tableName{"db": {"t"}}
+	mock.ExpectExec("LOCK TABLE").WillReturnError(errors.New(""))
+	ctrl, _ = NewConsistencyController(conf, db)
+	err = ctrl.Setup()
+	c.Assert(err, NotNil)
+}
diff --git a/v4/export/dump_test.go b/v4/export/dump_test.go
index 3306a4c..2d32e8c 100644
--- a/v4/export/dump_test.go
+++ b/v4/export/dump_test.go
@@ -3,6 +3,7 @@ package export
 import (
 	"context"
 	"fmt"
+
 	"github.com/DATA-DOG/go-sqlmock"
 	. "github.com/pingcap/check"
 )