diff --git a/sessionctx/binloginfo/binloginfo.go b/sessionctx/binloginfo/binloginfo.go index d4c7bb0be8ef1..67c1953fc9114 100644 --- a/sessionctx/binloginfo/binloginfo.go +++ b/sessionctx/binloginfo/binloginfo.go @@ -15,6 +15,7 @@ package binloginfo import ( "context" + "math" "regexp" "strings" "sync" @@ -29,7 +30,7 @@ import ( "github.com/pingcap/tidb/metrics" "github.com/pingcap/tidb/sessionctx" "github.com/pingcap/tidb/util/logutil" - binlog "github.com/pingcap/tipb/go-binlog" + "github.com/pingcap/tipb/go-binlog" "go.uber.org/zap" "google.golang.org/grpc" ) @@ -42,6 +43,8 @@ func init() { // shared by all sessions. var pumpsClient *pumpcli.PumpsClient var pumpsClientLock sync.RWMutex +var shardPat = regexp.MustCompile(`SHARD_ROW_ID_BITS\s*=\s*\d+\s*`) +var preSplitPat = regexp.MustCompile(`PRE_SPLIT_REGIONS\s*=\s*\d+\s*`) // BinlogInfo contains binlog data and binlog client. type BinlogInfo struct { @@ -136,7 +139,7 @@ func SetDDLBinlog(client *pumpcli.PumpsClient, txn kv.Transaction, jobID int64, return } - ddlQuery = addSpecialComment(ddlQuery) + ddlQuery = AddSpecialComment(ddlQuery) info := &BinlogInfo{ Data: &binlog.Binlog{ Tp: binlog.BinlogType_Prewrite, @@ -150,18 +153,52 @@ func SetDDLBinlog(client *pumpcli.PumpsClient, txn kv.Transaction, jobID int64, const specialPrefix = `/*!90000 ` -func addSpecialComment(ddlQuery string) string { +// AddSpecialComment uses to add comment for table option in DDL query. +// Export for testing. +func AddSpecialComment(ddlQuery string) string { if strings.Contains(ddlQuery, specialPrefix) { return ddlQuery } + return addSpecialCommentByRegexps(ddlQuery, shardPat, preSplitPat) +} + +// addSpecialCommentByRegexps uses to add special comment for the worlds in the ddlQuery with match the regexps. +func addSpecialCommentByRegexps(ddlQuery string, regs ...*regexp.Regexp) string { upperQuery := strings.ToUpper(ddlQuery) - reg, err := regexp.Compile(`SHARD_ROW_ID_BITS\s*=\s*\d+`) - terror.Log(err) - loc := reg.FindStringIndex(upperQuery) - if len(loc) < 2 { - return ddlQuery + var specialComments []string + minIdx := math.MaxInt64 + for i := 0; i < len(regs); { + reg := regs[i] + loc := reg.FindStringIndex(upperQuery) + if len(loc) < 2 { + i++ + continue + } + specialComments = append(specialComments, ddlQuery[loc[0]:loc[1]]) + if loc[0] < minIdx { + minIdx = loc[0] + } + ddlQuery = ddlQuery[:loc[0]] + ddlQuery[loc[1]:] + upperQuery = upperQuery[:loc[0]] + upperQuery[loc[1]:] + } + if minIdx != math.MaxInt64 { + query := ddlQuery[:minIdx] + specialPrefix + for _, comment := range specialComments { + if query[len(query)-1] != ' ' { + query += " " + } + query += comment + } + if query[len(query)-1] != ' ' { + query += " " + } + query += "*/" + if len(ddlQuery[minIdx:]) > 0 { + return query + " " + ddlQuery[minIdx:] + } + return query } - return ddlQuery[:loc[0]] + specialPrefix + ddlQuery[loc[0]:loc[1]] + ` */` + ddlQuery[loc[1]:] + return ddlQuery } // MockPumpsClient creates a PumpsClient, used for test. diff --git a/sessionctx/binloginfo/binloginfo_test.go b/sessionctx/binloginfo/binloginfo_test.go index 86c20c87751ba..a10b064811d9a 100644 --- a/sessionctx/binloginfo/binloginfo_test.go +++ b/sessionctx/binloginfo/binloginfo_test.go @@ -431,3 +431,44 @@ func (s *testBinlogSuite) TestDeleteSchema(c *C) { tk.MustExec("delete from b1 where job_id in (select job_id from b2 where batch_class = 'TEST') or split_job_id in (select job_id from b2 where batch_class = 'TEST');") tk.MustExec("delete b1 from b2 right join b1 on b1.job_id = b2.job_id and batch_class = 'TEST';") } + +func (s *testBinlogSuite) TestAddSpecialComment(c *C) { + testCase := []struct { + input string + result string + }{ + { + "create table t1 (id int ) shard_row_id_bits=2;", + "create table t1 (id int ) /*!90000 shard_row_id_bits=2 */ ;", + }, + { + "create table t1 (id int ) shard_row_id_bits=2 pre_split_regions=2;", + "create table t1 (id int ) /*!90000 shard_row_id_bits=2 pre_split_regions=2 */ ;", + }, + { + "create table t1 (id int ) shard_row_id_bits=2 pre_split_regions=2;", + "create table t1 (id int ) /*!90000 shard_row_id_bits=2 pre_split_regions=2 */ ;", + }, + + { + "create table t1 (id int ) shard_row_id_bits=2 engine=innodb pre_split_regions=2;", + "create table t1 (id int ) /*!90000 shard_row_id_bits=2 pre_split_regions=2 */ engine=innodb ;", + }, + { + "create table t1 (id int ) pre_split_regions=2 shard_row_id_bits=2;", + "create table t1 (id int ) /*!90000 shard_row_id_bits=2 pre_split_regions=2 */ ;", + }, + { + "create table t6 (id int ) shard_row_id_bits=2 shard_row_id_bits=3 pre_split_regions=2;", + "create table t6 (id int ) /*!90000 shard_row_id_bits=2 shard_row_id_bits=3 pre_split_regions=2 */ ;", + }, + { + "alter table t shard_row_id_bits=2 ", + "alter table t /*!90000 shard_row_id_bits=2 */", + }, + } + for _, ca := range testCase { + re := binloginfo.AddSpecialComment(ca.input) + c.Assert(re, Equals, ca.result) + } +}