parser: support new sql syntax to distribute table regions

pkg/parser/ast/dml.go

@@ -24,6 +24,7 @@ import (
 
 var (
 	_ DMLNode = &DeleteStmt{}
+	_ DMLNode = &DistributeTableStmt{}
 	_ DMLNode = &InsertStmt{}
 	_ DMLNode = &SetOprStmt{}
 	_ DMLNode = &UpdateStmt{}
@@ -3052,6 +3053,7 @@ const (
 	ShowReplicaStatus
 	ShowDistributions
 	ShowPlanForSQL
+	ShowDistributionJobs
 )
 
 const (
@@ -3102,6 +3104,8 @@ type ShowStmt struct {
 
 	ImportJobID *int64 // Used for `SHOW IMPORT JOB <ID>` syntax
 	SQLOrDigest string // Used for `SHOW PLAN FOR ...` syntax
+
+	DistributionJobID *int64 // Used for `SHOW DISTRIBUTION JOB <ID>` syntax
 }
 
 // Restore implements Node interface.
@@ -3320,6 +3324,14 @@ func (n *ShowStmt) Restore(ctx *format.RestoreCtx) error {
 			ctx.WriteKeyWord("IMPORT JOBS")
 			restoreShowLikeOrWhereOpt()
 		}
+	case ShowDistributionJobs:
+		if n.DistributionJobID != nil {
+			ctx.WriteKeyWord("DISTRIBUTION JOB ")
+			ctx.WritePlainf("%d", *n.DistributionJobID)
+		} else {
+			ctx.WriteKeyWord("DISTRIBUTION JOBS")
+			restoreShowLikeOrWhereOpt()
+		}
 	// ShowTargetFilterable
 	default:
 		switch n.Tp {
@@ -3842,6 +3854,62 @@ func (n *FrameBound) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+type DistributeTableStmt struct {
+	dmlNode
+	Table          *TableName
+	PartitionNames []CIStr
+	Rule           CIStr
+	Engine         CIStr
+}
+
+// Restore implements Node interface.
+func (n *DistributeTableStmt) Restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("DISTRIBUTE ")
+	ctx.WriteKeyWord("TABLE ")
+
+	if err := n.Table.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore SplitIndexRegionStmt.Table")
+	}
+	if len(n.PartitionNames) > 0 {
+		ctx.WriteKeyWord(" PARTITION")
+		ctx.WritePlain("(")
+		for i, v := range n.PartitionNames {
+			if i != 0 {
+				ctx.WritePlain(", ")
+			}
+			ctx.WriteName(v.String())
+		}
+		ctx.WritePlain(")")
+	}
+
+	if len(n.Rule.L) > 0 {
+		ctx.WriteKeyWord(" RULE = ")
+		ctx.WriteName(n.Rule.String())
+	}
+
+	if len(n.Engine.L) > 0 {
+		ctx.WriteKeyWord(" ENGINE = ")
+		ctx.WriteName(n.Engine.String())
+	}
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *DistributeTableStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+
+	n = newNode.(*DistributeTableStmt)
+	node, ok := n.Table.Accept(v)
+	if !ok {
+		return n, false
+	}
+	n.Table = node.(*TableName)
+	return v.Leave(n)
+}
+
 type SplitRegionStmt struct {
 	dmlNode
 

pkg/parser/keywords_test.go

@@ -36,7 +36,7 @@ func TestKeywords(t *testing.T) {
 }
 
 func TestKeywordsLength(t *testing.T) {
-	require.Equal(t, 657, len(parser.Keywords))
+	require.Equal(t, 660, len(parser.Keywords))
 
 	reservedNr := 0
 	for _, kw := range parser.Keywords {

pkg/parser/misc.go

@@ -343,6 +343,8 @@ var tokenMap = map[string]int{
 	"DISK":                     disk,
 	"DISTINCT":                 distinct,
 	"DISTINCTROW":              distinct,
+	"DISTRIBUTE":               distribute,
+	"DISTRIBUTION":             distribution,
 	"DISTRIBUTIONS":            distributions,
 	"DIV":                      div,
 	"DO":                       do,
@@ -679,6 +681,7 @@ var tokenMap = map[string]int{
 	"REPLICAS":                 replicas,
 	"REPLICATION":              replication,
 	"RU":                       ru,
+	"RULE":                     rule,
 	"REQUIRE":                  require,
 	"REQUIRED":                 required,
 	"RESET":                    reset,