generation-expression: add name resolve (#3676)

table/tables/gen_expr.go

@@ -0,0 +1,82 @@
+// Copyright 2017 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tables
+
+import (
+	"fmt"
+
+	"github.com/juju/errors"
+	"github.com/pingcap/tidb/ast"
+	"github.com/pingcap/tidb/model"
+	"github.com/pingcap/tidb/parser"
+)
+
+// getDefaultCharsetAndCollate is copyed from ddl/ddl_api.go.
+func getDefaultCharsetAndCollate() (string, string) {
+	return "utf8", "utf8_bin"
+}
+
+// nameResolver is the visitor to resolve table name and column name.
+// it combines TableInfo and ColumnInfo to a generation expression.
+type nameResolver struct {
+	tableInfo *model.TableInfo
+	err       error
+}
+
+// Enter implements ast.Visitor interface.
+func (nr *nameResolver) Enter(inNode ast.Node) (ast.Node, bool) {
+	return inNode, false
+}
+
+// Leave implements ast.Visitor interface.
+func (nr *nameResolver) Leave(inNode ast.Node) (node ast.Node, ok bool) {
+	switch v := inNode.(type) {
+	case *ast.ColumnNameExpr:
+		for _, col := range nr.tableInfo.Columns {
+			if col.Name.L == v.Name.Name.L {
+				v.Refer = &ast.ResultField{
+					Column: col,
+					Table:  nr.tableInfo,
+				}
+				return inNode, true
+			}
+		}
+		nr.err = errors.Errorf("can't find column %s in %s", v.Name.Name.O, nr.tableInfo.Name.O)
+		return inNode, false
+	}
+	return inNode, true
+}
+
+// ParseExpression parses an ExprNode from a string.
+// When TiDB loads infoschema from TiKV, `GeneratedExprString`
+// of `ColumnInfo` is a string field, so we need to parse
+// it into ast.ExprNode. This function is for that.
+func parseExpression(expr string) (node ast.ExprNode, err error) {
+	expr = fmt.Sprintf("select %s", expr)
+	charset, collation := getDefaultCharsetAndCollate()
+	stmts, err := parser.New().Parse(expr, charset, collation)
+	if err == nil {
+		node = stmts[0].(*ast.SelectStmt).Fields.Fields[0].Expr
+	}
+	return node, errors.Trace(err)
+}
+
+// SimpleResolveName resolves all column names in the expression node.
+func simpleResolveName(node ast.ExprNode, tblInfo *model.TableInfo) (ast.ExprNode, error) {
+	nr := nameResolver{tblInfo, nil}
+	if _, ok := node.Accept(&nr); !ok {
+		return nil, errors.Trace(nr.err)
+	}
+	return node, nil
+}

table/tables/gen_expr_test.go

@@ -11,25 +11,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package parser
+package tables
 
 import (
-	"testing"
-
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/ast"
 )
 
-func TestT(t *testing.T) {
-	CustomVerboseFlag = true
-	TestingT(t)
-}
+var _ = Suite(&testGenExprSuite{})
 
-var _ = Suite(&testParserSuite{})
+type testGenExprSuite struct{}
 
-type testParserSuite struct{}
-
-func (s *testParserSuite) TestParseExpression(c *C) {
+func (s *testGenExprSuite) TestParseExpression(c *C) {
 	tests := []struct {
 		input   string
 		output  string
@@ -38,7 +31,7 @@ func (s *testParserSuite) TestParseExpression(c *C) {
 		{"json_extract(a, '$.a')", "json_extract", true},
 	}
 	for _, tt := range tests {
-		node, err := ParseExpression(tt.input)
+		node, err := parseExpression(tt.input)
 		if tt.success {
 			fc := node.(*ast.FuncCallExpr)
 			c.Assert(fc.FnName.L, Equals, tt.output)

table/tables/tables.go

@@ -34,7 +34,6 @@ import (
 	"github.com/pingcap/tidb/terror"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/codec"
-	"github.com/pingcap/tidb/util/parser"
 	"github.com/pingcap/tidb/util/types"
 	"github.com/pingcap/tipb/go-binlog"
 )
@@ -80,7 +79,11 @@ func TableFromMeta(alloc autoid.Allocator, tblInfo *model.TableInfo) (table.Tabl
 
 		col := table.ToColumn(colInfo)
 		if len(colInfo.GeneratedExprString) != 0 {
-			expr, err := parser.ParseExpression(colInfo.GeneratedExprString)
+			expr, err := parseExpression(colInfo.GeneratedExprString)
+			if err != nil {
+				return nil, errors.Trace(err)
+			}
+			expr, err = simpleResolveName(expr, tblInfo)
 			if err != nil {
 				return nil, errors.Trace(err)
 			}

table/tables/tables_test.go

@@ -249,6 +249,7 @@ func (ts *testSuite) TestRowKeyCodec(c *C) {
 
 func (ts *testSuite) TestUnsignedPK(c *C) {
 	defer testleak.AfterTest(c)()
+	ts.se.Execute("DROP TABLE IF EXISTS test.tPK")
 	_, err := ts.se.Execute("CREATE TABLE test.tPK (a bigint unsigned primary key, b varchar(255))")
 	c.Assert(err, IsNil)
 	ctx := ts.se.(context.Context)
@@ -267,6 +268,7 @@ func (ts *testSuite) TestUnsignedPK(c *C) {
 
 func (ts *testSuite) TestIterRecords(c *C) {
 	defer testleak.AfterTest(c)()
+	ts.se.Execute("DROP TABLE IF EXISTS test.tIter")
 	_, err := ts.se.Execute("CREATE TABLE test.tIter (a int primary key, b int)")
 	c.Assert(err, IsNil)
 	_, err = ts.se.Execute("INSERT test.tIter VALUES (1, 2), (2, NULL)")

util/parser/expression.go

@@ -1,43 +0,0 @@
-// Copyright 2017 PingCAP, Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package parser
-
-import (
-	"fmt"
-
-	"github.com/juju/errors"
-	"github.com/pingcap/tidb/ast"
-	"github.com/pingcap/tidb/parser"
-)
-
-// getDefaultCharsetAndCollate is copyed from ddl/ddl_api.go.
-func getDefaultCharsetAndCollate() (string, string) {
-	return "utf8", "utf8_bin"
-}
-
-// ParseExpression parses an ExprNode from a string.
-// Where should we use this?
-//   When TiDB bootstraps, it'll load infoschema from TiKV.
-//   Because some ColumnInfos have attribute `GeneratedExprString`,
-//   we need to parse that string into ast.ExprNode.
-func ParseExpression(expr string) (node ast.ExprNode, err error) {
-	expr = fmt.Sprintf("select %s", expr)
-	charset, collation := getDefaultCharsetAndCollate()
-	stmts, err := parser.New().Parse(expr, charset, collation)
-	if err == nil {
-		sel := stmts[0].(*ast.SelectStmt)
-		node = sel.Fields.Fields[0].Expr
-	}
-	return node, errors.Trace(err)
-}