Golang ast.SetFlag函数代码示例

func (s *testEvaluatorSuite) TestEvaluatedFlag(c *C) {
	l := ast.NewValueExpr(int64(1))
	r := ast.NewValueExpr(int64(2))
	b := &ast.BinaryOperationExpr{L: l, R: r, Op: opcode.Plus}
	ast.SetFlag(b)
	c.Assert(ast.IsPreEvaluable(b), Equals, true)
	d, err := Eval(s.ctx, b)
	c.Assert(ast.IsEvaluated(b), Equals, true)
	c.Assert(err, IsNil)
	c.Assert(d, testutil.DatumEquals, types.NewIntDatum(3))

	funcCall := &ast.FuncCallExpr{
		FnName: model.NewCIStr("abs"),
		Args:   []ast.ExprNode{ast.NewValueExpr(int(-1))},
	}
	b = &ast.BinaryOperationExpr{L: funcCall, R: r, Op: opcode.Plus}
	ast.ResetEvaluatedFlag(b)
	ast.SetFlag(b)
	c.Assert(ast.IsPreEvaluable(b), Equals, true)
	d, err = Eval(s.ctx, b)
	c.Assert(ast.IsEvaluated(b), Equals, false)
	c.Assert(err, IsNil)
	c.Assert(d, testutil.DatumEquals, types.NewIntDatum(3))

	rf := &ast.ResultField{Expr: ast.NewValueExpr(int64(1))}
	colExpr := &ast.ColumnNameExpr{Refer: rf}
	b = &ast.BinaryOperationExpr{L: colExpr, R: r, Op: opcode.Plus}
	ast.ResetEvaluatedFlag(b)
	ast.SetFlag(b)
	c.Assert(ast.IsPreEvaluable(b), Equals, false)
	d, err = Eval(s.ctx, b)
	c.Assert(ast.IsEvaluated(b), Equals, false)
	c.Assert(err, IsNil)
	c.Assert(d, testutil.DatumEquals, types.NewIntDatum(3))
}

func (s *testPlanSuite) TestMultiColumnIndex(c *C) {
	defer testleak.AfterTest(c)()
	cases := []struct {
		sql              string
		accessEqualCount int
		usedColumnCount  int
	}{
		{"select * from t where c = 0 and d = 0 and e = 0", 3, 3},
		{"select * from t where c = 0 and d = 0 and e > 0", 2, 3},
		{"select * from t where d > 0 and e = 0 and c = 0", 1, 2},
	}
	for _, ca := range cases {
		comment := Commentf("for %s", ca.sql)
		s, err := s.ParseOneStmt(ca.sql, "", "")
		c.Assert(err, IsNil, comment)
		stmt := s.(*ast.SelectStmt)
		ast.SetFlag(stmt)
		mockResolve(stmt)
		b := &planBuilder{}
		p := b.buildFrom(stmt)
		err = Refine(p)
		c.Assert(err, IsNil)
		idxScan, ok := p.(*IndexScan)
		c.Assert(ok, IsTrue)
		c.Assert(idxScan.AccessEqualCount, Equals, ca.accessEqualCount)
		c.Assert(idxScan.Ranges[0].LowVal, HasLen, ca.usedColumnCount)
	}
}

func (s *testEvaluatorSuite) TestCast(c *C) {
	defer testleak.AfterTest(c)()
	f := types.NewFieldType(mysql.TypeLonglong)

	expr := &ast.FuncCastExpr{
		Expr: ast.NewValueExpr(1),
		Tp:   f,
	}
	ast.SetFlag(expr)
	v, err := Eval(s.ctx, expr)
	c.Assert(err, IsNil)
	c.Assert(v, testutil.DatumEquals, types.NewDatum(int64(1)))

	f.Flag |= mysql.UnsignedFlag
	v, err = Eval(s.ctx, expr)
	c.Assert(err, IsNil)
	c.Assert(v, testutil.DatumEquals, types.NewDatum(uint64(1)))

	f.Tp = mysql.TypeString
	f.Charset = charset.CharsetBin
	v, err = Eval(s.ctx, expr)
	c.Assert(err, IsNil)
	c.Assert(v, testutil.DatumEquals, types.NewDatum([]byte("1")))

	f.Tp = mysql.TypeString
	f.Charset = "utf8"
	v, err = Eval(s.ctx, expr)
	c.Assert(err, IsNil)
	c.Assert(v, testutil.DatumEquals, types.NewDatum("1"))

	expr.Expr = ast.NewValueExpr(nil)
	v, err = Eval(s.ctx, expr)
	c.Assert(err, IsNil)
	c.Assert(v.Kind(), Equals, types.KindNull)
}

