Repository: pingcap/dumpling
Pull request #196 (2020-11-10 02:35:37)

refine order by primary key
<!--
Thank you for contributing to Dumpling! Please read the [CONTRIBUTING](https://github.com/pingcap/dumpling/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.

If you want to open the **Challenge Program** pull request, please use the following template:
https://raw.githubusercontent.com/pingcap/.github/master/.github/PULL_REQUEST_TEMPLATE/challenge-program.md

You can use it with query parameters: https://github.com/pingcap/dumpling/compare/master...${you branch}?template=challenge-program.md
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->

https://github.com/pingcap/dumpling/issues/187

### What is changed and how it works?


### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Side effects

 - Possible performance regression
 - Increased code complexity
 - Breaking backward compatibility

Related changes

 - Need to cherry-pick to the release branch
 - Need to update the documentation
 
### Release note

<!-- bugfixes or new feature need a release note, must in the form of a list, such as

- support -T/--tables-list argument

or if no need to be included in the release note, just add the following line

- No release note
-->


jiyfhust: @lichunzhu @kennytm 
jiyfhust: hi, @lichunzhu Please have a check again.
lichunzhu: thanks for the contribution!

Patch:
diff --git a/go.sum b/go.sum
index 6f0d547..0d7ffc2 100644
--- a/go.sum
+++ b/go.sum
@@ -414,6 +414,7 @@ github.com/pingcap/br v0.0.0-20201027124415-c2ed897feada/go.mod h1:+zXr3GgwnjLWg
 github.com/pingcap/check v0.0.0-20190102082844-67f458068fc8/go.mod h1:B1+S9LNcuMyLH/4HMTViQOJevkGiik3wW2AN9zb2fNQ=
 github.com/pingcap/check v0.0.0-20191107115940-caf2b9e6ccf4/go.mod h1:PYMCGwN0JHjoqGr3HrZoD+b8Tgx8bKnArhSq8YVzUMc=
 github.com/pingcap/check v0.0.0-20191216031241-8a5a85928f12/go.mod h1:PYMCGwN0JHjoqGr3HrZoD+b8Tgx8bKnArhSq8YVzUMc=
+github.com/pingcap/check v0.0.0-20200212061837-5e12011dc712 h1:R8gStypOBmpnHEx1qi//SaqxJVI4inOqljg/Aj5/390=
 github.com/pingcap/check v0.0.0-20200212061837-5e12011dc712/go.mod h1:PYMCGwN0JHjoqGr3HrZoD+b8Tgx8bKnArhSq8YVzUMc=
 github.com/pingcap/errcode v0.0.0-20180921232412-a1a7271709d9/go.mod h1:4b2X8xSqxIroj/IZ9MX/VGZhAwc11wB9wRIzHvz6SeM=
 github.com/pingcap/errors v0.11.0/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTwv5KuLY8=
@@ -546,6 +547,7 @@ github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sasha-s/go-deadlock v0.2.0/go.mod h1:StQn567HiB1fF2yJ44N9au7wOhrPS3iZqiDbRupzT10=
 github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
+github.com/sergi/go-diff v1.0.1-0.20180205163309-da645544ed44 h1:tB9NOR21++IjLyVx3/PCPhWMwqGNCMQEH96A6dMZ/gc=
 github.com/sergi/go-diff v1.0.1-0.20180205163309-da645544ed44/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=
 github.com/shirou/gopsutil v2.19.10+incompatible h1:lA4Pi29JEVIQIgATSeftHSY0rMGI9CLrl2ZvDLiahto=
 github.com/shirou/gopsutil v2.19.10+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 58c9d1b..0eff0a3 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -276,17 +276,20 @@ func buildOrderByClause(conf *Config, db *sql.Conn, database, table string) (str
 		}
 		if ok {
 			return "ORDER BY _tidb_rowid", nil
-		} else {
-			return "", nil
 		}
 	}
-	pkName, err := GetPrimaryKeyName(db, database, table)
+	cols, err := GetPrimaryKeyColumns(db, database, table)
 	if err != nil {
 		return "", errors.Trace(err)
 	}
-	tableContainsPriKey := pkName != ""
+	tableContainsPriKey := len(cols) != 0
 	if tableContainsPriKey {
-		return fmt.Sprintf("ORDER BY `%s`", escapeString(pkName)), nil
+		separator := ", "
+		quotaCols := make([]string, len(cols))
+		for i, col := range cols {
+			quotaCols[i] = fmt.Sprintf("`%s`", escapeString(col))
+		}
+		return fmt.Sprintf("ORDER BY %s", strings.Join(quotaCols, separator)), nil
 	}
 	return "", nil
 }
@@ -315,6 +318,29 @@ func GetColumnTypes(db *sql.Conn, fields, database, table string) ([]*sql.Column
 	return rows.ColumnTypes()
 }
 
