*: rewrite the logic of set. (pingcap#2223)

ast/misc.go

@@ -264,7 +264,7 @@ type VariableAssignment struct {
 	// VariableAssignment should be able to store information for SetCharset/SetPWD Stmt.
 	// For SetCharsetStmt, Value is charset, ExtendValue is collation.
 	// TODO: Use SetStmt to implement set password statement.
-	ExtendValue ExprNode
+	ExtendValue *ValueExpr
 }
 
 // Accept implements Node interface.

executor/builder.go

@@ -78,6 +78,8 @@ func (b *executorBuilder) build(p plan.Plan) Executor {
 		return b.buildShow(v)
 	case *plan.Simple:
 		return b.buildSimple(v)
+	case *plan.Set:
+		return b.buildSet(v)
 	case *plan.Sort:
 		return b.buildSort(v)
 	case *plan.Union:
@@ -217,16 +219,14 @@ func (b *executorBuilder) buildSimple(v *plan.Simple) Executor {
 	switch s := v.Statement.(type) {
 	case *ast.GrantStmt:
 		return b.buildGrant(s)
-	case *ast.SetStmt:
-		return b.buildSet(s)
 	}
 	return &SimpleExec{Statement: v.Statement, ctx: b.ctx}
 }
 
-func (b *executorBuilder) buildSet(v *ast.SetStmt) Executor {
+func (b *executorBuilder) buildSet(v *plan.Set) Executor {
 	return &SetExecutor{
 		ctx:  b.ctx,
-		stmt: v,
+		vars: v.VarAssigns,
 	}
 }
 

executor/executor_set.go

@@ -21,7 +21,6 @@ import (
 	"github.com/ngaut/log"
 	"github.com/pingcap/tidb/ast"
 	"github.com/pingcap/tidb/context"
-	"github.com/pingcap/tidb/evaluator"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
@@ -32,7 +31,7 @@ import (
 
 // SetExecutor executes set statement.
 type SetExecutor struct {
-	stmt *ast.SetStmt
+	vars []*expression.VarAssignment
 	ctx  context.Context
 	done bool
 }
@@ -52,14 +51,14 @@ func (e *SetExecutor) Next() (*Row, error) {
 
 func (e *SetExecutor) executeSet() error {
 	sessionVars := e.ctx.GetSessionVars()
-	for _, v := range e.stmt.Variables {
+	for _, v := range e.vars {
 		// Variable is case insensitive, we use lower case.
 		if v.Name == ast.SetNames {
 			// This is set charset stmt.
-			cs := v.Value.GetValue().(string)
+			cs := v.Expr.(*expression.Constant).Value.GetString()
 			var co string
 			if v.ExtendValue != nil {
-				co = v.ExtendValue.GetValue().(string)
+				co = v.ExtendValue.Value.GetString()
 			}
 			err := e.setCharset(cs, co)
 			if err != nil {
@@ -70,7 +69,7 @@ func (e *SetExecutor) executeSet() error {
 		name := strings.ToLower(v.Name)
 		if !v.IsSystem {
 			// Set user variable.
-			value, err := evaluator.Eval(e.ctx, v.Value)
+			value, err := v.Expr.Eval(nil, e.ctx)
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -161,9 +160,8 @@ func (e *SetExecutor) setCharset(cs, co string) error {
 	return nil
 }
 
-func (e *SetExecutor) getVarValue(v *ast.VariableAssignment, sysVar *variable.SysVar) (value types.Datum, err error) {
-	switch v.Value.(type) {
-	case *ast.DefaultExpr:
+func (e *SetExecutor) getVarValue(v *expression.VarAssignment, sysVar *variable.SysVar) (value types.Datum, err error) {
+	if v.IsDefault {
 		// To set a SESSION variable to the GLOBAL value or a GLOBAL value
 		// to the compiled-in MySQL default value, use the DEFAULT keyword.
 		// See http://dev.mysql.com/doc/refman/5.7/en/set-statement.html
@@ -176,9 +174,9 @@ func (e *SetExecutor) getVarValue(v *ast.VariableAssignment, sysVar *variable.Sy
 			}
 			value = types.NewStringDatum(s)
 		}
-	default:
-		value, err = evaluator.Eval(e.ctx, v.Value)
+		return
 	}
+	value, err = v.Expr.Eval(nil, e.ctx)
 	return value, errors.Trace(err)
 }
 

expression/expression.go

@@ -172,6 +172,16 @@ type Assignment struct {
 	Expr Expression
 }
 
+// VarAssignment represents a variable assignment in Set, such as set global a = 1.
+type VarAssignment struct {
+	Name        string
+	Expr        Expression
+	IsDefault   bool
+	IsGlobal    bool
+	IsSystem    bool
+	ExtendValue *Constant
+}
+
 // splitNormalFormItems split CNF(conjunctive normal form) like "a and b and c", or DNF(disjunctive normal form) like "a or b or c"
 func splitNormalFormItems(onExpr Expression, funcName string) []Expression {
 	switch v := onExpr.(type) {

plan/plan.go

@@ -28,6 +28,8 @@ import (
 const (
 	// Sel is the type of Selection.
 	Sel = "Selection"
+	// St is the type of Set.
+	St = "Set"
 	// Proj is the type of Projection.
 	Proj = "Projection"
 	// Agg is the type of Aggregation.

plan/planbuilder.go

@@ -95,7 +95,9 @@ func (b *planBuilder) build(node ast.Node) Plan {
 		return b.buildShow(x)
 	case *ast.DoStmt:
 		return b.buildDo(x)
-	case *ast.AnalyzeTableStmt, *ast.BinlogStmt, *ast.FlushTableStmt, *ast.UseStmt, *ast.SetStmt,
+	case *ast.SetStmt:
+		return b.buildSet(x)
+	case *ast.AnalyzeTableStmt, *ast.BinlogStmt, *ast.FlushTableStmt, *ast.UseStmt,
 		*ast.BeginStmt, *ast.CommitStmt, *ast.RollbackStmt, *ast.CreateUserStmt, *ast.SetPwdStmt,
 		*ast.GrantStmt, *ast.DropUserStmt, *ast.AlterUserStmt:
 		return b.buildSimple(node.(ast.StmtNode))
@@ -143,6 +145,36 @@ func (b *planBuilder) buildDo(v *ast.DoStmt) Plan {
 	return p
 }
 
+func (b *planBuilder) buildSet(v *ast.SetStmt) Plan {
+	p := &Set{}
+	p.tp = St
+	p.allocator = b.allocator
+	for _, vars := range v.Variables {
+		assign := &expression.VarAssignment{
+			Name:     vars.Name,
+			IsGlobal: vars.IsGlobal,
+			IsSystem: vars.IsSystem,
+		}
+		if _, ok := vars.Value.(*ast.DefaultExpr); !ok {
+			assign.Expr, _, b.err = b.rewrite(vars.Value, nil, nil, true)
+			if b.err != nil {
+				return nil
+			}
+		} else {
+			assign.IsDefault = true
+		}
+		if vars.ExtendValue != nil {
+			assign.ExtendValue = &expression.Constant{
+				Value:   vars.ExtendValue.Datum,
+				RetType: &vars.ExtendValue.Type,
+			}
+		}
+		p.VarAssigns = append(p.VarAssigns, assign)
+	}
+	p.initIDAndContext(b.ctx)
+	return p
+}
+
 // Detect aggregate function or groupby clause.
 func (b *planBuilder) detectSelectAgg(sel *ast.SelectStmt) bool {
 	if sel.GroupBy != nil {

plan/plans.go

@@ -158,6 +158,13 @@ type Show struct {
 	GlobalScope bool
 }
 
+// Set represents a plan for set stmt.
+type Set struct {
+	basePlan
+
+	VarAssigns []*expression.VarAssignment
+}
+
 // Simple represents a simple statement plan which doesn't need any optimization.
 type Simple struct {
 	basePlan