*: Fix issue#1114 (pingcap#1231)

* *: supprot assignment * *: fix issue 1114
rjmcguire · May 13, 2016 · 2e355f3 · 2e355f3
1 parent dcec06d
commit 2e355f3
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 2 deletions.
diff --git a/ast/expressions.go b/ast/expressions.go
@@ -724,6 +724,8 @@ type VariableExpr struct {
 	IsGlobal bool
 	// IsSystem indicates whether this variable is a system variable in current session.
 	IsSystem bool
+	// Value is the variable value.
+	Value ExprNode
 }
 
 // Accept implements Node Accept interface.
@@ -733,5 +735,14 @@ func (n *VariableExpr) Accept(v Visitor) (Node, bool) {
 		return v.Leave(newNode)
 	}
 	n = newNode.(*VariableExpr)
+	if n.Value == nil {
+		return v.Leave(n)
+	}
+
+	node, ok := n.Value.Accept(v)
+	if !ok {
+		return n, false
+	}
+	n.Value = node.(ExprNode)
 	return v.Leave(n)
 }
diff --git a/ast/flag.go b/ast/flag.go
@@ -121,7 +121,11 @@ func (f *flagSetter) Leave(in Node) (Node, bool) {
 	case *ValuesExpr:
 		x.SetFlag(FlagHasReference)
 	case *VariableExpr:
-		x.SetFlag(FlagHasVariable)
+		if x.Value == nil {
+			x.SetFlag(FlagHasVariable)
+		} else {
+			x.SetFlag(FlagHasVariable | x.Value.GetFlag())
+		}
 	}
 
 	return in, true

diff --git a/evaluator/evaluator.go b/evaluator/evaluator.go
@@ -577,6 +577,16 @@ func (e *Evaluator) variable(v *ast.VariableExpr) bool {
 	sessionVars := variable.GetSessionVars(e.ctx)
 	globalVars := variable.GetGlobalVarAccessor(e.ctx)
 	if !v.IsSystem {
+		if v.Value != nil && v.Value.GetDatum().Kind() != types.KindNull {
+			strVal, err := v.Value.GetDatum().ToString()
+			if err != nil {
+				e.err = errors.Trace(err)
+				return false
+			}
+			sessionVars.Users[name] = strings.ToLower(strVal)
+			v.SetString(strVal)
+			return true
+		}
 		// user vars
 		if value, ok := sessionVars.Users[name]; ok {
 			v.SetString(value)

diff --git a/parser/parser.y b/parser/parser.y
@@ -70,6 +70,7 @@ import (
 	as		"AS"
 	asc		"ASC"
 	ascii		"ASCII"
+	assignmentEq	":="
 	at		"AT"
 	autoIncrement	"AUTO_INCREMENT"
 	avg		"AVG"
@@ -663,6 +664,7 @@ import (
 %left   tableRefPriority
 %precedence lowerThanOn
 %precedence on
+%right  assignmentEq
 %left 	oror or
 %left 	xor
 %left 	andand and
@@ -1537,7 +1539,18 @@ NUM:
 	intLit
 
 Expression:
-	Expression logOr Expression %prec oror
+	"USER_VAR" assignmentEq Expression %prec assignmentEq
+	{
+		v := $1.(string)
+		v = strings.TrimPrefix(v, "@")
+		$$ = &ast.VariableExpr{
+				Name: 	  v,
+				IsGlobal: false,
+				IsSystem: false,
+				Value:	  $3.(ast.ExprNode),
+		}
+	}
+|	Expression logOr Expression %prec oror
 	{
 		$$ = &ast.BinaryOperationExpr{Op: opcode.OrOr, L: $1.(ast.ExprNode), R: $3.(ast.ExprNode)}
 	}
@@ -1613,6 +1626,18 @@ Factor:
 	{
 		$$ = &ast.BinaryOperationExpr{Op: $2.(opcode.Op), L: $1.(ast.ExprNode), R: $3.(ast.ExprNode)}
 	}
+|	Factor CompareOp "USER_VAR" assignmentEq PredicateExpr %prec assignmentEq
+	{
+		v := $3.(string)
+		v = strings.TrimPrefix(v, "@")
+		variable := &ast.VariableExpr{
+				Name: 	  v,
+				IsGlobal: false,
+				IsSystem: false,
+				Value:	  $5.(ast.ExprNode),
+		}
+		$$ = &ast.BinaryOperationExpr{Op: $2.(opcode.Op), L: $1.(ast.ExprNode), R: variable}
+	}
 |	Factor CompareOp AnyOrAll SubSelect %prec eq
 	{
 		sq := $4.(*ast.SubqueryExpr)
@@ -3664,6 +3689,12 @@ VariableAssignment:
 		v = strings.TrimPrefix(v, "@")
 		$$ = &ast.VariableAssignment{Name: v, Value: $3.(ast.ExprNode)}
 	}
+|	"USER_VAR" assignmentEq Expression
+	{
+		v := $1.(string)
+		v = strings.TrimPrefix(v, "@")
+		$$ = &ast.VariableAssignment{Name: v, Value: $3.(ast.ExprNode)}
+	}
 
 VariableAssignmentList:
 	{

diff --git a/parser/parser_test.go b/parser/parser_test.go
@@ -198,6 +198,7 @@ func (s *testParserSuite) TestDMLStmt(c *C) {
 		// set
 		// user defined
 		{"SET @a = 1", true},
+		{"SET @b := 1", true},
 		// session system variables
 		{"SET SESSION autocommit = 1", true},
 		{"SET @@session.autocommit = 1", true},

diff --git a/parser/scanner.l b/parser/scanner.l
@@ -653,6 +653,7 @@ action		{a}{c}{t}{i}{o}{n}
 "<<"			return lsh
 "<="			return le
 "=" 			return eq
+":="			return assignmentEq
 ">="			return ge
 "!="			return neq
 "<>"			return neq

diff --git a/session_test.go b/session_test.go
@@ -687,6 +687,54 @@ func (s *testSessionSuite) TestIssue1135(c *C) {
 	c.Assert(err, IsNil)
 }
 
+func (s *testSessionSuite) TestIssue1114(c *C) {
+	defer testleak.AfterTest(c)()
+	store := newStore(c, s.dbName)
+	se := newSession(c, store, s.dbName)
+
+	mustExecSQL(c, se, "set @tmp = 0")
+	mustExecSQL(c, se, "set @tmp := @tmp + 1")
+	r := mustExecSQL(c, se, "select @tmp")
+	row, err := r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 1)
+
+	r = mustExecSQL(c, se, "select @tmp1 = 1, @tmp2 := 2")
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, nil, 2)
+
+	r = mustExecSQL(c, se, "select @tmp1 := 11, @tmp2")
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 11, 2)
+
+	mustExecSQL(c, se, "drop table if exists t")
+	mustExecSQL(c, se, "create table t (c int);")
+	mustExecSQL(c, se, "insert into t values (1),(2);")
+	mustExecSQL(c, se, "update t set c = 3 WHERE c = @var:= 1")
+	r = mustExecSQL(c, se, "select * from t")
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 3)
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 2)
+
+	r = mustExecSQL(c, se, "select @tmp := count(*) from t")
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 2)
+
+	r = mustExecSQL(c, se, "select @tmp := c-2 from t where c=3")
+	row, err = r.Next()
+	c.Assert(err, IsNil)
+	match(c, row.Data, 1)
+
+	err = store.Close()
+	c.Assert(err, IsNil)
+}
+
 func (s *testSessionSuite) TestSelectForUpdate(c *C) {
 	defer testleak.AfterTest(c)()
 	store := newStore(c, s.dbName)