Skip to content

Commit

Permalink
pkg/parser: use clearer structures for RESOURCE GROUP ASTs (#54881)
Browse files Browse the repository at this point in the history 
ref #54434
  • Loading branch information
@JmPotato
JmPotato authored Jul 25, 2024
1 parent 9897ddd commit 98b7858
Show file tree
Hide file tree
Showing 8 changed files with 355 additions and 102 deletions.
16 changes: 8 additions & 8 deletions pkg/ddl/executor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3723,7 +3723,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
}
}
Expand All @@ -3747,25 +3747,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
}
Expand Down
19 changes: 11 additions & 8 deletions pkg/executor/internal/querywatch/query_watch.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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
}
Expand All @@ -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 {
Expand All @@ -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:
Expand Down
83 changes: 63 additions & 20 deletions pkg/parser/ast/ddl.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
52 changes: 52 additions & 0 deletions pkg/parser/ast/ddl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Loading

0 comments on commit 98b7858

Please sign in to comment.