// Parse parses a query string to raw ast.StmtNode.
// If charset or collation is "", default charset and collation will be used.
func (parser *Parser) Parse(sql, charset, collation string) ([]ast.StmtNode, error) {
	if charset == "" {
		charset = mysql.DefaultCharset
	}
	if collation == "" {
		collation = mysql.DefaultCollationName
	}
	parser.charset = charset
	parser.collation = collation
	parser.src = sql
	parser.result = parser.result[:0]

	sql = handleMySQLSpecificCode(sql)

	var l yyLexer
	parser.lexer.reset(sql)
	l = &parser.lexer
	yyParse(l, parser)

	if len(l.Errors()) != 0 {
		return nil, errors.Trace(l.Errors()[0])
	}
	for _, stmt := range parser.result {
		ast.SetFlag(stmt)
	}
	return parser.result, nil
}

// Compile compiles an ast.StmtNode to a stmt.Statement.
// If it is supported to use new plan and executer, it optimizes the node to
// a plan, and we wrap the plan in an adapter as stmt.Statement.
// If it is not supported, the node will be converted to old statement.
func (c *Compiler) Compile(ctx context.Context, node ast.StmtNode) (stmt.Statement, error) {
	if optimizer.IsSupported(node) {
		ast.SetFlag(node)

		is := sessionctx.GetDomain(ctx).InfoSchema()
		if err := optimizer.Preprocess(node, is, ctx); err != nil {
			return nil, errors.Trace(err)
		}
		// Validate should be after NameResolve.
		if err := optimizer.Validate(node, false); err != nil {
			return nil, errors.Trace(err)
		}
		sb := NewSubQueryBuilder(is)
		p, err := optimizer.Optimize(ctx, node, sb)
		if err != nil {
			return nil, errors.Trace(err)
		}
		sa := &statementAdapter{
			is:   is,
			plan: p,
		}
		return sa, nil
	}
	c.converter = &converter.Converter{}
	s, err := c.converter.Convert(node)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return s, nil
}

func (b *planBuilder) buildTablePlanFromJoinPath(path *joinPath) Plan {
	for _, equiv := range path.eqConds {
		columnNameExpr := &ast.ColumnNameExpr{}
		columnNameExpr.Name = &ast.ColumnName{}
		columnNameExpr.Name.Name = equiv.left.Column.Name
		columnNameExpr.Name.Table = equiv.left.Table.Name
		columnNameExpr.Refer = equiv.left
		condition := &ast.BinaryOperationExpr{L: columnNameExpr, R: equiv.right.Expr, Op: opcode.EQ}
		ast.SetFlag(condition)
		path.conditions = append(path.conditions, condition)
	}
	candidates := b.buildAllAccessMethodsPlan(path.table, path.conditions)
	var p Plan
	var lowestCost float64
	for _, can := range candidates {
		cost := EstimateCost(can)
		if p == nil {
			p = can
			lowestCost = cost
		}
		if cost < lowestCost {
			p = can
			lowestCost = cost
		}
	}
	return p
}