+func GetPrimaryKeyColumns(db *sql.Conn, database, table string) ([]string, error) {
+	priKeyColsQuery := "SELECT column_name FROM information_schema.KEY_COLUMN_USAGE " +
+		"WHERE table_schema = ? AND table_name = ? AND CONSTRAINT_NAME = 'PRIMARY' order by ORDINAL_POSITION;"
+	rows, err := db.QueryContext(context.Background(), priKeyColsQuery, database, table)
+	if err != nil {
+		return nil, errors.Annotatef(err, "sql: %s", priKeyColsQuery)
+	}
+	defer rows.Close()
+	var cols []string
+	var col string
+	for rows.Next() {
+		err = rows.Scan(&col)
+		if err != nil {
+			return nil, errors.Annotatef(err, "sql: %s", priKeyColsQuery)
+		}
+		cols = append(cols, col)
+	}
+	if err = rows.Err(); err != nil {
+		return nil, errors.Annotatef(err, "sql: %s", priKeyColsQuery)
+	}
+	return cols, nil
+}
+
 func GetPrimaryKeyName(db *sql.Conn, database, table string) (string, error) {
 	priKeyQuery := "SELECT column_name FROM information_schema.columns " +
 		"WHERE table_schema = ? AND table_name = ? AND column_key = 'PRI';"


Test patch:
diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go
index 9185b18..d95b315 100644
--- a/v4/export/sql_test.go
+++ b/v4/export/sql_test.go
@@ -84,6 +84,10 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 	mock.ExpectExec("SELECT _tidb_rowid from `test`.`t`").
 		WillReturnError(errors.New(`1054, "Unknown column '_tidb_rowid' in 'field list'"`))
 
+	mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
+		WithArgs("test", "t").
+		WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id"))
+
 	orderByClause, err = buildOrderByClause(mockConf, conn, "test", "t")
 	c.Assert(err, IsNil)
 
@@ -94,7 +98,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 	selectedField, err = buildSelectField(conn, "test", "t", false)
 	c.Assert(err, IsNil)
 	q = buildSelectQuery("test", "t", selectedField, "", orderByClause)
-	c.Assert(q, Equals, "SELECT * FROM `test`.`t`")
+	c.Assert(q, Equals, "SELECT * FROM `test`.`t` ORDER BY `id`")
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 
 	// Test other servers.
@@ -104,7 +108,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 	for _, serverTp := range otherServers {
 		mockConf.ServerInfo.ServerType = serverTp
 		cmt := Commentf("server type: %s", serverTp)
-		mock.ExpectQuery("SELECT column_name FROM information_schema.columns").
+		mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
 			WithArgs("test", "t").
 			WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id"))
 		orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
@@ -127,7 +131,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 	for _, serverTp := range otherServers {
 		mockConf.ServerInfo.ServerType = serverTp
 		cmt := Commentf("server type: %s", serverTp)
-		mock.ExpectQuery("SELECT column_name FROM information_schema.columns").
+		mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
 			WithArgs("test", "t").
 			WillReturnRows(sqlmock.NewRows([]string{"column_name"}))
 
@@ -165,6 +169,91 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 	}
 }
 
+func (s *testDumpSuite) TestBuildOrderByClause(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+	conn, err := db.Conn(context.Background())
+	c.Assert(err, IsNil)
+
+	mockConf := DefaultConfig()
+	mockConf.SortByPk = true
+
+	// Test TiDB server.
+	mockConf.ServerInfo.ServerType = ServerTypeTiDB
+
+	// _tidb_rowid is available.
+	mock.ExpectExec("SELECT _tidb_rowid from `test`.`t`").
+		WillReturnResult(sqlmock.NewResult(0, 0))
+
+	orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
+	c.Assert(err, IsNil)
+	c.Assert(orderByClause, Equals, "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'"`))
+
+	mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
+		WithArgs("test", "t").
+		WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id"))
+
+	orderByClause, err = buildOrderByClause(mockConf, conn, "test", "t")
+	c.Assert(err, IsNil)
+	c.Assert(orderByClause, Equals, "ORDER BY `id`")
+
+	// Test other servers.
+	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.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
+			WithArgs("test", "t").
+			WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id"))
+		orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(orderByClause, Equals, "ORDER BY `id`", cmt)
+	}
+
+	// Test table with joint primary key.
+	for _, serverTp := range otherServers {
+		mockConf.ServerInfo.ServerType = serverTp
+		cmt := Commentf("server type: %s", serverTp)
+		mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
+			WithArgs("test", "t").
+			WillReturnRows(sqlmock.NewRows([]string{"column_name"}).AddRow("id").AddRow("name"))
+		orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(orderByClause, Equals, "ORDER BY `id`, `name`", cmt)
+	}
+
+	// Test table without primary key.
+	for _, serverTp := range otherServers {
+		mockConf.ServerInfo.ServerType = serverTp
+		cmt := Commentf("server type: %s", serverTp)
+		mock.ExpectQuery("SELECT column_name FROM information_schema.KEY_COLUMN_USAGE").
+			WithArgs("test", "t").
+			WillReturnRows(sqlmock.NewRows([]string{"column_name"}))
+
+		orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(orderByClause, Equals, "", cmt)
+	}
+
+	// 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)
+
+		orderByClause, err := buildOrderByClause(mockConf, conn, "test", "t")
+		c.Assert(err, IsNil, cmt)
+		c.Assert(orderByClause, Equals, "", cmt)
+	}
+}
+
 func (s *testDumpSuite) TestBuildSelectField(c *C) {
 	db, mock, err := sqlmock.New()
 	c.Assert(err, IsNil)