restore: support restoring SQL wrapping with TiDB special comments

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 ")

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) {

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 */")
+}

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

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},
 		},
 		{

misc.go

@@ -993,38 +993,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
-}

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