func (s *testPlanSuite) TestTableScanWithOrder(c *C) {
	defer testleak.AfterTest(c)()
	// Sort result by scanning PKHandle column.
	sql := "select * from t order by a limit 1;"
	stmt, err := s.ParseOneStmt(sql, "", "")
	c.Assert(err, IsNil)
	ast.SetFlag(stmt)

	err = newMockResolve(stmt)
	c.Assert(err, IsNil)

	builder := &planBuilder{
		allocator: new(idAllocator),
		ctx:       mock.NewContext(),
		colMapper: make(map[*ast.ColumnNameExpr]int),
	}
	p := builder.build(stmt)
	c.Assert(builder.err, IsNil)
	logic, ok := p.(LogicalPlan)
	c.Assert(ok, IsTrue)
	// Get physical plan.
	_, pp, _, err := logic.convert2PhysicalPlan(nil)
	c.Assert(err, IsNil)
	// Limit->Projection->PhysicalTableScan
	// Get PhysicalTableScan plan.
	cpp, ok := pp.p.GetChildByIndex(0).GetChildByIndex(0).(*PhysicalTableScan)
	c.Assert(cpp, NotNil)
	c.Assert(ok, IsTrue)
	// Make sure KeepOrder is true.
	c.Assert(cpp.KeepOrder, IsTrue)
}

// Compile compiles an ast.StmtNode to a stmt.Statement.
// If it is supported to use new plan and executer, it optimizes the node to
// a plan, and we wrap the plan in an adapter as stmt.Statement.
// If it is not supported, the node will be converted to old statement.
func (c *Compiler) Compile(ctx context.Context, node ast.StmtNode) (stmt.Statement, error) {
	if optimizer.IsSupported(node) {
		ast.SetFlag(node)
		if err := optimizer.Validate(node, false); err != nil {
			return nil, errors.Trace(err)
		}
		is := sessionctx.GetDomain(ctx).InfoSchema()
		if err := optimizer.ResolveName(node, is, ctx); err != nil {
			return nil, errors.Trace(err)
		}
		p, err := optimizer.Optimize(ctx, node)
		if err != nil {
			return nil, errors.Trace(err)
		}
		sa := &statementAdapter{
			is:   is,
			plan: p,
		}
		return sa, nil
	}
	c.converter = &converter.Converter{}
	s, err := c.converter.Convert(node)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return s, nil
}

// Compile compiles an ast.StmtNode to a stmt.Statement.
// If it is supported to use new plan and executer, it optimizes the node to
// a plan, and we wrap the plan in an adapter as stmt.Statement.
// If it is not supported, the node will be converted to old statement.
func (c *Compiler) Compile(ctx context.Context, node ast.StmtNode) (ast.Statement, error) {
	ast.SetFlag(node)
	if _, ok := node.(*ast.UpdateStmt); ok {
		sVars := variable.GetSessionVars(ctx)
		sVars.InUpdateStmt = true
		defer func() {
			sVars.InUpdateStmt = false
		}()
	}

	is := sessionctx.GetDomain(ctx).InfoSchema()
	if err := plan.Preprocess(node, is, ctx); err != nil {
		return nil, errors.Trace(err)
	}
	// Validate should be after NameResolve.
	if err := plan.Validate(node, false); err != nil {
		return nil, errors.Trace(err)
	}
	sb := NewSubQueryBuilder(is)

	p, err := plan.Optimize(ctx, node, sb, is)
	if err != nil {
		return nil, errors.Trace(err)
	}
	_, isDDL := node.(ast.DDLNode)
	sa := &statement{
		is:    is,
		plan:  p,
		text:  node.Text(),
		isDDL: isDDL,
	}
	return sa, nil
}

func (s *testPlanSuite) TestVisitCount(c *C) {
	sqls := []string{
		"select t1.c1, t2.c2 from t1, t2",
		"select * from t1 left join t2 on t1.c1 = t2.c1",
		"select * from t1 group by t1.c1 having sum(t1.c2) = 1",
		"select * from t1 where t1.c1 > 2 order by t1.c2 limit 100",
		"insert t1 values (1), (2)",
		"delete from t1 where false",
		"truncate table t1",
		"do 1",
		"show databases",
	}
	for _, sql := range sqls {
		stmt, err := parser.ParseOneStmt(sql, "", "")
		c.Assert(err, IsNil, Commentf(sql))
		ast.SetFlag(stmt)
		mockJoinResolve(c, stmt)
		b := &planBuilder{}
		p := b.build(stmt)
		c.Assert(b.err, IsNil)
		visitor := &countVisitor{}
		for i := 0; i < 5; i++ {
			visitor.skipAt = i
			visitor.enterCount = 0
			visitor.leaveCount = 0
			p.Accept(visitor)
			c.Assert(visitor.enterCount, Equals, visitor.leaveCount, Commentf(sql))
		}
	}
}

