Skip to content

Commit

Permalink
refine order by primary key (pingcap#196)
Browse files Browse the repository at this point in the history
* refine order by primary key
  • Loading branch information
jiyfhust authored Nov 10, 2020
1 parent f269d59 commit e1ed4ff
Show file tree
Hide file tree
Showing 3 changed files with 125 additions and 8 deletions.
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down
36 changes: 31 additions & 5 deletions v4/export/sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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';"
Expand Down
95 changes: 92 additions & 3 deletions v4/export/sql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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.
Expand All @@ -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")
Expand All @@ -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"}))

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit e1ed4ff

Please sign in to comment.