Support SWITCH_GROUP syntax for runaway watch

Signed-off-by: JmPotato <[email protected]>
pingcap · Jul 22, 2024 · ce699f9 · ce699f9
1 parent 669bf4d
commit ce699f9
Show file tree

Hide file tree

Showing 6 changed files with 11,049 additions and 10,989 deletions.
diff --git a/pkg/parser/ast/ddl.go b/pkg/parser/ast/ddl.go
@@ -2276,6 +2276,8 @@ type ResourceGroupRunawayOption struct {
 	Tp       RunawayOptionType
 	StrValue string
 	IntValue int32
+
+	WrappedAction *WrappedQueryWatchAction
 }
 
 func (n *ResourceGroupRunawayOption) Restore(ctx *format.RestoreCtx) error {
@@ -2285,9 +2287,9 @@ func (n *ResourceGroupRunawayOption) Restore(ctx *format.RestoreCtx) error {
 		ctx.WritePlain("= ")
 		ctx.WriteString(n.StrValue)
 	case RunawayAction:
-		ctx.WriteKeyWord("ACTION ")
-		ctx.WritePlain("= ")
-		ctx.WriteKeyWord(model.RunawayActionType(n.IntValue).String())
+		if err := n.WrappedAction.Restore(ctx); err != nil {
+			return err
+		}
 	case RunawayWatch:
 		ctx.WriteKeyWord("WATCH ")
 		ctx.WritePlain("= ")

diff --git a/pkg/parser/ast/misc.go b/pkg/parser/ast/misc.go
@@ -4091,6 +4091,8 @@ type QueryWatchOption struct {
 	IntValue  int32
 	ExprValue ExprNode
 	BoolValue bool
+
+	WrappedAction *WrappedQueryWatchAction
 }
 
 func (n *QueryWatchOption) Restore(ctx *format.RestoreCtx) error {
@@ -4105,9 +4107,9 @@ func (n *QueryWatchOption) Restore(ctx *format.RestoreCtx) error {
 			ctx.WriteName(n.StrValue.O)
 		}
 	case QueryWatchAction:
-		ctx.WriteKeyWord("ACTION ")
-		ctx.WritePlain("= ")
-		ctx.WriteKeyWord(model.RunawayActionType(n.IntValue).String())
+		if err := n.WrappedAction.Restore(ctx); err != nil {
+			return err
+		}
 	case QueryWatchType:
 		if n.BoolValue {
 			ctx.WriteKeyWord("SQL TEXT ")
@@ -4153,3 +4155,28 @@ func CheckQueryWatchAppend(ops []*QueryWatchOption, newOp *QueryWatchOption) boo
 	}
 	return true
 }
+
+// WrappedQueryWatchAction is a wrapper for the query watch action.
+type WrappedQueryWatchAction struct {
+	Tp          model.RunawayActionType
+	SwitchGroup model.CIStr
+}
+
+func (n *WrappedQueryWatchAction) Restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("ACTION ")
+	ctx.WritePlain("= ")
+	switch n.Tp {
+	case model.RunawayActionNone, model.RunawayActionDryRun, model.RunawayActionCooldown, model.RunawayActionKill:
+		ctx.WriteKeyWord(n.Tp.String())
+	case model.RunawayActionSwitchGroup:
+		switchGroup := n.SwitchGroup.String()
+		if len(switchGroup) == 0 {
+			return errors.New("SWITCH_GROUP runaway watch action requires a non-empty group name")
+		}
+		ctx.WriteKeyWord("SWITCH_GROUP")
+		ctx.WritePlain("(")
+		ctx.WriteName(switchGroup)
+		ctx.WritePlain(")")
+	}
+	return nil
+}
diff --git a/pkg/parser/keywords.go b/pkg/parser/keywords.go
@@ -622,6 +622,7 @@ var Keywords = []KeywordsType{
 	{"VALIDATION", false, "unreserved"},
 	{"VALUE", false, "unreserved"},
 	{"VARIABLES", false, "unreserved"},
+	{"VECTOR", false, "unreserved"},
 	{"VIEW", false, "unreserved"},
 	{"VISIBLE", false, "unreserved"},
 	{"WAIT", false, "unreserved"},

diff --git a/pkg/parser/model/model.go b/pkg/parser/model/model.go
@@ -1951,6 +1951,7 @@ const (
 	RunawayActionDryRun
 	RunawayActionCooldown
 	RunawayActionKill
+	RunawayActionSwitchGroup
 )
 
 // RunawayWatchType is the type of runaway watch.
@@ -1995,6 +1996,8 @@ func (t RunawayActionType) String() string {
 		return "COOLDOWN"
 	case RunawayActionKill:
 		return "KILL"
+	case RunawayActionSwitchGroup:
+		return "SWITCH_GROUP"
 	default:
 		return "DRYRUN"
 	}