Repository: pingcap/dumpling
Pull request #262 (2021-04-01 04:21:14)

sql: fix table sample sql generator bug
<!--
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/pull/260 introduced some bugs.

### What is changed and how it works?
Fix them.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test

### 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
-->
- No release note


ti-chi-bot: @lichunzhu: PR needs rebase.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
lichunzhu: > can you add a test case showing what is fixed by this?

Already addressed, PTAL again.

Patch:
diff --git a/go.mod b/go.mod
index 9655a7d..d600d9e 100644
--- a/go.mod
+++ b/go.mod
@@ -22,6 +22,7 @@ require (
 	github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible
 	github.com/prometheus/client_golang v1.5.1
 	github.com/prometheus/client_model v0.2.0
+	github.com/siddontang/go-mysql v0.0.0-20200222075837-12e89848f047
 	github.com/soheilhy/cmux v0.1.4
 	github.com/spf13/pflag v1.0.5
 	github.com/syndtr/goleveldb v1.0.1-0.20190625010220-02440ea7a285 // indirect
diff --git a/go.sum b/go.sum
index 3538c38..9da22f5 100644
--- a/go.sum
+++ b/go.sum
@@ -766,6 +766,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/satori/go.uuid v1.2.1-0.20181028125025-b2ce2384e17b h1:gQZ0qzfKHQIybLANtM3mBXNUtOfsCFXeTsnBqCsx1KM=
 github.com/satori/go.uuid v1.2.1-0.20181028125025-b2ce2384e17b/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
 github.com/segmentio/kafka-go v0.1.0/go.mod h1:X6itGqS9L4jDletMsxZ7Dz+JFWxM6JHfPOCvTvk+EJo=
 github.com/segmentio/kafka-go v0.2.0/go.mod h1:X6itGqS9L4jDletMsxZ7Dz+JFWxM6JHfPOCvTvk+EJo=
@@ -789,9 +790,12 @@ github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeV
 github.com/shurcooL/vfsgen v0.0.0-20181020040650-a97a25d856ca/go.mod h1:TrYk7fJVaAttu97ZZKrO9UbRa8izdowaMIZcxYMbVaw=
 github.com/shurcooL/vfsgen v0.0.0-20181202132449-6a9ea43bcacd h1:ug7PpSOB5RBPK1Kg6qskGBoP3Vnj/aNYFTznWvlkGo0=
 github.com/shurcooL/vfsgen v0.0.0-20181202132449-6a9ea43bcacd/go.mod h1:TrYk7fJVaAttu97ZZKrO9UbRa8izdowaMIZcxYMbVaw=
+github.com/siddontang/go v0.0.0-20180604090527-bdc77568d726 h1:xT+JlYxNGqyT+XcU8iUrN18JYed2TvG9yN5ULG2jATM=
 github.com/siddontang/go v0.0.0-20180604090527-bdc77568d726/go.mod h1:3yhqj7WBBfRhbBlzyOC3gUxftwsU0u8gqevxwIHQpMw=
 github.com/siddontang/go-log v0.0.0-20180807004314-8d05993dda07/go.mod h1:yFdBgwXP24JziuRl2NMUahT7nGLNOKi1SIiFxMttVD4=
+github.com/siddontang/go-log v0.0.0-20190221022429-1e957dd83bed h1:KMgQoLJGCq1IoZpLZE3AIffh9veYWoVlsvA4ib55TMM=
 github.com/siddontang/go-log v0.0.0-20190221022429-1e957dd83bed/go.mod h1:yFdBgwXP24JziuRl2NMUahT7nGLNOKi1SIiFxMttVD4=
+github.com/siddontang/go-mysql v0.0.0-20200222075837-12e89848f047 h1:boyJ8EgQN/aC3grvx8QUoJrptt7RvneezSJSCbW25a4=
 github.com/siddontang/go-mysql v0.0.0-20200222075837-12e89848f047/go.mod h1:+W4RCzesQDI11HvIkaDjS8yM36SpAnGNQ7jmTLn5BnU=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/sirupsen/logrus v1.3.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 6eb5670..1e8f613 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -616,6 +616,9 @@ func selectTiDBTableSample(conn *sql.Conn, dbName, tableName string) (pkFields [
 	if hasImplicitRowID {
 		pkFields, pkColTypes = []string{"_tidb_rowid"}, []string{"BIGINT"}
 	}
+	if len(pkFields) == 0 {
+		return pkFields, pkVals, nil
+	}
 	query := buildTiDBTableSampleQuery(pkFields, dbName, tableName)
 	rows, err := conn.QueryContext(context.Background(), query)
 	if err != nil {
@@ -638,7 +641,6 @@ func selectTiDBTableSample(conn *sql.Conn, dbName, tableName string) (pkFields [
 			pkValRow = append(pkValRow, buf.String())
 			buf.Reset()
 		}
-		rowRec.WriteToBuffer(buf, true)
 		pkVals = append(pkVals, pkValRow)
 		iter.Next()
 	}
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 78c8652..2e160d8 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -514,8 +514,20 @@ func isUnknownSystemVariableErr(err error) bool {
 func resetDBWithSessionParams(tctx *tcontext.Context, db *sql.DB, dsn string, params map[string]interface{}) (*sql.DB, error) {
 	support := make(map[string]interface{})
 	for k, v := range params {
+		var pv interface{}
+		if str, ok := v.(string); ok {
+			if pvi, err := strconv.ParseInt(str, 10, 64); err == nil {
+				pv = pvi
+			} else if pvf, err := strconv.ParseFloat(str, 64); err == nil {
+				pv = pvf
+			} else {
+				pv = str
+			}
+		} else {
+			pv = v
+		}
 		s := fmt.Sprintf("SET SESSION %s = ?", k)
-		_, err := db.ExecContext(tctx, s, v)
+		_, err := db.ExecContext(tctx, s, pv)
 		if err != nil {
 			if isUnknownSystemVariableErr(err) {
 				tctx.L().Info("session variable is not supported by db", zap.String("variable", k), zap.Reflect("value", v))
@@ -524,7 +536,7 @@ func resetDBWithSessionParams(tctx *tcontext.Context, db *sql.DB, dsn string, pa
 			return nil, errors.Trace(err)
 		}
 
-		support[k] = v
+		support[k] = pv
 	}
 
 	for k, v := range support {
@@ -663,26 +675,14 @@ func buildCompareClause(buf *bytes.Buffer, quotaCols []string, bound []string, c
 	}
 }
 
-// Compare returns an integer comparing two string arrays lexicographically.
-// return (compare result, common prefix length)
-// return -1 if low < up
-// return  0 if low == up
-// return  1 if low > up
-func compareArrString(low []string, up []string) (compareRes, commonLength int) {
-	l := len(low)
-	c := -1
-	if u := len(up); u < l {
-		l = u
-		c = 1
-	} else if u == l {
-		c = 0
-	}
-	for i := 0; i < l; i++ {
-		if d := strings.Compare(low[i], up[i]); d != 0 {
-			return d, i
+// getCommonLength returns the common length of low and up
+func getCommonLength(low []string, up []string) int {
+	for i := range low {
+		if low[i] != up[i] {
+			return i
 		}
 	}
-	return c, l
+	return len(low)
 }
 
 // buildBetweenClause build clause in a specified table range.
@@ -701,13 +701,13 @@ func buildBetweenClause(buf *bytes.Buffer, quotaCols []string, low []string, up
 		buf.WriteString(up[0])
 	}
 	// handle special cases with common prefix
-	compare, commonLen := compareArrString(low, up)
-	// unexpected case for low >= up, return empty result
-	if compare >= 0 {
-		buf.WriteString("false")
-		return
-	}
+	commonLen := getCommonLength(low, up)
 	if commonLen > 0 {
+		// unexpected case for low == up, return empty result
+		if commonLen == len(low) {
+			buf.WriteString("false")
+			return
+		}
 		for i := 0; i < commonLen; i++ {
 			if i > 0 {
 				buf.WriteString(" and ")


Test patch:
diff --git a/v4/export/dump_test.go b/v4/export/dump_test.go
index b9d8ec6..64cccc2 100644
--- a/v4/export/dump_test.go
+++ b/v4/export/dump_test.go
@@ -25,7 +25,7 @@ func (s *testSQLSuite) TestDumpBlock(c *C) {
 		WillReturnRows(sqlmock.NewRows([]string{"Database", "Create Database"}).
 			AddRow("test", "CREATE DATABASE `test` /*!40100 DEFAULT CHARACTER SET utf8mb4 */"))
 
-	tctx, cancel := tcontext.Background().WithCancel()
+	tctx, cancel := tcontext.Background().WithLogger(appLogger).WithCancel()
 	defer cancel()
 	conn, err := db.Conn(tctx)
 	c.Assert(err, IsNil)
diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go
index bdecade..e346fa1 100644
--- a/v4/export/sql_test.go
+++ b/v4/export/sql_test.go
@@ -5,13 +5,17 @@ package export
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
+	"fmt"
+	"strings"
 
 	tcontext "github.com/pingcap/dumpling/v4/context"
 
 	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/coreos/go-semver/semver"
 	. "github.com/pingcap/check"
+	"github.com/siddontang/go-mysql/mysql"
 )
 
 var _ = Suite(&testSQLSuite{})
@@ -406,41 +410,66 @@ func (s *testSQLSuite) TestGetSuitableRows(c *C) {
 }
 
 func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+	conn, err := db.Conn(context.Background())
+	c.Assert(err, IsNil)
+	tctx, cancel := tcontext.Background().WithLogger(appLogger).WithCancel()
+
+	d := &Dumper{
+		tctx:      tctx,
+		conf:      DefaultConfig(),
+		cancelCtx: cancel,
+	}
+	d.conf.ServerInfo = ServerInfo{
+		ServerType:    ServerTypeTiDB,
+		ServerVersion: tableSampleVersion,
+	}
+	database := "foo"
+	table := "bar"
+
 	testCases := []struct {
 		handleColNames       []string
-		handleVals           [][]string
+		handleColTypes       []string
+		handleVals           [][]driver.Value
 		expectedWhereClauses []string
 	}{
 		{
 			[]string{},
-			[][]string{},
+			[]string{},
+			[][]driver.Value{},
 			nil,
 		},
 		{
 			[]string{"a"},
-			[][]string{{"1"}},
+			[]string{"bigint"},
+			[][]driver.Value{{1}},
 			[]string{"`a`<1", "`a`>=1"},
 		},
 		{
 			[]string{"a"},
-			[][]string{
-				{"1"},
-				{"2"},
-				{"3"},
+			[]string{"bigint"},
+			[][]driver.Value{
+				{1},
+				{2},
+				{3},
 			},
 			[]string{"`a`<1", "`a`>=1 and `a`<2", "`a`>=2 and `a`<3", "`a`>=3"},
 		},
 		{
 			[]string{"a", "b"},
-			[][]string{{"1", "2"}},
+			[]string{"bigint", "bigint"},
+			[][]driver.Value{{1, 2}},
 			[]string{"`a`<1 or(`a`=1 and `b`<2)", "`a`>1 or(`a`=1 and `b`>=2)"},
 		},
 		{
 			[]string{"a", "b"},
-			[][]string{
-				{"1", "2"},
-				{"3", "4"},
-				{"5", "6"},
+			[]string{"bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2},
+				{3, 4},
+				{5, 6},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)",
@@ -451,9 +480,10 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 		},
 		{
 			[]string{"a", "b", "c"},
-			[][]string{
-				{"1", "2", "3"},
-				{"4", "5", "6"},
+			[]string{"bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2, 3},
+				{4, 5, 6},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
@@ -463,9 +493,10 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 		},
 		{
 			[]string{"a", "b", "c"},
-			[][]string{
-				{"1", "2", "3"},
-				{"1", "4", "5"},
+			[]string{"bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2, 3},
+				{1, 4, 5},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
@@ -475,9 +506,10 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 		},
 		{
 			[]string{"a", "b", "c"},
-			[][]string{
-				{"1", "2", "3"},
-				{"1", "2", "8"},
+			[]string{"bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2, 3},
+				{1, 2, 8},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
@@ -488,9 +520,10 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 		// special case: avoid return same samples
 		{
 			[]string{"a", "b", "c"},
-			[][]string{
-				{"1", "2", "3"},
-				{"1", "2", "3"},
+			[]string{"bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2, 3},
+				{1, 2, 3},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
@@ -498,24 +531,40 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 				"`a`>1 or(`a`=1 and `b`>2)or(`a`=1 and `b`=2 and `c`>=3)",
 			},
 		},
+		// special case: numbers has bigger lexicographically order but lower number
+		{
+			[]string{"a", "b", "c"},
+			[]string{"bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{12, 2, 3},
+				{111, 4, 5},
+			},
+			[]string{
+				"`a`<12 or(`a`=12 and `b`<2)or(`a`=12 and `b`=2 and `c`<3)",
+				"(`a`>12 and `a`<111)or(`a`=12 and(`b`>2 or(`b`=2 and `c`>=3)))or(`a`=111 and(`b`<4 or(`b`=4 and `c`<5)))", // should return sql correctly
+				"`a`>111 or(`a`=111 and `b`>4)or(`a`=111 and `b`=4 and `c`>=5)",
+			},
+		},
 		// test string fields
 		{
 			[]string{"a", "b", "c"},
-			[][]string{
-				{"1", "2", "\"3\""},
-				{"1", "4", "\"5\""},
+			[]string{"bigint", "bigint", "varchar"},
+			[][]driver.Value{
+				{1, 2, "3"},
+				{1, 4, "5"},
 			},
 			[]string{
-				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<\"3\")",
-				"`a`=1 and((`b`>2 and `b`<4)or(`b`=2 and(`c`>=\"3\"))or(`b`=4 and(`c`<\"5\")))",
-				"`a`>1 or(`a`=1 and `b`>4)or(`a`=1 and `b`=4 and `c`>=\"5\")",
+				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<'3')",
+				"`a`=1 and((`b`>2 and `b`<4)or(`b`=2 and(`c`>='3'))or(`b`=4 and(`c`<'5')))",
+				"`a`>1 or(`a`=1 and `b`>4)or(`a`=1 and `b`=4 and `c`>='5')",
 			},
 		},
 		{
 			[]string{"a", "b", "c", "d"},
-			[][]string{
-				{"1", "2", "3", "4"},
-				{"5", "6", "7", "8"},
+			[]string{"bigint", "bigint", "bigint", "bigint"},
+			[][]driver.Value{
+				{1, 2, 3, 4},
+				{5, 6, 7, 8},
 			},
 			[]string{
 				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)or(`a`=1 and `b`=2 and `c`=3 and `d`<4)",
@@ -524,9 +573,96 @@ func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 			},
 		},
 	}
-	for _, testCase := range testCases {
-		whereClauses := buildWhereClauses(testCase.handleColNames, testCase.handleVals)
+	transferHandleValStrings := func(handleColTypes []string, handleVals [][]driver.Value) [][]string {
+		handleValStrings := make([][]string, 0, len(handleVals))
+		for _, handleVal := range handleVals {
+			handleValString := make([]string, 0, len(handleVal))
+			for i, val := range handleVal {
+				rec := colTypeRowReceiverMap[strings.ToUpper(handleColTypes[i])]()
+				var valStr string
+				switch rec.(type) {
+				case *SQLTypeString:
+					valStr = fmt.Sprintf("'%s'", val)
+				case *SQLTypeBytes:
+					valStr = fmt.Sprintf("x'%x'", val)
+				case *SQLTypeNumber:
+					valStr = fmt.Sprintf("%d", val)
+				}
+				handleValString = append(handleValString, valStr)
+			}
+			handleValStrings = append(handleValStrings, handleValString)
+		}
+		return handleValStrings
+	}
+
+	for i, testCase := range testCases {
+		c.Log(fmt.Sprintf("case #%d", i))
+		handleColNames := testCase.handleColNames
+		handleColTypes := testCase.handleColTypes
+		handleVals := testCase.handleVals
+		handleValStrings := transferHandleValStrings(handleColTypes, handleVals)
+
+		// Test build whereClauses
+		whereClauses := buildWhereClauses(handleColNames, handleValStrings)
 		c.Assert(whereClauses, DeepEquals, testCase.expectedWhereClauses)
+
+		// Test build tasks through table sample
+		if len(handleColNames) > 0 {
+			taskChan := make(chan Task, 128)
+			quotaCols := make([]string, 0, len(handleColNames))
+			for _, col := range quotaCols {
+				quotaCols = append(quotaCols, wrapBackTicks(col))
+			}
+			selectFields := strings.Join(quotaCols, ",")
+			meta := &tableMeta{
+				database:      database,
+				table:         table,
+				selectedField: selectFields,
+				specCmts: []string{
+					"/*!40101 SET NAMES binary*/;",
+				},
+			}
+
+			rows := sqlmock.NewRows([]string{"COLUMN_NAME", "DATA_TYPE"})
+			for i := range handleColNames {
+				rows.AddRow(handleColNames[i], handleColTypes[i])
+			}
+			mock.ExpectQuery("SELECT c.COLUMN_NAME, DATA_TYPE FROM").WithArgs(database, table).WillReturnRows(rows)
+			mock.ExpectExec(fmt.Sprintf("SELECT _tidb_rowid from `%s`.`%s` LIMIT 0", database, table)).
+				WillReturnError(&mysql.MyError{
+					Code:    mysql.ER_BAD_FIELD_ERROR,
+					State:   "42S22",
+					Message: "Unknown column '_tidb_rowid' in 'field list'",
+				})
+
+			rows = sqlmock.NewRows(handleColNames)
+			for _, handleVal := range handleVals {
+				rows.AddRow(handleVal...)
+			}
+			mock.ExpectQuery(fmt.Sprintf("SELECT .* FROM `%s`.`%s` TABLESAMPLE REGIONS", database, table)).WillReturnRows(rows)
+
+			rows = sqlmock.NewRows([]string{"COLUMN_NAME", "EXTRA"})
+			for _, handleCol := range handleColNames {
+				rows.AddRow(handleCol, "")
+			}
+			mock.ExpectQuery("SELECT COLUMN_NAME,EXTRA FROM INFORMATION_SCHEMA.COLUMNS").WithArgs(database, table).
+				WillReturnRows(rows)
+
+			c.Assert(d.concurrentDumpTable(tctx, conn, meta, taskChan), IsNil)
+			c.Assert(mock.ExpectationsWereMet(), IsNil)
+			orderByClause := buildOrderByClauseString(handleColNames)
+
+			for i, w := range testCase.expectedWhereClauses {
+				query := buildSelectQuery(database, table, "*", buildWhereCondition(d.conf, w), orderByClause)
+				task := <-taskChan
+				taskTableData, ok := task.(*TaskTableData)
+				c.Assert(ok, IsTrue)
+				c.Assert(taskTableData.ChunkIndex, Equals, i)
+				data, ok := taskTableData.Data.(*tableData)
+				c.Assert(ok, IsTrue)
+				c.Assert(data.query, Equals, query)
+			}
+		}
 	}
 }
 
diff --git a/v4/export/writer_util_test.go b/v4/export/writer_util_test.go
index ae6453e..cac38c9 100644
--- a/v4/export/writer_util_test.go
+++ b/v4/export/writer_util_test.go
@@ -9,14 +9,27 @@ import (
 	"testing"
 
 	tcontext "github.com/pingcap/dumpling/v4/context"
+	"github.com/pingcap/dumpling/v4/log"
 
 	"github.com/pingcap/br/pkg/storage"
 	. "github.com/pingcap/check"
 	"github.com/pingcap/errors"
 )
 
+var appLogger log.Logger
+
 func TestT(t *testing.T) {
 	initColTypeRowReceiverMap()
+	logger, err := log.InitAppLogger(&log.Config{
+		Level:  "info",
+		File:   "",
+		Format: "text",
+	})
+	if err != nil {
+		t.Log("fail to init logger, err: " + err.Error())
+		t.Fail()
+	}
+	appLogger = logger
 	TestingT(t)
 }