diff --git a/go.sum b/go.sum index 6f0d54720ea0e..0d7ffc264eb85 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 58c9d1b67ff07..0eff0a37032a6 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';" diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go index 9185b186f574d..d95b3150ba2a1 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)