func (s *testPlanSuite) TestJoinReOrder(c *C) {
	defer testleak.AfterTest(c)()
	cases := []struct {
		sql  string
		best string
	}{
		{
			sql:  "select * from t t1, t t2, t t3, t t4, t t5, t t6 where t1.a = t2.b and t2.a = t3.b and t3.c = t4.a and t4.d = t2.c and t5.d = t6.d",
			best: "LeftHashJoin{LeftHashJoin{LeftHashJoin{LeftHashJoin{Table(t)->Table(t)}(t1.a,t2.b)->Table(t)}(t2.a,t3.b)->Table(t)}(t3.c,t4.a)(t2.c,t4.d)->LeftHashJoin{Table(t)->Table(t)}(t5.d,t6.d)}->Projection",
		},
		{
			sql:  "select * from t t1, t t2, t t3, t t4, t t5, t t6, t t7, t t8 where t1.a = t8.a",
			best: "LeftHashJoin{LeftHashJoin{LeftHashJoin{LeftHashJoin{Table(t)->Table(t)}(t1.a,t8.a)->Table(t)}->LeftHashJoin{Table(t)->Table(t)}}->LeftHashJoin{LeftHashJoin{Table(t)->Table(t)}->Table(t)}}->Projection",
		},
		{
			sql:  "select * from t t1, t t2, t t3, t t4, t t5 where t1.a = t5.a and t5.a = t4.a and t4.a = t3.a and t3.a = t2.a and t2.a = t1.a and t1.a = t3.a and t2.a = t4.a and t5.b < 8",
			best: "LeftHashJoin{LeftHashJoin{LeftHashJoin{RightHashJoin{Table(t)->Selection->Table(t)}(t5.a,t1.a)->Table(t)}(t1.a,t2.a)->Table(t)}(t2.a,t3.a)(t1.a,t3.a)->Table(t)}(t5.a,t4.a)(t3.a,t4.a)(t2.a,t4.a)->Projection",
		},
		{
			sql:  "select * from t t1, t t2, t t3, t t4, t t5 where t1.a = t5.a and t5.a = t4.a and t4.a = t3.a and t3.a = t2.a and t2.a = t1.a and t1.a = t3.a and t2.a = t4.a and t3.b = 1 and t4.a = 1",
			best: "LeftHashJoin{LeftHashJoin{LeftHashJoin{Table(t)->Selection->Table(t)}->LeftHashJoin{Table(t)->Table(t)}}->Table(t)}->Projection",
		},
		{
			sql:  "select * from t o where o.b in (select t3.c from t t1, t t2, t t3 where t1.a = t3.a and t2.a = t3.a and t2.a = o.a)",
			best: "Table(t)->Apply(LeftHashJoin{RightHashJoin{Table(t)->Selection->Table(t)}(t2.a,t3.a)->Table(t)}(t3.a,t1.a)->Projection)->Selection->Projection",
		},
		{
			sql:  "select * from t o where o.b in (select t3.c from t t1, t t2, t t3 where t1.a = t3.a and t2.a = t3.a and t2.a = o.a and t1.a = 1)",
			best: "Table(t)->Apply(LeftHashJoin{LeftHashJoin{Table(t)->Table(t)}->Table(t)->Selection}->Projection)->Selection->Projection",
		},
	}
	for _, ca := range cases {
		comment := Commentf("for %s", ca.sql)
		stmt, err := s.ParseOneStmt(ca.sql, "", "")
		c.Assert(err, IsNil, comment)
		ast.SetFlag(stmt)

		err = mockResolve(stmt)
		c.Assert(err, IsNil)

		builder := &planBuilder{
			allocator: new(idAllocator),
			ctx:       mock.NewContext(),
			colMapper: make(map[*ast.ColumnNameExpr]int),
		}
		p := builder.build(stmt)
		c.Assert(builder.err, IsNil)
		lp := p.(LogicalPlan)

		_, lp, err = lp.PredicatePushDown(nil)
		c.Assert(err, IsNil)
		_, err = lp.PruneColumnsAndResolveIndices(lp.GetSchema())
		c.Assert(err, IsNil)
		_, res, _, err := lp.convert2PhysicalPlan(nil)
		c.Assert(err, IsNil)
		p = res.p.PushLimit(nil)
		c.Assert(ToString(p), Equals, ca.best, Commentf("for %s", ca.sql))
	}
}

