parser: Refine LoadDataStmt (#288)

* Code Refine * Add multi set test case * revert test case loss
pingcap · May 1, 2019 · 366172d · 366172d
1 parent 3de53cf
commit 366172d
Show file tree

Hide file tree

Showing 4 changed files with 4,252 additions and 4,103 deletions.
diff --git a/ast/dml.go b/ast/dml.go
@@ -1113,18 +1113,26 @@ func (n *Assignment) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+type ColumnNameOrUserVar struct {
+	ColumnName *ColumnName
+	UserVar    *VariableExpr
+}
+
 // LoadDataStmt is a statement to load data from a specified file, then insert this rows into an existing table.
 // See https://dev.mysql.com/doc/refman/5.7/en/load-data.html
 type LoadDataStmt struct {
 	dmlNode
 
-	IsLocal     bool
-	Path        string
-	Table       *TableName
-	Columns     []*ColumnName
-	FieldsInfo  *FieldsClause
-	LinesInfo   *LinesClause
-	IgnoreLines uint64
+	IsLocal           bool
+	Path              string
+	Table             *TableName
+	Columns           []*ColumnName
+	FieldsInfo        *FieldsClause
+	LinesInfo         *LinesClause
+	IgnoreLines       uint64
+	ColumnAssignments []*Assignment
+
+	ColumnsAndUserVars []*ColumnNameOrUserVar
 }
 
 // Restore implements Node interface.
@@ -1146,18 +1154,39 @@ func (n *LoadDataStmt) Restore(ctx *RestoreCtx) error {
 		ctx.WritePlainf("%d", n.IgnoreLines)
 		ctx.WriteKeyWord(" LINES")
 	}
-	if len(n.Columns) != 0 {
+	if len(n.ColumnsAndUserVars) != 0 {
 		ctx.WritePlain(" (")
-		for i, column := range n.Columns {
+		for i, c := range n.ColumnsAndUserVars {
 			if i != 0 {
 				ctx.WritePlain(",")
 			}
-			if err := column.Restore(ctx); err != nil {
-				return errors.Annotate(err, "An error occurred while restore LoadDataStmt.Columns")
+			if c.ColumnName != nil {
+				if err := c.ColumnName.Restore(ctx); err != nil {
+					return errors.Annotate(err, "An error occurred while restore LoadDataStmt.ColumnsAndUserVars")
+				}
 			}
+			if c.UserVar != nil {
+				if err := c.UserVar.Restore(ctx); err != nil {
+					return errors.Annotate(err, "An error occurred while restore LoadDataStmt.ColumnsAndUserVars")
+				}
+			}
+
 		}
 		ctx.WritePlain(")")
 	}
+
+	if n.ColumnAssignments != nil {
+		ctx.WriteKeyWord(" SET")
+		for i, assign := range n.ColumnAssignments {
+			if i != 0 {
+				ctx.WritePlain(",")
+			}
+			ctx.WritePlain(" ")
+			if err := assign.Restore(ctx); err != nil {
+				return errors.Annotate(err, "An error occurred while restore LoadDataStmt.ColumnAssignments")
+			}
+		}
+	}
 	return nil
 }
 
@@ -1182,6 +1211,14 @@ func (n *LoadDataStmt) Accept(v Visitor) (Node, bool) {
 		}
 		n.Columns[i] = node.(*ColumnName)
 	}
+
+	for i, assignment := range n.ColumnAssignments {
+		node, ok := assignment.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.ColumnAssignments[i] = node.(*Assignment)
+	}
 	return v.Leave(n)
 }