parser: add traffic capture/replay statements in the parser (pingcap#…

…57808) close pingcap#57807
joccau · Dec 5, 2024 · 30069c5 · 30069c5
1 parent a4278e0
commit 30069c5
Show file tree

Hide file tree

Showing 5 changed files with 11,385 additions and 10,903 deletions.
diff --git a/pkg/parser/ast/misc.go b/pkg/parser/ast/misc.go
@@ -61,6 +61,7 @@ var (
 	_ StmtNode = &PlanReplayerStmt{}
 	_ StmtNode = &CompactTableStmt{}
 	_ StmtNode = &SetResourceGroupStmt{}
+	_ StmtNode = &TrafficStmt{}
 
 	_ Node = &PrivElem{}
 	_ Node = &VariableAssignment{}
@@ -415,6 +416,112 @@ func (n *PlanReplayerStmt) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+// TrafficOpType is traffic operation type.
+type TrafficOpType int
+
+const (
+	TrafficOpCapture TrafficOpType = iota
+	TrafficOpReplay
+	TrafficOpShow
+	TrafficOpCancel
+)
+
+// TrafficOptionType is traffic option type.
+type TrafficOptionType int
+
+const (
+	// capture options
+	TrafficOptionDuration TrafficOptionType = iota
+	TrafficOptionEncryptionMethod
+	TrafficOptionCompress
+	// replay options
+	TrafficOptionUsername
+	TrafficOptionPassword
+	TrafficOptionSpeed
+	TrafficOptionReadOnly
+)
+
+// TrafficStmt is traffic operation statement.
+type TrafficStmt struct {
+	stmtNode
+	OpType  TrafficOpType
+	Options []*TrafficOption
+	Dir     string
+}
+
+// TrafficOption is traffic option.
+type TrafficOption struct {
+	OptionType TrafficOptionType
+	FloatValue ValueExpr
+	StrValue   string
+	BoolValue  bool
+}
+
+// Restore implements Node interface.
+func (n *TrafficStmt) Restore(ctx *format.RestoreCtx) error {
+	switch n.OpType {
+	case TrafficOpCapture:
+		ctx.WriteKeyWord("TRAFFIC CAPTURE TO ")
+		ctx.WriteString(n.Dir)
+		for _, option := range n.Options {
+			ctx.WritePlain(" ")
+			switch option.OptionType {
+			case TrafficOptionDuration:
+				ctx.WriteKeyWord("DURATION ")
+				ctx.WritePlain("= ")
+				ctx.WriteString(option.StrValue)
+			case TrafficOptionEncryptionMethod:
+				ctx.WriteKeyWord("ENCRYPTION_METHOD ")
+				ctx.WritePlain("= ")
+				ctx.WriteString(option.StrValue)
+			case TrafficOptionCompress:
+				ctx.WriteKeyWord("COMPRESS ")
+				ctx.WritePlain("= ")
+				ctx.WritePlain(strings.ToUpper(fmt.Sprintf("%v", option.BoolValue)))
+			}
+		}
+	case TrafficOpReplay:
+		ctx.WriteKeyWord("TRAFFIC REPLAY FROM ")
+		ctx.WriteString(n.Dir)
+		for _, option := range n.Options {
+			ctx.WritePlain(" ")
+			switch option.OptionType {
+			case TrafficOptionUsername:
+				ctx.WriteKeyWord("USER ")
+				ctx.WritePlain("= ")
+				ctx.WriteString(option.StrValue)
+			case TrafficOptionPassword:
+				ctx.WriteKeyWord("PASSWORD ")
+				ctx.WritePlain("= ")
+				ctx.WriteString(option.StrValue)
+			case TrafficOptionSpeed:
+				ctx.WriteKeyWord("SPEED ")
+				ctx.WritePlain("= ")
+				ctx.WritePlainf("%v", option.FloatValue.GetValue())
+			case TrafficOptionReadOnly:
+				ctx.WriteKeyWord("READONLY ")
+				ctx.WritePlain("= ")
+				ctx.WritePlain(strings.ToUpper(fmt.Sprintf("%v", option.BoolValue)))
+			}
+		}
+	case TrafficOpShow:
+		ctx.WriteKeyWord("SHOW TRAFFIC JOBS")
+	case TrafficOpCancel:
+		ctx.WriteKeyWord("CANCEL TRAFFIC JOBS")
+	}
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *TrafficStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	n = newNode.(*TrafficStmt)
+	return v.Leave(n)
+}
+
 type CompactReplicaKind string
 
 const (

diff --git a/pkg/parser/misc.go b/pkg/parser/misc.go
@@ -262,6 +262,7 @@ var tokenMap = map[string]int{
 	"COMMIT":                   commit,
 	"COMMITTED":                committed,
 	"COMPACT":                  compact,
+	"COMPRESS":                 compress,
 	"COMPRESSED":               compressed,
 	"COMPRESSION":              compression,
 	"CONCURRENCY":              concurrency,
@@ -647,6 +648,7 @@ var tokenMap = map[string]int{
 	"RANGE":                    rangeKwd,
 	"RATE_LIMIT":               rateLimit,
 	"READ":                     read,
+	"READ_ONLY":                readOnly,
 	"REAL":                     realType,
 	"REBUILD":                  rebuild,
 	"RECENT":                   recent,
@@ -667,6 +669,7 @@ var tokenMap = map[string]int{
 	"REPEAT":                   repeat,
 	"REPEATABLE":               repeatable,
 	"REPLACE":                  replace,
+	"REPLAY":                   replay,
 	"REPLAYER":                 replayer,
 	"REPLICA":                  replica,
 	"REPLICAS":                 replicas,
@@ -741,6 +744,7 @@ var tokenMap = map[string]int{
 	"SOME":                     some,
 	"SOURCE":                   source,
 	"SPATIAL":                  spatial,
+	"SPEED":                    speed,
 	"SPLIT":                    split,
 	"SQL_BIG_RESULT":           sqlBigResult,
 	"SQL_BUFFER_RESULT":        sqlBufferResult,
@@ -847,6 +851,7 @@ var tokenMap = map[string]int{
 	"TPCC":                     tpcc,
 	"TRACE":                    trace,
 	"TRADITIONAL":              traditional,
+	"TRAFFIC":                  traffic,
 	"TRAILING":                 trailing,
 	"TRANSACTION":              transaction,
 	"TRIGGER":                  trigger,