pkg/parser: use clearer structures for RESOURCE GROUP ASTs

pkg/ddl/executor.go

@@ -3734,7 +3734,7 @@ func SetDirectResourceGroupSettings(groupInfo *model.ResourceGroupInfo, opt *ast
 			resourceGroupSettings.Runaway = nil
 		}
 		for _, opt := range opt.RunawayOptionList {
-			if err := SetDirectResourceGroupRunawayOption(resourceGroupSettings, opt.Tp, opt.StrValue, opt.IntValue); err != nil {
+			if err := SetDirectResourceGroupRunawayOption(resourceGroupSettings, opt); err != nil {
 				return err
 			}
 		}
@@ -3758,25 +3758,25 @@ func SetDirectResourceGroupSettings(groupInfo *model.ResourceGroupInfo, opt *ast
 }
 
 // SetDirectResourceGroupRunawayOption tries to set runaway part of the ResourceGroupSettings.
-func SetDirectResourceGroupRunawayOption(resourceGroupSettings *model.ResourceGroupSettings, typ ast.RunawayOptionType, stringVal string, intVal int32) error {
+func SetDirectResourceGroupRunawayOption(resourceGroupSettings *model.ResourceGroupSettings, opt *ast.ResourceGroupRunawayOption) error {
 	if resourceGroupSettings.Runaway == nil {
 		resourceGroupSettings.Runaway = &model.ResourceGroupRunawaySettings{}
 	}
 	settings := resourceGroupSettings.Runaway
-	switch typ {
+	switch opt.Tp {
 	case ast.RunawayRule:
 		// because execute time won't be too long, we use `time` pkg which does not support to parse unit 'd'.
-		dur, err := time.ParseDuration(stringVal)
+		dur, err := time.ParseDuration(opt.RuleOption.ExecElapsed)
 		if err != nil {
 			return err
 		}
 		settings.ExecElapsedTimeMs = uint64(dur.Milliseconds())
 	case ast.RunawayAction:
-		settings.Action = model.RunawayActionType(intVal)
+		settings.Action = opt.ActionOption.Type
 	case ast.RunawayWatch:
-		settings.WatchType = model.RunawayWatchType(intVal)
-		if len(stringVal) > 0 {
-			dur, err := time.ParseDuration(stringVal)
+		settings.WatchType = opt.WatchOption.Type
+		if dur := opt.WatchOption.Duration; len(dur) > 0 {
+			dur, err := time.ParseDuration(dur)
 			if err != nil {
 				return err
 			}

pkg/executor/internal/querywatch/query_watch.go

@@ -42,8 +42,9 @@ func setWatchOption(ctx context.Context,
 ) error {
 	switch op.Tp {
 	case ast.QueryWatchResourceGroup:
-		if op.ExprValue != nil {
-			expr, err := plannerutil.RewriteAstExprWithPlanCtx(sctx.GetPlanCtx(), op.ExprValue, nil, nil, false)
+		resourceGroupOption := op.ResourceGroupOption
+		if resourceGroupOption.GroupNameExpr != nil {
+			expr, err := plannerutil.RewriteAstExprWithPlanCtx(sctx.GetPlanCtx(), resourceGroupOption.GroupNameExpr, nil, nil, false)
 			if err != nil {
 				return err
 			}
@@ -56,12 +57,13 @@ func setWatchOption(ctx context.Context,
 			}
 			record.ResourceGroupName = name
 		} else {
-			record.ResourceGroupName = op.StrValue.L
+			record.ResourceGroupName = resourceGroupOption.GroupNameStr.L
 		}
 	case ast.QueryWatchAction:
-		record.Action = rmpb.RunawayAction(op.IntValue)
+		record.Action = rmpb.RunawayAction(op.ActionOption.Type)
 	case ast.QueryWatchType:
-		expr, err := plannerutil.RewriteAstExprWithPlanCtx(sctx.GetPlanCtx(), op.ExprValue, nil, nil, false)
+		textOption := op.TextOption
+		expr, err := plannerutil.RewriteAstExprWithPlanCtx(sctx.GetPlanCtx(), textOption.PatternExpr, nil, nil, false)
 		if err != nil {
 			return err
 		}
@@ -72,8 +74,9 @@ func setWatchOption(ctx context.Context,
 		if isNull {
 			return errors.Errorf("invalid watch text expression")
 		}
-		record.Watch = rmpb.RunawayWatchType(op.IntValue)
-		if op.BoolValue {
+		watchType := textOption.Type
+		record.Watch = rmpb.RunawayWatchType(watchType)
+		if textOption.TypeSpecified {
 			p := parser.New()
 			stmts, _, err := p.ParseSQL(strval)
 			if err != nil {
@@ -83,7 +86,7 @@ func setWatchOption(ctx context.Context,
 				return errors.Errorf("only support one SQL")
 			}
 			sql := stmts[0].Text()
-			switch model.RunawayWatchType(op.IntValue) {
+			switch watchType {
 			case model.WatchNone:
 				return errors.Errorf("watch type must be specified")
 			case model.WatchExact:

pkg/parser/ast/ddl.go

@@ -2273,39 +2273,82 @@ const (
 
 // ResourceGroupRunawayOption is used for parsing resource group runaway rule option.
 type ResourceGroupRunawayOption struct {
-	Tp       RunawayOptionType
-	StrValue string
-	IntValue int32
+	Tp           RunawayOptionType
+	RuleOption   *ResourceGroupRunawayRuleOption
+	ActionOption *ResourceGroupRunawayActionOption
+	WatchOption  *ResourceGroupRunawayWatchOption
 }
 
 func (n *ResourceGroupRunawayOption) Restore(ctx *format.RestoreCtx) error {
 	switch n.Tp {
 	case RunawayRule:
-		ctx.WriteKeyWord("EXEC_ELAPSED ")
-		ctx.WritePlain("= ")
-		ctx.WriteString(n.StrValue)
+		n.RuleOption.restore(ctx)
 	case RunawayAction:
-		ctx.WriteKeyWord("ACTION ")
-		ctx.WritePlain("= ")
-		ctx.WriteKeyWord(model.RunawayActionType(n.IntValue).String())
+		n.ActionOption.Restore(ctx)
 	case RunawayWatch:
-		ctx.WriteKeyWord("WATCH ")
-		ctx.WritePlain("= ")
-		ctx.WriteKeyWord(model.RunawayWatchType(n.IntValue).String())
-		ctx.WritePlain(" ")
-		ctx.WriteKeyWord("DURATION ")
-		ctx.WritePlain("= ")
-		if len(n.StrValue) > 0 {
-			ctx.WriteString(n.StrValue)
-		} else {
-			ctx.WriteKeyWord("UNLIMITED")
-		}
+		n.WatchOption.restore(ctx)
 	default:
 		return errors.Errorf("invalid ResourceGroupRunawayOption: %d", n.Tp)
 	}
 	return nil
 }
 
+// ResourceGroupRunawayRuleOption is used for parsing the resource group/query watch runaway rule.
+type ResourceGroupRunawayRuleOption struct {
+	ExecElapsed string
+}
+
+func (n *ResourceGroupRunawayRuleOption) restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("EXEC_ELAPSED ")
+	ctx.WritePlain("= ")
+	ctx.WriteString(n.ExecElapsed)
+	return nil
+}
+
+// ResourceGroupRunawayActionOption is used for parsing the resource group runaway action.
+type ResourceGroupRunawayActionOption struct {
+	node
+	Type model.RunawayActionType
+}
+
+// Restore implements Node interface.
+func (n *ResourceGroupRunawayActionOption) Restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("ACTION ")
+	ctx.WritePlain("= ")
+	ctx.WriteKeyWord(n.Type.String())
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *ResourceGroupRunawayActionOption) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	return v.Leave(n)
+}
+
+// ResourceGroupRunawayWatchOption is used for parsing the resource group runaway watch.
+type ResourceGroupRunawayWatchOption struct {
+	Type     model.RunawayWatchType
+	Duration string
+}
+
+func (n *ResourceGroupRunawayWatchOption) restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("WATCH ")
+	ctx.WritePlain("= ")
+	ctx.WriteKeyWord(n.Type.String())
+	ctx.WritePlain(" ")
+	ctx.WriteKeyWord("DURATION ")
+	ctx.WritePlain("= ")
+	if len(n.Duration) > 0 {
+		ctx.WriteString(n.Duration)
+	} else {
+		ctx.WriteKeyWord("UNLIMITED")
+	}
+	return nil
+}
+
 type BackgroundOptionType int
 
 const (

pkg/parser/ast/ddl_test.go

@@ -933,3 +933,55 @@ func TestTableOptionTTLRestoreWithTTLEnableOffFlag(t *testing.T) {
 		runNodeRestoreTestWithFlagsStmtChange(t, testCases, "%s", extractNodeFunc, ca.flags)
 	}
 }
+
+func TestResourceGroupDDLStmtRestore(t *testing.T) {
+	createTestCases := []NodeRestoreTestCase{
+		{
+			"CREATE RESOURCE GROUP IF NOT EXISTS rg1 RU_PER_SEC = 500 BURSTABLE",
+			"CREATE RESOURCE GROUP IF NOT EXISTS `rg1` RU_PER_SEC = 500, BURSTABLE = TRUE",
+		},
+		{
+			"CREATE RESOURCE GROUP IF NOT EXISTS rg2 RU_PER_SEC = 600",
+			"CREATE RESOURCE GROUP IF NOT EXISTS `rg2` RU_PER_SEC = 600",
+		},
+		{
+			"CREATE RESOURCE GROUP IF NOT EXISTS rg3 RU_PER_SEC = 100 PRIORITY = HIGH",
+			"CREATE RESOURCE GROUP IF NOT EXISTS `rg3` RU_PER_SEC = 100, PRIORITY = HIGH",
+		},
+		{
+			"CREATE RESOURCE GROUP IF NOT EXISTS rg1 RU_PER_SEC = 500 QUERY_LIMIT=(EXEC_ELAPSED='60s', ACTION=COOLDOWN)",
+			"CREATE RESOURCE GROUP IF NOT EXISTS `rg1` RU_PER_SEC = 500, QUERY_LIMIT = (EXEC_ELAPSED = '60s' ACTION = COOLDOWN)",
+		},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*CreateResourceGroupStmt)
+	}
+	runNodeRestoreTest(t, createTestCases, "%s", extractNodeFunc)
+
+	alterTestCase := []NodeRestoreTestCase{
+		{
+			"ALTER RESOURCE GROUP rg1 QUERY_LIMIT=(EXEC_ELAPSED='60s', ACTION=KILL, WATCH=SIMILAR DURATION='10m')",
+			"ALTER RESOURCE GROUP `rg1` QUERY_LIMIT = (EXEC_ELAPSED = '60s' ACTION = KILL WATCH = SIMILAR DURATION = '10m')",
+		},
+		{
+			"ALTER RESOURCE GROUP rg1 QUERY_LIMIT=NULL",
+			"ALTER RESOURCE GROUP `rg1` QUERY_LIMIT = NULL",
+		},
+		{
+			"ALTER RESOURCE GROUP `default` BACKGROUND=(TASK_TYPES='br,ddl')",
+			"ALTER RESOURCE GROUP `default` BACKGROUND = (TASK_TYPES = 'br,ddl')",
+		},
+		{
+			"ALTER RESOURCE GROUP `default` BACKGROUND=NULL",
+			"ALTER RESOURCE GROUP `default` BACKGROUND = NULL",
+		},
+		{
+			"ALTER RESOURCE GROUP `default` BACKGROUND=(TASK_TYPES='')",
+			"ALTER RESOURCE GROUP `default` BACKGROUND = (TASK_TYPES = '')",
+		},
+	}
+	extractNodeFunc = func(node Node) Node {
+		return node.(*AlterResourceGroupStmt)
+	}
+	runNodeRestoreTest(t, alterTestCase, "%s", extractNodeFunc)
+}