Support --replace_regex (#112)

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
 )

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=

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:<num>, loops:<num>, RU:<num>, cop_task: {num:<num>, max:<num>, proc_keys:<num>, rpc_num:<num>, rpc_time:<num>, copr_cache_hit_ratio:<num>, build_task_duration:<num>, max_distsql_concurrency:<num>}	data:TableFullScan	<num> Bytes	N/A
+└─TableFullScan	10000.00	5	cop[tikv]	table:t	tikv_task:{time:<num>, loops:<num>}	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:<num>, loops:<num>, RU:<num>, cop_task: {num:<num>, max:<num>, proc_keys:<num>, rpc_num:<num>, rpc_time:<num>, copr_cache_hit_ratio:<num>, build_task_duration:<num>, max_distsql_concurrency:<num>}	data:TableFullScan_4	<num> Bytes	N/A
+└─TableFullScan_4	10000.00	5	cop[tikv]	table:t	tikv_task:{time:<num>, loops:<num>}	keep order:false, stats:pseudo	N/A	N/A

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 <static data> 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)
 		}

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
+}

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=<num>/`,
+			input:     "Some infos [conn=2097154]",
+			output:    "Some infos [conn=<num>]",
+		},
+		{
+			succ:      true,
+			regexpStr: `/conn=[0-9]+/conn=<num>/`,
+			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=<num>`,
+			input:     "",
+			output:    "",
+		},
+		{
+			succ:      false,
+			regexpStr: `conn=[0-9]+/conn=<num>/`,
+			input:     "",
+			output:    "",
+		},
+		{
+			succ:      false,
+			regexpStr: `/*/conn=<num>/`,
+			input:     "",
+			output:    "",
+		},
+		{
+			succ:      false,
+			regexpStr: `/abc\/conn=<num>\/`,
+			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)
+	}
+}

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]+.*?,/:<num>,/ /:[ ]?[.0-9]+.*?}/:<num>}/ /[0-9]+ Bytes/<num> Bytes/
+explain analyze format='brief' select * from t;
+
+--replace_regex /:[ ]?[.0-9]+.*?,/:<num>,/ /:[ ]?[.0-9]+.*?}/:<num>}/ /[0-9]+ Bytes/<num> Bytes/
+explain analyze select * from t;