parser: support RETURNING clause syntax (#65633)

Author: bb7133

pkg/parser/ast/dml.go

@@ -871,6 +871,18 @@ func (n *FieldList) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+func restoreReturningFields(ctx *format.RestoreCtx, fields []*SelectField) error {
+	for i, field := range fields {
+		if i != 0 {
+			ctx.WritePlain(", ")
+		}
+		if err := field.Restore(ctx); err != nil {
+			return errors.Annotatef(err, "An error occurred while restore Returning[%d]", i)
+		}
+	}
+	return nil
+}
+
 // Accept implements Node Accept interface.
 func (n *FieldList) Accept(v Visitor) (Node, bool) {
 	newNode, skipChildren := v.Enter(n)
@@ -2401,6 +2413,8 @@ type InsertStmt struct {
 	// TableHints represents the table level Optimizer Hint for join type.
 	TableHints     []*TableOptimizerHint
 	PartitionNames []CIStr
+	// Returning represents the RETURNING select_expr list for INSERT statement.
+	Returning []*SelectField
 }
 
 // Restore implements Node interface.
@@ -2526,6 +2540,12 @@ func (n *InsertStmt) Restore(ctx *format.RestoreCtx) error {
 			}
 		}
 	}
+	if len(n.Returning) > 0 {
+		ctx.WriteKeyWord(" RETURNING ")
+		if err := restoreReturningFields(ctx, n.Returning); err != nil {
+			return errors.Annotate(err, "An error occurred while restore InsertStmt.Returning")
+		}
+	}
 
 	return nil
 }
@@ -2575,6 +2595,13 @@ func (n *InsertStmt) Accept(v Visitor) (Node, bool) {
 		}
 		n.OnDuplicate[i] = node.(*Assignment)
 	}
+	for i, field := range n.Returning {
+		node, ok := field.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.Returning[i] = node.(*SelectField)
+	}
 	return v.Leave(n)
 }
 
@@ -2634,6 +2661,8 @@ type DeleteStmt struct {
 	// TableHints represents the table level Optimizer Hint for join type.
 	TableHints []*TableOptimizerHint
 	With       *WithClause
+	// Returning represents the RETURNING select_expr list for DELETE statement.
+	Returning []*SelectField
 }
 
 // Restore implements Node interface.
@@ -2723,6 +2752,12 @@ func (n *DeleteStmt) Restore(ctx *format.RestoreCtx) error {
 			return errors.Annotate(err, "An error occurred while restore DeleteStmt.Limit")
 		}
 	}
+	if len(n.Returning) > 0 {
+		ctx.WriteKeyWord(" RETURNING ")
+		if err := restoreReturningFields(ctx, n.Returning); err != nil {
+			return errors.Annotate(err, "An error occurred while restore DeleteStmt.Returning")
+		}
+	}
 
 	return nil
 }
@@ -2777,6 +2812,13 @@ func (n *DeleteStmt) Accept(v Visitor) (Node, bool) {
 		}
 		n.Limit = node.(*Limit)
 	}
+	for i, field := range n.Returning {
+		node, ok = field.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.Returning[i] = node.(*SelectField)
+	}
 	return v.Leave(n)
 }
 
@@ -2886,6 +2928,8 @@ type UpdateStmt struct {
 	MultipleTable bool
 	TableHints    []*TableOptimizerHint
 	With          *WithClause
+	// Returning represents the RETURNING select_expr list for UPDATE statement.
+	Returning []*SelectField
 }
 
 // Restore implements Node interface.
@@ -2964,6 +3008,12 @@ func (n *UpdateStmt) Restore(ctx *format.RestoreCtx) error {
 			return errors.Annotate(err, "An error occur while restore UpdateStmt.Limit")
 		}
 	}
+	if len(n.Returning) > 0 {
+		ctx.WriteKeyWord(" RETURNING ")
+		if err := restoreReturningFields(ctx, n.Returning); err != nil {
+			return errors.Annotate(err, "An error occurred while restore UpdateStmt.Returning")
+		}
+	}
 
 	return nil
 }
@@ -3015,6 +3065,13 @@ func (n *UpdateStmt) Accept(v Visitor) (Node, bool) {
 		}
 		n.Limit = node.(*Limit)
 	}
+	for i, field := range n.Returning {
+		node, ok = field.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.Returning[i] = node.(*SelectField)
+	}
 	return v.Leave(n)
 }
 

pkg/parser/keywords.go

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

pkg/parser/keywords_test.go

@@ -709,6 +709,7 @@ var tokenMap = map[string]int{
 	"RESTORES":                       restores,
 	"RESTORED_TS":                    restoredTS,
 	"RESTRICT":                       restrict,
+	"RETURNING":                      returning,
 	"REVERSE":                        reverse,
 	"REVOKE":                         revoke,
 	"RIGHT":                          right,