Repository: pingcap/dumpling
Pull request #260 (2021-03-31 06:09:16)

dump: support new where sql generator
<!--
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-->
close https://github.com/pingcap/dumpling/issues/259

### What is changed and how it works?
support new where sql generator.

### 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)


Side effects

 - Increased code complexity

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
-->
- No release note


ti-chi-bot: [REVIEW NOTIFICATION]

This pull request has not been approved.


To complete the [pull request process](https://book.prow.tidb.io/#/en/workflows/pr), please ask the reviewers in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/260/owners) to review by filling `/cc @reviewer` in the comment.
After your PR has acquired the required number of LGTMs, you can assign this pull request to the committer in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/260/owners) by filling  `/assign @committer` in the comment to help you merge this pull request.

The full list of commands accepted by this bot can be found [here](https://prow.tidb.io/command-help?repo=pingcap%2Fdumpling).

<details>

Reviewer can indicate their review by writing `/lgtm` in a comment.
Reviewer can cancel approval by writing `/lgtm cancel` in a comment.
</details>

<!--Review Notification Identifier-->
lichunzhu: cc @tangenta 
lichunzhu: /hold
I will do a test on the v5.0 cluster and paste the result here later.
lichunzhu: /unhold
ti-chi-bot: [REVIEW NOTIFICATION]

This pull request has been approved by:

- leoppro




To complete the [pull request process](https://book.prow.tidb.io/#/en/workflows/pr), please ask the reviewers in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/260/owners) to review by filling `/cc @reviewer` in the comment.
After your PR has acquired the required number of LGTMs, you can assign this pull request to the committer in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/260/owners) by filling  `/assign @committer` in the comment to help you merge this pull request.

The full list of commands accepted by this bot can be found [here](https://prow.tidb.io/command-help?repo=pingcap%2Fdumpling).

<details>

Reviewer can indicate their review by writing `/lgtm` in a comment.
Reviewer can cancel approval by writing `/lgtm cancel` in a comment.
</details>

<!--Review Notification Identifier-->
lichunzhu: /merge
ti-chi-bot: This pull request has been accepted and is ready to merge. <details>Commit hash: d9ab10cba543a279e37551d384e798697e02e1ff</details>

Patch:
diff --git a/v4/export/dump.go b/v4/export/dump.go
index e0b5f55..81f7b19 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -521,7 +521,7 @@ func (d *Dumper) L() log.Logger {
 	return d.tctx.L()
 }
 
-func selectTiDBTableSample(conn *sql.Conn, dbName, tableName string) (pkFields []string, pkVals []string, err error) {
+func selectTiDBTableSample(conn *sql.Conn, dbName, tableName string) (pkFields []string, pkVals [][]string, err error) {
 	pkFields, pkColTypes, err := GetPrimaryKeyAndColumnTypes(conn, dbName, tableName)
 	if err != nil {
 		return nil, nil, errors.Trace(err)
@@ -538,18 +538,25 @@ func selectTiDBTableSample(conn *sql.Conn, dbName, tableName string) (pkFields [
 	if err != nil {
 		return nil, nil, errors.Trace(err)
 	}
-	iter := newRowIter(rows, len(pkFields))
+	pkValNum := len(pkFields)
+	iter := newRowIter(rows, pkValNum)
 	defer iter.Close()
 	rowRec := MakeRowReceiver(pkColTypes)
 	buf := new(bytes.Buffer)
+
 	for iter.HasNext() {
 		err = iter.Decode(rowRec)
 		if err != nil {
 			return nil, nil, errors.Trace(err)
 		}
+		pkValRow := make([]string, 0, pkValNum)
+		for _, rec := range rowRec.receivers {
+			rec.WriteToBuffer(buf, true)
+			pkValRow = append(pkValRow, buf.String())
+			buf.Reset()
+		}
 		rowRec.WriteToBuffer(buf, true)
-		pkVals = append(pkVals, buf.String())
-		buf.Reset()
+		pkVals = append(pkVals, pkValRow)
 		iter.Next()
 	}
 	return pkFields, pkVals, nil
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 7c93d1b..5374830 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -570,25 +570,149 @@ func buildSelectField(db *sql.Conn, dbName, tableName string, completeInsert boo
 	return "*", len(availableFields), nil
 }
 
-func buildWhereClauses(handleColNames, handleVals []string) []string {
-	if len(handleVals) == 0 {
+func buildWhereClauses(handleColNames []string, handleVals [][]string) []string {
+	if len(handleColNames) == 0 {
 		return nil
 	}
 	quotaCols := make([]string, len(handleColNames))
 	for i, s := range handleColNames {
 		quotaCols[i] = fmt.Sprintf("`%s`", escapeString(s))
 	}
-	handleCols := strings.Join(quotaCols, ",")
 	where := make([]string, 0, len(handleVals)+1)
-	where = append(where, fmt.Sprintf("(%s) < %s", handleCols, handleVals[0]))
+	buf := &bytes.Buffer{}
+	buildCompareClause(buf, quotaCols, handleVals[0], less, false)
+	where = append(where, buf.String())
+	buf.Reset()
 	for i := 1; i < len(handleVals); i++ {
 		low, up := handleVals[i-1], handleVals[i]
-		where = append(where, fmt.Sprintf("(%s) >= (%s) AND (%s) < (%s)", handleCols, low, handleCols, up))
+		buildBetweenClause(buf, quotaCols, low, up)
+		where = append(where, buf.String())
+		buf.Reset()
 	}
-	where = append(where, fmt.Sprintf("(%s) >= (%s)", handleCols, handleVals[len(handleVals)-1]))
+	buildCompareClause(buf, quotaCols, handleVals[len(handleVals)-1], greater, true)
+	where = append(where, buf.String())
+	buf.Reset()
 	return where
 }
 
+// return greater than TableRangeScan where clause
+// the result doesn't contain brackets
+const (
+	greater = '>'
+	less    = '<'
+	equal   = '='
+)
+
+// buildCompareClause build clause with specified bounds. Usually we will use the following two conditions:
+// (compare, writeEqual) == (less, false), return quotaCols < bound clause. In other words, (-inf, bound)
+// (compare, writeEqual) == (greater, true), return quotaCols >= bound clause. In other words, [bound, +inf)
+func buildCompareClause(buf *bytes.Buffer, quotaCols []string, bound []string, compare byte, writeEqual bool) { // revive:disable-line:flag-parameter
+	for i, col := range quotaCols {
+		if i > 0 {
+			buf.WriteString("or(")
+		}
+		for j := 0; j < i; j++ {
+			buf.WriteString(quotaCols[j])
+			buf.WriteByte(equal)
+			buf.WriteString(bound[j])
+			buf.WriteString(" and ")
+		}
+		buf.WriteString(col)
+		buf.WriteByte(compare)
+		if writeEqual && i == len(quotaCols)-1 {
+			buf.WriteByte(equal)
+		}
+		buf.WriteString(bound[i])
+		if i > 0 {
+			buf.WriteByte(')')
+		} else if i != len(quotaCols)-1 {
+			buf.WriteByte(' ')
+		}
+	}
+}
+
+// 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
+		}
+	}
+	return c, l
+}
+
+// buildBetweenClause build clause in a specified table range.
+// the result where clause will be low <= quotaCols < up. In other words, [low, up)
+func buildBetweenClause(buf *bytes.Buffer, quotaCols []string, low []string, up []string) {
+	singleBetween := func(writeEqual bool) {
+		buf.WriteString(quotaCols[0])
+		buf.WriteByte(greater)
+		if writeEqual {
+			buf.WriteByte(equal)
+		}
+		buf.WriteString(low[0])
+		buf.WriteString(" and ")
+		buf.WriteString(quotaCols[0])
+		buf.WriteByte(less)
+		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
+	}
+	if commonLen > 0 {
+		for i := 0; i < commonLen; i++ {
+			if i > 0 {
+				buf.WriteString(" and ")
+			}
+			buf.WriteString(quotaCols[i])
+			buf.WriteByte(equal)
+			buf.WriteString(low[i])
+		}
+		buf.WriteString(" and(")
+		defer buf.WriteByte(')')
+		quotaCols = quotaCols[commonLen:]
+		low = low[commonLen:]
+		up = up[commonLen:]
+	}
+
+	// handle special cases with only one column
+	if len(quotaCols) == 1 {
+		singleBetween(true)
+		return
+	}
+	buf.WriteByte('(')
+	singleBetween(false)
+	buf.WriteString(")or(")
+	buf.WriteString(quotaCols[0])
+	buf.WriteByte(equal)
+	buf.WriteString(low[0])
+	buf.WriteString(" and(")
+	buildCompareClause(buf, quotaCols[1:], low[1:], greater, true)
+	buf.WriteString("))or(")
+	buf.WriteString(quotaCols[0])
+	buf.WriteByte(equal)
+	buf.WriteString(up[0])
+	buf.WriteString(" and(")
+	buildCompareClause(buf, quotaCols[1:], up[1:], less, false)
+	buf.WriteString("))")
+}
+
 func buildOrderByClauseString(handleColNames []string) string {
 	if len(handleColNames) == 0 {
 		return ""
diff --git a/v4/export/sql_type.go b/v4/export/sql_type.go
index 04791f0..6e117f0 100644
--- a/v4/export/sql_type.go
+++ b/v4/export/sql_type.go
@@ -159,7 +159,7 @@ func SQLTypeNumberMaker() RowReceiverStringer {
 }
 
 // MakeRowReceiver constructs RowReceiverArr from column types
-func MakeRowReceiver(colTypes []string) RowReceiverStringer {
+func MakeRowReceiver(colTypes []string) RowReceiverArr {
 	rowReceiverArr := make([]RowReceiverStringer, len(colTypes))
 	for i, colTp := range colTypes {
 		recMaker, ok := colTypeRowReceiverMap[colTp]


Test patch:
diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go
index 61597f8..0402d42 100644
--- a/v4/export/sql_test.go
+++ b/v4/export/sql_test.go
@@ -347,6 +347,131 @@ func (s *testSQLSuite) TestShowCreateView(c *C) {
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 }
 
+func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
+	testCases := []struct {
+		handleColNames       []string
+		handleVals           [][]string
+		expectedWhereClauses []string
+	}{
+		{
+			[]string{},
+			[][]string{},
+			nil,
+		},
+		{
+			[]string{"a"},
+			[][]string{{"1"}},
+			[]string{"`a`<1", "`a`>=1"},
+		},
+		{
+			[]string{"a"},
+			[][]string{
+				{"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{"`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{
+				"`a`<1 or(`a`=1 and `b`<2)",
+				"(`a`>1 and `a`<3)or(`a`=1 and(`b`>=2))or(`a`=3 and(`b`<4))",
+				"(`a`>3 and `a`<5)or(`a`=3 and(`b`>=4))or(`a`=5 and(`b`<6))",
+				"`a`>5 or(`a`=5 and `b`>=6)",
+			},
+		},
+		{
+			[]string{"a", "b", "c"},
+			[][]string{
+				{"1", "2", "3"},
+				{"4", "5", "6"},
+			},
+			[]string{
+				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
+				"(`a`>1 and `a`<4)or(`a`=1 and(`b`>2 or(`b`=2 and `c`>=3)))or(`a`=4 and(`b`<5 or(`b`=5 and `c`<6)))",
+				"`a`>4 or(`a`=4 and `b`>5)or(`a`=4 and `b`=5 and `c`>=6)",
+			},
+		},
+		{
+			[]string{"a", "b", "c"},
+			[][]string{
+				{"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)",
+			},
+		},
+		{
+			[]string{"a", "b", "c"},
+			[][]string{
+				{"1", "2", "3"},
+				{"1", "2", "8"},
+			},
+			[]string{
+				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
+				"`a`=1 and `b`=2 and(`c`>=3 and `c`<8)",
+				"`a`>1 or(`a`=1 and `b`>2)or(`a`=1 and `b`=2 and `c`>=8)",
+			},
+		},
+		// special case: avoid return same samples
+		{
+			[]string{"a", "b", "c"},
+			[][]string{
+				{"1", "2", "3"},
+				{"1", "2", "3"},
+			},
+			[]string{
+				"`a`<1 or(`a`=1 and `b`<2)or(`a`=1 and `b`=2 and `c`<3)",
+				"false",
+				"`a`>1 or(`a`=1 and `b`>2)or(`a`=1 and `b`=2 and `c`>=3)",
+			},
+		},
+		// test string fields
+		{
+			[]string{"a", "b", "c"},
+			[][]string{
+				{"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\")",
+			},
+		},
+		{
+			[]string{"a", "b", "c", "d"},
+			[][]string{
+				{"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)",
+				"(`a`>1 and `a`<5)or(`a`=1 and(`b`>2 or(`b`=2 and `c`>3)or(`b`=2 and `c`=3 and `d`>=4)))or(`a`=5 and(`b`<6 or(`b`=6 and `c`<7)or(`b`=6 and `c`=7 and `d`<8)))",
+				"`a`>5 or(`a`=5 and `b`>6)or(`a`=5 and `b`=6 and `c`>7)or(`a`=5 and `b`=6 and `c`=7 and `d`>=8)",
+			},
+		},
+	}
+	for _, testCase := range testCases {
+		whereClauses := buildWhereClauses(testCase.handleColNames, testCase.handleVals)
+		c.Assert(whereClauses, DeepEquals, testCase.expectedWhereClauses)
+	}
+}
+
 func makeVersion(major, minor, patch int64, preRelease string) *semver.Version {
 	return &semver.Version{
 		Major:      major,