parser: add RemovePlacementRuleFlag (#33161)

close #33160

parser/ast/ddl.go

@@ -175,6 +175,9 @@ type AlterDatabaseStmt struct {
 
 // Restore implements Node interface.
 func (n *AlterDatabaseStmt) Restore(ctx *format.RestoreCtx) error {
+	if ctx.Flags.HasSkipPlacementRuleForRestoreFlag() && n.isAllPlacementOptions() {
+		return nil
+	}
 	// If all options placement options and RestoreTiDBSpecialComment flag is on,
 	// we should restore the whole node in special comment. For example, the restore result should be:
 	// /*T![placement] ALTER DATABASE `db1` PLACEMENT POLICY = `p1` */
@@ -1936,6 +1939,9 @@ type PlacementOption struct {
 }
 
 func (n *PlacementOption) Restore(ctx *format.RestoreCtx) error {
+	if ctx.Flags.HasSkipPlacementRuleForRestoreFlag() {
+		return nil
+	}
 	fn := func() error {
 		switch n.Tp {
 		case PlacementOptionPrimaryRegion:
@@ -2324,6 +2330,9 @@ func (n *TableOption) Restore(ctx *format.RestoreCtx) error {
 		ctx.WritePlain("= ")
 		ctx.WriteString(n.StrValue)
 	case TableOptionPlacementPolicy:
+		if ctx.Flags.HasSkipPlacementRuleForRestoreFlag() {
+			return nil
+		}
 		placementOpt := PlacementOption{
 			Tp:        PlacementOptionPolicy,
 			UintValue: n.UintValue,
@@ -2685,8 +2694,25 @@ func (n *AlterOrderItem) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+func (n *AlterTableSpec) IsAllPlacementRule() bool {
+	switch n.Tp {
+	case AlterTablePartitionAttributes, AlterTablePartitionOptions, AlterTableOption, AlterTableAttributes:
+		for _, o := range n.Options {
+			if o.Tp != TableOptionPlacementPolicy {
+				return false
+			}
+		}
+		return true
+	default:
+		return false
+	}
+}
+
 // Restore implements Node interface.
 func (n *AlterTableSpec) Restore(ctx *format.RestoreCtx) error {
+	if n.IsAllPlacementRule() && ctx.Flags.HasSkipPlacementRuleForRestoreFlag() {
+		return nil
+	}
 	switch n.Tp {
 	case AlterTableSetTiFlashReplica:
 		ctx.WriteKeyWord("SET TIFLASH REPLICA ")
@@ -3275,13 +3301,36 @@ type AlterTableStmt struct {
 	Specs []*AlterTableSpec
 }
 
+func (n *AlterTableStmt) HaveOnlyPlacementOptions() bool {
+	for _, n := range n.Specs {
+		if n.Tp == AlterTablePartitionOptions {
+			if !n.IsAllPlacementRule() {
+				return false
+			}
+		} else {
+			return false
+		}
+
+	}
+	return true
+}
+
 // Restore implements Node interface.
 func (n *AlterTableStmt) Restore(ctx *format.RestoreCtx) error {
+	if ctx.Flags.HasSkipPlacementRuleForRestoreFlag() && n.HaveOnlyPlacementOptions() {
+		return nil
+	}
 	ctx.WriteKeyWord("ALTER TABLE ")
 	if err := n.Table.Restore(ctx); err != nil {
 		return errors.Annotate(err, "An error occurred while restore AlterTableStmt.Table")
 	}
-	for i, spec := range n.Specs {
+	var specs []*AlterTableSpec
+	for _, spec := range n.Specs {
+		if !(spec.IsAllPlacementRule() && ctx.Flags.HasSkipPlacementRuleForRestoreFlag()) {
+			specs = append(specs, spec)
+		}
+	}
+	for i, spec := range specs {
 		if i == 0 || spec.Tp == AlterTablePartition || spec.Tp == AlterTableRemovePartitioning || spec.Tp == AlterTableImportTablespace || spec.Tp == AlterTableDiscardTablespace {
 			ctx.WritePlain(" ")
 		} else {
@@ -3991,6 +4040,9 @@ type AlterPlacementPolicyStmt struct {
 }
 
 func (n *AlterPlacementPolicyStmt) Restore(ctx *format.RestoreCtx) error {
+	if ctx.Flags.HasSkipPlacementRuleForRestoreFlag() {
+		return nil
+	}
 	if ctx.Flags.HasTiDBSpecialCommentFlag() {
 		return restorePlacementStmtInSpecialComment(ctx, n)
 	}

parser/ast/ddl_test.go

@@ -754,3 +754,76 @@ func TestDropPlacementPolicyRestore(t *testing.T) {
 		runNodeRestoreTestWithFlags(t, testCases, "%s", extractNodeFunc, ca.flags)
 	}
 }
+
+func TestRemovePlacementRestore(t *testing.T) {
+	f := format.DefaultRestoreFlags | format.SkipPlacementRuleForRestore
+	cases := []struct {
+		sourceSQL string
+		expectSQL string
+	}{
+		{
+			"CREATE TABLE t1 (id BIGINT NOT NULL PRIMARY KEY auto_increment, b varchar(255)) PLACEMENT POLICY=placement1;",
+			"CREATE TABLE `t1` (`id` BIGINT NOT NULL PRIMARY KEY AUTO_INCREMENT,`b` VARCHAR(255)) ",
+		},
+		{
+			"CREATE TABLE `t1` (\n  `a` int(11) DEFAULT NULL\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin /*T![placement] PLACEMENT POLICY=`p2` */",
+			"CREATE TABLE `t1` (`a` INT(11) DEFAULT NULL) ENGINE = InnoDB DEFAULT CHARACTER SET = UTF8MB4 DEFAULT COLLATE = UTF8MB4_BIN ",
+		},
+		{
+			"CREATE TABLE t4 (firstname VARCHAR(25) NOT NULL,lastname VARCHAR(25) NOT NULL,username VARCHAR(16) NOT NULL,email VARCHAR(35),joined DATE NOT NULL) PARTITION BY RANGE( YEAR(joined) ) (PARTITION p0 VALUES LESS THAN (1960) PLACEMENT POLICY=p1,PARTITION p1 VALUES LESS THAN (1970),PARTITION p2 VALUES LESS THAN (1980),PARTITION p3 VALUES LESS THAN (1990),PARTITION p4 VALUES LESS THAN MAXVALUE);",
+			"CREATE TABLE `t4` (`firstname` VARCHAR(25) NOT NULL,`lastname` VARCHAR(25) NOT NULL,`username` VARCHAR(16) NOT NULL,`email` VARCHAR(35),`joined` DATE NOT NULL) PARTITION BY RANGE (YEAR(`joined`)) (PARTITION `p0` VALUES LESS THAN (1960) ,PARTITION `p1` VALUES LESS THAN (1970),PARTITION `p2` VALUES LESS THAN (1980),PARTITION `p3` VALUES LESS THAN (1990),PARTITION `p4` VALUES LESS THAN (MAXVALUE))",
+		},
+		{
+			"ALTER TABLE t3 PLACEMENT POLICY=DEFAULT;",
+			"ALTER TABLE `t3`",
+		},
+		{
+			"ALTER TABLE t1 PLACEMENT POLICY=p10",
+			"ALTER TABLE `t1`",
+		},
+		{
+			"ALTER TABLE t1 PLACEMENT POLICY=p10, add d text(50)",
+			"ALTER TABLE `t1` ADD COLUMN `d` TEXT(50)",
+		},
+		{
+			"alter table tp PARTITION p1 placement policy p2",
+			"",
+		},
+		{
+			"alter table t add d text(50) PARTITION p1 placement policy p2",
+			"ALTER TABLE `t` ADD COLUMN `d` TEXT(50)",
+		},
+		{
+			"alter table tp set tiflash replica 1 PARTITION p1 placement policy p2",
+			"ALTER TABLE `tp` SET TIFLASH REPLICA 1",
+		},
+		{
+			"ALTER DATABASE TestResetPlacementDB PLACEMENT POLICY SET DEFAULT",
+			"",
+		},
+
+		{
+			"ALTER DATABASE TestResetPlacementDB PLACEMENT POLICY p1 charset utf8mb4",
+			"ALTER DATABASE `TestResetPlacementDB`  CHARACTER SET = utf8mb4",
+		},
+		{
+			"/*T![placement] ALTER DATABASE `db1` PLACEMENT POLICY = `p1` */",
+			"",
+		},
+		{
+			"ALTER PLACEMENT POLICY p3 PRIMARY_REGION='us-east-1' REGIONS='us-east-1,us-east-2,us-west-1';",
+			"",
+		},
+	}
+
+	extractNodeFunc := func(node Node) Node {
+		return node
+	}
+
+	for _, ca := range cases {
+		testCases := []NodeRestoreTestCase{
+			{ca.sourceSQL, ca.expectSQL},
+		}
+		runNodeRestoreTestWithFlagsStmtChange(t, testCases, "%s", extractNodeFunc, f)
+	}
+}

parser/ast/dml_test.go

@@ -17,6 +17,7 @@ import (
 	"testing"
 
 	. "github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/format"
 	"github.com/stretchr/testify/require"
 )
 
@@ -238,7 +239,7 @@ func TestJoinRestore(t *testing.T) {
 		return node.(*SelectStmt).From.TableRefs
 	}
 	runNodeRestoreTest(t, testCases, "select * from %s", extractNodeFunc)
-	runNodeRestoreTestWithFlagsStmtChange(t, testChangedCases, "select * from %s", extractNodeFunc)
+	runNodeRestoreTestWithFlagsStmtChange(t, testChangedCases, "select * from %s", extractNodeFunc, format.DefaultRestoreFlags)
 }
 
 func TestTableRefsClauseRestore(t *testing.T) {

parser/ast/util_test.go

@@ -190,7 +190,7 @@ func runNodeRestoreTestWithFlags(t *testing.T, nodeTestCases []NodeRestoreTestCa
 
 // runNodeRestoreTestWithFlagsStmtChange likes runNodeRestoreTestWithFlags but not check if the ASTs are same.
 // Sometimes the AST are different and it's expected.
-func runNodeRestoreTestWithFlagsStmtChange(t *testing.T, nodeTestCases []NodeRestoreTestCase, template string, extractNodeFunc func(node Node) Node) {
+func runNodeRestoreTestWithFlagsStmtChange(t *testing.T, nodeTestCases []NodeRestoreTestCase, template string, extractNodeFunc func(node Node) Node, flags RestoreFlags) {
 	p := parser.New()
 	p.EnableWindowFunc(true)
 	for _, testCase := range nodeTestCases {
@@ -200,7 +200,7 @@ func runNodeRestoreTestWithFlagsStmtChange(t *testing.T, nodeTestCases []NodeRes
 		comment := fmt.Sprintf("source %#v", testCase)
 		require.NoError(t, err, comment)
 		var sb strings.Builder
-		err = extractNodeFunc(stmt).Restore(NewRestoreCtx(DefaultRestoreFlags, &sb))
+		err = extractNodeFunc(stmt).Restore(NewRestoreCtx(flags, &sb))
 		require.NoError(t, err, comment)
 		restoreSql := fmt.Sprintf(template, sb.String())
 		comment = fmt.Sprintf("source %#v; restore %v", testCase, restoreSql)

parser/format/format.go

@@ -224,6 +224,7 @@ const (
 	RestoreStringWithoutDefaultCharset
 
 	RestoreTiDBSpecialComment
+	SkipPlacementRuleForRestore
 )
 
 const (
@@ -300,6 +301,10 @@ func (rf RestoreFlags) HasTiDBSpecialCommentFlag() bool {
 	return rf.has(RestoreTiDBSpecialComment)
 }
 
+func (rf RestoreFlags) HasSkipPlacementRuleForRestoreFlag() bool {
+	return rf.has(SkipPlacementRuleForRestore)
+}
+
 // RestoreCtx is `Restore` context to hold flags and writer.
 type RestoreCtx struct {
 	Flags     RestoreFlags