func (s *testPlanSuite) TestRefine(c *C) {
	UseNewPlanner = true
	defer testleak.AfterTest(c)()
	cases := []struct {
		sql  string
		best string
	}{
		{
			sql:  "select a from t where c = 4 and d = 5 and e = 6",
			best: "Index(t.c_d_e)[[4 5 6,4 5 6]]->Selection->Projection",
		},
		{
			sql:  "select a from t where d = 4 and c = 5",
			best: "Index(t.c_d_e)[[5 4,5 4]]->Selection->Projection",
		},
		{
			sql:  "select a from t where c = 4 and e < 5",
			best: "Index(t.c_d_e)[[4,4]]->Selection->Projection",
		},
		{
			sql:  "select a from t where c = 4 and d <= 5 and d > 3",
			best: "Index(t.c_d_e)[[4 3,4 5]]->Selection->Projection",
		},
		{
			sql:  "select a from t where d <= 5 and d > 3",
			best: "Index(t.c_d_e)[[<nil>,<nil>]]->Selection->Projection",
		},
		{
			sql:  "select a from t where c <= 5 and c >= 3 and d = 1",
			best: "Index(t.c_d_e)[[3,5]]->Selection->Projection",
		},
	}
	for _, ca := range cases {
		comment := Commentf("for %s", ca.sql)
		stmt, err := s.ParseOneStmt(ca.sql, "", "")
		c.Assert(err, IsNil, comment)
		ast.SetFlag(stmt)

		err = newMockResolve(stmt)
		c.Assert(err, IsNil)

		builder := &planBuilder{
			allocator: new(idAllocator),
		}
		p := builder.build(stmt).(LogicalPlan)
		c.Assert(builder.err, IsNil)

		_, p, err = p.PredicatePushDown(nil)
		c.Assert(err, IsNil)
		_, err = p.PruneColumnsAndResolveIndices(p.GetSchema())
		c.Assert(err, IsNil)
		np := p.Convert2PhysicalPlan()
		err = refine(np)
		c.Assert(err, IsNil)
		c.Assert(ToString(np), Equals, ca.best, Commentf("for %s", ca.sql))
	}
	UseNewPlanner = false
}

// Prepare prepares a raw statement parsed from parser.
// The statement must be prepared before it can be passed to optimize function.
// We pass InfoSchema instead of getting from Context in case it is changed after resolving name.
func Prepare(is infoschema.InfoSchema, ctx context.Context, node ast.Node) error {
	ast.SetFlag(node)
	if err := Preprocess(node, is, ctx); err != nil {
		return errors.Trace(err)
	}
	if err := Validate(node, true); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// ParseOneStmt parses a query and returns an ast.StmtNode.
// The query must have one statement, otherwise ErrSyntax is returned.
func (parser *Parser) ParseOneStmt(sql, charset, collation string) (ast.StmtNode, error) {
	stmts, err := parser.Parse(sql, charset, collation)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if len(stmts) != 1 {
		return nil, ErrSyntax
	}
	ast.SetFlag(stmts[0])
	return stmts[0], nil
}

// Prepare prepares a raw statement parsed from parser.
// The statement must be prepared before it can be passed to optimize function.
// We pass InfoSchema instead of getting from Context in case it is changed after resolving name.
func Prepare(is infoschema.InfoSchema, ctx context.Context, node ast.Node) error {
	if err := Validate(node, true); err != nil {
		return errors.Trace(err)
	}
	ast.SetFlag(node)
	if err := ResolveName(node, is, ctx); err != nil {
		return errors.Trace(err)
	}
	if err := InferType(node); err != nil {
		return errors.Trace(err)
	}
	return nil
}