diff --git a/ast/ddl.go b/ast/ddl.go index a0b26877f..ec75906c7 100644 --- a/ast/ddl.go +++ b/ast/ddl.go @@ -20,6 +20,7 @@ import ( "github.com/pingcap/parser/model" "github.com/pingcap/parser/mysql" "github.com/pingcap/parser/terror" + "github.com/pingcap/parser/tidb" "github.com/pingcap/parser/types" ) @@ -501,7 +502,9 @@ func (n *ColumnOption) Restore(ctx *format.RestoreCtx) error { pkTp := n.PrimaryKeyTp.String() if len(pkTp) != 0 { ctx.WritePlain(" ") - ctx.WriteKeyWord(pkTp) + ctx.WriteWithSpecialComments(tidb.FeatureIDClusteredIndex, func() { + ctx.WriteKeyWord(pkTp) + }) } case ColumnOptionNotNull: ctx.WriteKeyWord("NOT NULL") @@ -574,10 +577,12 @@ func (n *ColumnOption) Restore(ctx *format.RestoreCtx) error { ctx.WriteKeyWord("STORAGE ") ctx.WriteKeyWord(n.StrValue) case ColumnOptionAutoRandom: - ctx.WriteKeyWord("AUTO_RANDOM") - if n.AutoRandomBitLength != types.UnspecifiedLength { - ctx.WritePlainf("(%d)", n.AutoRandomBitLength) - } + ctx.WriteWithSpecialComments(tidb.FeatureIDAutoRandom, func() { + ctx.WriteKeyWord("AUTO_RANDOM") + if n.AutoRandomBitLength != types.UnspecifiedLength { + ctx.WritePlainf("(%d)", n.AutoRandomBitLength) + } + }) default: return errors.New("An error occurred while splicing ColumnOption") } @@ -1891,22 +1896,32 @@ func (n *TableOption) Restore(ctx *format.RestoreCtx) error { ctx.WriteKeyWord(n.StrValue) case TableOptionAutoIncrement: if n.BoolValue { - ctx.WriteKeyWord("FORCE ") + ctx.WriteWithSpecialComments(tidb.FeatureIDForceAutoInc, func() { + ctx.WriteKeyWord("FORCE") + }) + ctx.WritePlain(" ") } ctx.WriteKeyWord("AUTO_INCREMENT ") ctx.WritePlain("= ") ctx.WritePlainf("%d", n.UintValue) case TableOptionAutoIdCache: - ctx.WriteKeyWord("AUTO_ID_CACHE ") - ctx.WritePlain("= ") - ctx.WritePlainf("%d", n.UintValue) + ctx.WriteWithSpecialComments(tidb.FeatureIDAutoIDCache, func() { + ctx.WriteKeyWord("AUTO_ID_CACHE ") + ctx.WritePlain("= ") + ctx.WritePlainf("%d", n.UintValue) + }) case TableOptionAutoRandomBase: if n.BoolValue { - ctx.WriteKeyWord("FORCE ") + ctx.WriteWithSpecialComments(tidb.FeatureIDForceAutoInc, func() { + ctx.WriteKeyWord("FORCE") + }) + ctx.WritePlain(" ") } - ctx.WriteKeyWord("AUTO_RANDOM_BASE ") - ctx.WritePlain("= ") - ctx.WritePlainf("%d", n.UintValue) + ctx.WriteWithSpecialComments(tidb.FeatureIDAutoRandomBase, func() { + ctx.WriteKeyWord("AUTO_RANDOM_BASE ") + ctx.WritePlain("= ") + ctx.WritePlainf("%d", n.UintValue) + }) case TableOptionComment: ctx.WriteKeyWord("COMMENT ") ctx.WritePlain("= ") @@ -1997,11 +2012,15 @@ func (n *TableOption) Restore(ctx *format.RestoreCtx) error { ctx.WritePlainf("%d", n.UintValue) } case TableOptionShardRowID: - ctx.WriteKeyWord("SHARD_ROW_ID_BITS ") - ctx.WritePlainf("= %d", n.UintValue) + ctx.WriteWithSpecialComments(tidb.FeatureIDTiDB, func() { + ctx.WriteKeyWord("SHARD_ROW_ID_BITS ") + ctx.WritePlainf("= %d", n.UintValue) + }) case TableOptionPreSplitRegion: - ctx.WriteKeyWord("PRE_SPLIT_REGIONS ") - ctx.WritePlainf("= %d", n.UintValue) + ctx.WriteWithSpecialComments(tidb.FeatureIDTiDB, func() { + ctx.WriteKeyWord("PRE_SPLIT_REGIONS ") + ctx.WritePlainf("= %d", n.UintValue) + }) case TableOptionPackKeys: // TODO: not support ctx.WriteKeyWord("PACK_KEYS ") diff --git a/format/format.go b/format/format.go index 575adfb7a..f1a9147c7 100644 --- a/format/format.go +++ b/format/format.go @@ -221,6 +221,8 @@ const ( RestoreStringWithoutCharset RestoreStringWithoutDefaultCharset + + RestoreTiDBSpecialComment ) const ( @@ -289,6 +291,10 @@ func (rf RestoreFlags) HasStringWithoutCharset() bool { return rf.has(RestoreStringWithoutCharset) } +func (rf RestoreFlags) HasTiDBSpecialCommentFlag() bool { + return rf.has(RestoreTiDBSpecialComment) +} + // RestoreCtx is `Restore` context to hold flags and writer. type RestoreCtx struct { Flags RestoreFlags @@ -314,6 +320,20 @@ func (ctx *RestoreCtx) WriteKeyWord(keyWord string) { fmt.Fprint(ctx.In, keyWord) } +func (ctx *RestoreCtx) WriteWithSpecialComments(featureID string, fn func()) { + if !ctx.Flags.HasTiDBSpecialCommentFlag() { + fn() + return + } + ctx.WritePlain("/*T!") + if len(featureID) != 0 { + ctx.WritePlainf("[%s]", featureID) + } + ctx.WritePlain(" ") + fn() + ctx.WritePlain(" */") +} + // WriteString writes the string into writer // `str` may be wrapped in quotes and escaped according to RestoreFlags. func (ctx *RestoreCtx) WriteString(str string) { diff --git a/format/format_test.go b/format/format_test.go index 835654f89..9122edbe2 100644 --- a/format/format_test.go +++ b/format/format_test.go @@ -95,3 +95,19 @@ func (s *testRestoreCtxSuite) TestRestoreCtx(c *C) { c.Assert(sb.String(), Equals, testCase.expect, Commentf("case: %#v", testCase)) } } + +func (s *testRestoreCtxSuite) TestRestoreSpecialComment(c *C) { + var sb strings.Builder + sb.Reset() + ctx := NewRestoreCtx(RestoreTiDBSpecialComment, &sb) + ctx.WriteWithSpecialComments("fea_id", func() { + ctx.WritePlain("content") + }) + c.Assert(sb.String(), Equals, "/*T![fea_id] content */") + + sb.Reset() + ctx.WriteWithSpecialComments("", func() { + ctx.WritePlain("shard_row_id_bits") + }) + c.Assert(sb.String(), Equals, "/*T! shard_row_id_bits */") +} diff --git a/lexer.go b/lexer.go index f5a0efc70..78cd4432f 100644 --- a/lexer.go +++ b/lexer.go @@ -22,6 +22,7 @@ import ( "unicode/utf8" "github.com/pingcap/parser/mysql" + tidbfeature "github.com/pingcap/parser/tidb" ) var _ = yyLexer(&Scanner{}) @@ -397,11 +398,10 @@ func startWithSlash(s *Scanner) (tok int, pos Pos, lit string) { s.r.inc() // in '/*T!', try to match the pattern '/*T![feature1,feature2,...]'. features := s.scanFeatureIDs() - if SpecialCommentsController.ContainsAll(features) { + if tidbfeature.CanParseFeature(features...) { s.inBangComment = true return s.scan() } - case 'M': // '/*M' maybe MariaDB-specific comments // no special treatment for now. break diff --git a/lexer_test.go b/lexer_test.go index 459b1d34d..f9e4c2720 100644 --- a/lexer_test.go +++ b/lexer_test.go @@ -159,7 +159,6 @@ func runTest(c *C, table []testCaseItem) { } func (s *testLexerSuite) TestComment(c *C) { - SpecialCommentsController.Register("test") table := []testCaseItem{ {"-- select --\n1", intLit}, {"/*!40101 SET character_set_client = utf8 */;", set}, @@ -178,8 +177,8 @@ SELECT`, selectKwd}, // The odd behavior of '*/' inside conditional comment is the same as // that of MySQL. - {"/*T![unsupported] '*/0 -- ' */", intLit}, // equivalent to 0 - {"/*T![test] '*/0 -- ' */", stringLit}, // equivalent to '*/0 -- ' + {"/*T![unsupported] '*/0 -- ' */", intLit}, // equivalent to 0 + {"/*T![auto_rand] '*/0 -- ' */", stringLit}, // equivalent to '*/0 -- ' } runTest(c, table) } @@ -273,7 +272,6 @@ func (s *testLexerSuite) TestSpecialComment(c *C) { } func (s *testLexerSuite) TestFeatureIDsComment(c *C) { - SpecialCommentsController.Register("auto_rand") l := NewScanner("/*T![auto_rand] auto_random(5) */") tok, pos, lit := l.scan() c.Assert(tok, Equals, identifier) @@ -317,7 +315,6 @@ func (s *testLexerSuite) TestOptimizerHint(c *C) { } func (s *testLexerSuite) TestOptimizerHintAfterCertainKeywordOnly(c *C) { - SpecialCommentsController.Register("test") tests := []struct { input string tokens []int @@ -359,7 +356,7 @@ func (s *testLexerSuite) TestOptimizerHintAfterCertainKeywordOnly(c *C) { tokens: []int{selectKwd, '*', 0}, }, { - input: "SELECT /*T![test] * */ /*+ hint */", + input: "SELECT /*T![auto_rand] * */ /*+ hint */", tokens: []int{selectKwd, '*', 0}, }, { diff --git a/misc.go b/misc.go index aaf5ecc90..be3b1e6ad 100644 --- a/misc.go +++ b/misc.go @@ -997,38 +997,3 @@ func handleIdent(lval *yySymType) int { lval.ident = cs return underscoreCS } - -// SpecialCommentsController controls whether special comments like `/*T![xxx] yyy */` -// can be parsed as `yyy`. To add such rules, please use SpecialCommentsController.Register(). -// For example: -// SpecialCommentsController.Register("30100"); -// Now the parser will treat -// select a, /*T![30100] mysterious_keyword */ from t; -// and -// select a, mysterious_keyword from t; -// equally. -// Similar special comments without registration are ignored by parser. -var SpecialCommentsController = specialCommentsCtrl{ - supportedFeatures: map[string]struct{}{}, -} - -type specialCommentsCtrl struct { - supportedFeatures map[string]struct{} -} - -func (s *specialCommentsCtrl) Register(featureID string) { - s.supportedFeatures[featureID] = struct{}{} -} - -func (s *specialCommentsCtrl) Unregister(featureID string) { - delete(s.supportedFeatures, featureID) -} - -func (s *specialCommentsCtrl) ContainsAll(featureIDs []string) bool { - for _, f := range featureIDs { - if _, found := s.supportedFeatures[f]; !found { - return false - } - } - return true -} diff --git a/tidb/features.go b/tidb/features.go new file mode 100644 index 000000000..5cfc14d00 --- /dev/null +++ b/tidb/features.go @@ -0,0 +1,46 @@ +// Copyright 2021 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 tidb + +const ( + // FeatureIDTiDB represents the general TiDB-specific features. + FeatureIDTiDB = "" + // FeatureIDAutoRandom is the `auto_random` feature. + FeatureIDAutoRandom = "auto_rand" + // FeatureIDAutoIDCache is the `auto_id_cache` feature. + FeatureIDAutoIDCache = "auto_id_cache" + // FeatureIDAutoRandomBase is the `auto_random_base` feature. + FeatureIDAutoRandomBase = "auto_rand_base" + // FeatureIDClusteredIndex is the `clustered_index` feature. + FeatureIDClusteredIndex = "clustered_index" + // FeatureIDForceAutoInc is the `force auto_increment` feature. + FeatureIDForceAutoInc = "force_inc" +) + +var featureIDs = map[string]struct{}{ + FeatureIDAutoRandom: {}, + FeatureIDAutoIDCache: {}, + FeatureIDAutoRandomBase: {}, + FeatureIDClusteredIndex: {}, + FeatureIDForceAutoInc: {}, +} + +func CanParseFeature(fs ...string) bool { + for _, f := range fs { + if _, ok := featureIDs[f]; !ok { + return false + } + } + return true +}