support SELECT FROM TABLESAMPLE syntax

ast/ddl.go

@@ -330,13 +330,15 @@ func (n *ReferenceDef) Accept(v Visitor) (Node, bool) {
 		return v.Leave(newNode)
 	}
 	n = newNode.(*ReferenceDef)
-	node, ok := n.Table.Accept(v)
-	if !ok {
-		return n, false
+	if n.Table != nil {
+		node, ok := n.Table.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.Table = node.(*TableName)
 	}
-	n.Table = node.(*TableName)
 	for i, val := range n.IndexPartSpecifications {
-		node, ok = val.Accept(v)
+		node, ok := val.Accept(v)
 		if !ok {
 			return n, false
 		}

ast/dml.go

@@ -184,6 +184,7 @@ type TableName struct {
 
 	IndexHints     []*IndexHint
 	PartitionNames []model.CIStr
+	TableSample    *TableSample
 }
 
 // Restore implements Node interface.
@@ -216,14 +217,25 @@ func (n *TableName) restoreIndexHints(ctx *format.RestoreCtx) error {
 			return errors.Annotate(err, "An error occurred while splicing IndexHints")
 		}
 	}
+	return nil
+}
 
+func (n *TableName) restoreTableSample(ctx *format.RestoreCtx) error {
+	if n.TableSample != nil {
+		if err := n.TableSample.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while splicing TableName.TableSample")
+		}
+	}
 	return nil
 }
 
 func (n *TableName) Restore(ctx *format.RestoreCtx) error {
 	n.restoreName(ctx)
 	n.restorePartitions(ctx)
-	return n.restoreIndexHints(ctx)
+	if err := n.restoreIndexHints(ctx); err != nil {
+		return err
+	}
+	return n.restoreTableSample(ctx)
 }
 
 // IndexHintType is the type for index hint use, ignore or force.
@@ -302,6 +314,13 @@ func (n *TableName) Accept(v Visitor) (Node, bool) {
 		return v.Leave(newNode)
 	}
 	n = newNode.(*TableName)
+	if n.TableSample != nil {
+		newTs, ok := n.TableSample.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.TableSample = newTs.(*TableSample)
+	}
 	return v.Leave(n)
 }
 
@@ -409,6 +428,9 @@ func (n *TableSource) Restore(ctx *format.RestoreCtx) error {
 		if err := tn.restoreIndexHints(ctx); err != nil {
 			return errors.Annotate(err, "An error occurred while restore TableSource.Source.(*TableName).IndexHints")
 		}
+		if err := tn.restoreTableSample(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore TableSource.Source.(*TableName).TableSample")
+		}
 
 		if needParen {
 			ctx.WritePlain(")")
@@ -775,6 +797,90 @@ func (n *OrderByClause) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+type SampleMethodType int8
+
+const (
+	SampleMethodTypeNone SampleMethodType = iota
+	SampleMethodTypeSystem
+	SampleMethodTypeBernoulli
+	SampleMethodTypeTiDBRegion
+)
+
+type SampleClauseUnitType int8
+
+const (
+	SampleClauseUnitTypeDefault SampleClauseUnitType = iota
+	SampleClauseUnitTypeRow
+	SampleClauseUnitTypePercent
+)
+
+type TableSample struct {
+	node
+	SampleMethod     SampleMethodType
+	Expr             ExprNode
+	SampleClauseUnit SampleClauseUnitType
+	RepeatableSeed   ExprNode
+}
+
+func (s *TableSample) Restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord(" TABLESAMPLE ")
+	switch s.SampleMethod {
+	case SampleMethodTypeBernoulli:
+		ctx.WriteKeyWord("BERNOULLI ")
+	case SampleMethodTypeSystem:
+		ctx.WriteKeyWord("SYSTEM ")
+	case SampleMethodTypeTiDBRegion:
+		ctx.WriteKeyWord("REGION ")
+	}
+	ctx.WritePlain("(")
+	if s.Expr != nil {
+		if err := s.Expr.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore TableSample.Expr")
+		}
+	}
+	switch s.SampleClauseUnit {
+	case SampleClauseUnitTypeDefault:
+	case SampleClauseUnitTypePercent:
+		ctx.WriteKeyWord(" PERCENT")
+	case SampleClauseUnitTypeRow:
+		ctx.WriteKeyWord(" ROWS")
+
+	}
+	ctx.WritePlain(")")
+	if s.RepeatableSeed != nil {
+		ctx.WriteKeyWord(" REPEATABLE")
+		ctx.WritePlain("(")
+		if err := s.RepeatableSeed.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore TableSample.Expr")
+		}
+		ctx.WritePlain(")")
+	}
+	return nil
+}
+
+func (s *TableSample) Accept(v Visitor) (node Node, ok bool) {
+	newNode, skipChildren := v.Enter(s)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	s = newNode.(*TableSample)
+	if s.Expr != nil {
+		node, ok = s.Expr.Accept(v)
+		if !ok {
+			return s, false
+		}
+		s.Expr = node.(ExprNode)
+	}
+	if s.RepeatableSeed != nil {
+		node, ok = s.RepeatableSeed.Accept(v)
+		if !ok {
+			return s, false
+		}
+		s.RepeatableSeed = node.(ExprNode)
+	}
+	return v.Leave(s)
+}
+
 type SelectStmtKind uint8
 
 const (
@@ -909,6 +1015,7 @@ func (n *SelectStmt) Restore(ctx *format.RestoreCtx) error {
 		if n.From == nil && n.Where != nil {
 			ctx.WriteKeyWord(" FROM DUAL")
 		}
+
 		if n.Where != nil {
 			ctx.WriteKeyWord(" WHERE ")
 			if err := n.Where.Restore(ctx); err != nil {

misc.go

@@ -174,6 +174,7 @@ var tokenMap = map[string]int{
 	"BACKUPS":                  backups,
 	"BEGIN":                    begin,
 	"BETWEEN":                  between,
+	"BERNOULLI":                bernoulli,
 	"BIGINT":                   bigIntType,
 	"BINARY":                   binaryType,
 	"BINDING":                  binding,
@@ -507,6 +508,7 @@ var tokenMap = map[string]int{
 	"PARTITIONING":             partitioning,
 	"PARTITIONS":               partitions,
 	"PASSWORD":                 password,
+	"PERCENT":                  percent,
 	"PER_DB":                   per_db,
 	"PER_TABLE":                per_table,
 	"PESSIMISTIC":              pessimistic,
@@ -656,10 +658,12 @@ var tokenMap = map[string]int{
 	"SUPER":                    super,
 	"SWAPS":                    swaps,
 	"SWITCHES":                 switchesSym,
+	"SYSTEM":                   system,
 	"SYSTEM_TIME":              systemTime,
 	"TABLE_CHECKSUM":           tableChecksum,
 	"TABLE":                    tableKwd,
 	"TABLES":                   tables,
+	"TABLESAMPLE":              tableSample,
 	"TABLESPACE":               tablespace,
 	"TELEMETRY":                telemetry,
 	"TELEMETRY_ID":             telemetryID,