diff --git a/go.mod b/go.mod index 1b525a0..8551ab8 100644 --- a/go.mod +++ b/go.mod @@ -6,10 +6,13 @@ require ( github.com/go-sql-driver/mysql v1.7.1 github.com/pingcap/errors v0.11.5-0.20221009092201-b66cddb77c32 github.com/sirupsen/logrus v1.8.1 + github.com/stretchr/testify v1.8.4 ) require ( - github.com/stretchr/testify v1.8.4 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect go.uber.org/atomic v1.11.0 // indirect golang.org/x/sys v0.5.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 78207a7..85ff952 100644 --- a/go.sum +++ b/go.sum @@ -31,6 +31,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/r/example.result b/r/example.result index a3f5181..64da5aa 100644 --- a/r/example.result +++ b/r/example.result @@ -3,3 +3,11 @@ insert into t values(1, 1), (2, 2), (3, 3), (4, 4), (5, 5); select * from t where a = 1; a b 1 1 +explain analyze format='brief' select * from t; +id estRows actRows task access object execution info operator info memory disk +TableReader 10000.00 5 root NULL time:, loops:, RU:, cop_task: {num:, max:, proc_keys:, rpc_num:, rpc_time:, copr_cache_hit_ratio:, build_task_duration:, max_distsql_concurrency:} data:TableFullScan Bytes N/A +└─TableFullScan 10000.00 5 cop[tikv] table:t tikv_task:{time:, loops:} keep order:false, stats:pseudo N/A N/A +explain analyze select * from t; +id estRows actRows task access object execution info operator info memory disk +TableReader_5 10000.00 5 root NULL time:, loops:, RU:, cop_task: {num:, max:, proc_keys:, rpc_num:, rpc_time:, copr_cache_hit_ratio:, build_task_duration:, max_distsql_concurrency:} data:TableFullScan_4 Bytes N/A +└─TableFullScan_4 10000.00 5 cop[tikv] table:t tikv_task:{time:, loops:} keep order:false, stats:pseudo N/A N/A diff --git a/src/main.go b/src/main.go index 592365f..692ddcf 100644 --- a/src/main.go +++ b/src/main.go @@ -21,6 +21,7 @@ import ( "fmt" "os" "path/filepath" + "regexp" "sort" "strconv" "strings" @@ -90,6 +91,11 @@ type ReplaceColumn struct { replace []byte } +type ReplaceRegex struct { + regex *regexp.Regexp + replace string +} + type tester struct { mdb *sql.DB name string @@ -131,6 +137,9 @@ type tester struct { // replace output column through --replace_column 1 3 # replaceColumn []ReplaceColumn + + // replace output result through --replace_regex /\.dll/.so/ + replaceRegex []*ReplaceRegex } func newTester(name string) *tester { @@ -378,6 +387,7 @@ func (t *tester) Run() error { t.sortedResult = false t.replaceColumn = nil + t.replaceRegex = nil case Q_SORTED_RESULT: t.sortedResult = true case Q_REPLACE_COLUMN: @@ -426,6 +436,13 @@ func (t *tester) Run() error { if err != nil { return errors.Annotate(err, "failed to remove file") } + case Q_REPLACE_REGEX: + t.replaceRegex = nil + regex, err := ParseReplaceRegex(q.Query) + if err != nil { + return errors.Annotate(err, fmt.Sprintf("Could not parse regex in --replace_regex: sql:%v", q.Query)) + } + t.replaceRegex = regex default: log.WithFields(log.Fields{"command": q.firstWord, "arguments": q.Query, "line": q.Line}).Warn("command not implemented") } @@ -604,8 +621,13 @@ func (t *tester) execute(query query) error { // TODO: check whether this err is expected. // but now we think it is. + errStr := err.Error() + for _, reg := range t.replaceRegex { + errStr = reg.regex.ReplaceAllString(errStr, reg.replace) + } + // output expected err - fmt.Fprintf(&t.buf, "%s\n", strings.ReplaceAll(err.Error(), "\r", "")) + fmt.Fprintf(&t.buf, "%s\n", strings.ReplaceAll(errStr, "\r", "")) err = nil } // clear expected errors after we execute the first query @@ -633,6 +655,21 @@ func (t *tester) execute(query query) error { } func (t *tester) writeQueryResult(rows *byteRows) error { + if t.sortedResult { + sort.Sort(rows) + } + + if len(t.replaceColumn) > 0 { + for _, row := range rows.data { + for _, r := range t.replaceColumn { + if len(row.data) < r.col { + continue + } + row.data[r.col-1] = r.replace + } + } + } + cols := rows.cols for i, c := range cols { t.buf.WriteString(c) @@ -645,6 +682,11 @@ func (t *tester) writeQueryResult(rows *byteRows) error { for _, row := range rows.data { var value string for i, col := range row.data { + // replace result by regex + for _, reg := range t.replaceRegex { + col = reg.regex.ReplaceAll(col, []byte(reg.replace)) + } + // Here we can check if the value is nil (NULL value) if col == nil { value = "NULL" @@ -751,21 +793,7 @@ func (t *tester) executeStmt(query string) error { return nil } - if len(t.replaceColumn) > 0 { - for _, row := range rows.data { - for _, r := range t.replaceColumn { - if len(row.data) < r.col { - continue - } - row.data[r.col-1] = r.replace - } - } - } - if len(rows.cols) > 0 || len(rows.data) > 0 { - if t.sortedResult { - sort.Sort(rows) - } if err = t.writeQueryResult(rows); err != nil { return errors.Trace(err) } diff --git a/src/util.go b/src/util.go index 12a6689..b62c64b 100644 --- a/src/util.go +++ b/src/util.go @@ -15,6 +15,8 @@ package main import ( "database/sql" + "regexp" + "strings" "time" "github.com/pingcap/errors" @@ -49,3 +51,56 @@ func OpenDBWithRetry(driverName, dataSourceName string, retryCount int) (mdb *sq return } + +func processEscapes(str string) string { + escapeMap := map[string]string{ + `\n`: "\n", + `\t`: "\t", + `\r`: "\r", + `\/`: "/", + `\\`: "\\", // better be the last one + } + + for escape, replacement := range escapeMap { + str = strings.ReplaceAll(str, escape, replacement) + } + + return str +} + +func ParseReplaceRegex(originalString string) ([]*ReplaceRegex, error) { + var begin, middle, end, cnt int + ret := make([]*ReplaceRegex, 0) + for i, c := range originalString { + if c != '/' { + continue + } + if i != 0 && originalString[i-1] == '\\' { + continue + } + cnt++ + switch cnt % 3 { + // The first '/' + case 1: + begin = i + // The second '/' + case 2: + middle = i + // The last '/', we could compile regex and process replace string + case 0: + end = i + reg, err := regexp.Compile(originalString[begin+1 : middle]) + if err != nil { + return nil, err + } + ret = append(ret, &ReplaceRegex{ + regex: reg, + replace: processEscapes(originalString[middle+1 : end]), + }) + } + } + if cnt%3 != 0 { + return nil, errors.Errorf("Could not parse regex in --replace_regex: sql:%v", originalString) + } + return ret, nil +} diff --git a/src/util_test.go b/src/util_test.go new file mode 100644 index 0000000..2b0e428 --- /dev/null +++ b/src/util_test.go @@ -0,0 +1,106 @@ +// Copyright 2023 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParseReplaceRegex(t *testing.T) { + testCases := []struct { + succ bool + regexpStr string + input string + output string + }{ + { + succ: true, + regexpStr: `/dll/so/`, + input: "a.dll.dll", + output: "a.so.so", + }, + { + succ: true, + regexpStr: `/\.dll/.so/`, + input: "a.dlldll", + output: "a.sodll", + }, + { + succ: true, + regexpStr: `/\.\.dll/..so/`, + input: "a.dll..dlldll", + output: "a.dll..sodll", + }, + { + succ: true, + regexpStr: `/conn=[0-9]+/conn=/`, + input: "Some infos [conn=2097154]", + output: "Some infos [conn=]", + }, + { + succ: true, + regexpStr: `/conn=[0-9]+/conn=/`, + input: "Some infos [conn=xxx]", + output: "Some infos [conn=xxx]", + }, + { + succ: true, + regexpStr: `/a/\/b\r\t/`, + input: "a", + output: "/b\r\t", + }, + { + succ: false, + regexpStr: `/conn=[0-9]+/conn=`, + input: "", + output: "", + }, + { + succ: false, + regexpStr: `conn=[0-9]+/conn=/`, + input: "", + output: "", + }, + { + succ: false, + regexpStr: `/*/conn=/`, + input: "", + output: "", + }, + { + succ: false, + regexpStr: `/abc\/conn=\/`, + input: "", + output: "", + }, + } + + for _, testCase := range testCases { + regexs, err := ParseReplaceRegex(testCase.regexpStr) + if !testCase.succ { + require.NotNil(t, err) + continue + } + require.Nil(t, err) + require.NotNil(t, regexs) + + result := testCase.input + for _, reg := range regexs { + result = reg.regex.ReplaceAllString(result, reg.replace) + } + require.Equal(t, testCase.output, result) + } +} diff --git a/t/example.test b/t/example.test index 5e02fbd..ba718a7 100644 --- a/t/example.test +++ b/t/example.test @@ -1,3 +1,9 @@ create table t(a bigint, b bigint); insert into t values(1, 1), (2, 2), (3, 3), (4, 4), (5, 5); select * from t where a = 1; + +--replace_regex /:[ ]?[.0-9]+.*?,/:,/ /:[ ]?[.0-9]+.*?}/:}/ /[0-9]+ Bytes/ Bytes/ +explain analyze format='brief' select * from t; + +--replace_regex /:[ ]?[.0-9]+.*?,/:,/ /:[ ]?[.0-9]+.*?}/:}/ /[0-9]+ Bytes/ Bytes/ +explain analyze select * from t;