Merge pull request #6708 from planetscale/evaluate-sysvars

Evaluate system variables

go/sqltypes/bind_variables.go

@@ -75,6 +75,14 @@ func Int32BindVariable(v int32) *querypb.BindVariable {
  return ValueBindVariable(NewInt32(v))
 }
 
+// BoolBindVariable converts an bool to a int32 bind var.
+func BoolBindVariable(v bool) *querypb.BindVariable {
+ if v {
+ return Int32BindVariable(1)
+ }
+ return Int32BindVariable(0)
+}
+
 // Int64BindVariable converts an int64 to a bind var.
 func Int64BindVariable(v int64) *querypb.BindVariable {
  return ValueBindVariable(NewInt64(v))

go/vt/sqlparser/bind_var_needs.go

@@ -0,0 +1,72 @@
+/*
+Copyright 2020 The Vitess Authors.
+
+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,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sqlparser
+
+// BindVarNeeds represents the bind vars that need to be provided as the result of expression rewriting.
+type BindVarNeeds struct {
+ NeedFunctionResult,
+ NeedSystemVariable,
+ // NeedUserDefinedVariables keeps track of all user defined variables a query is using
+ NeedUserDefinedVariables []string
+}
+
+//MergeWith adds bind vars needs coming from sub scopes
+func (bvn *BindVarNeeds) MergeWith(other *BindVarNeeds) {
+ bvn.NeedFunctionResult = append(bvn.NeedFunctionResult, other.NeedFunctionResult...)
+ bvn.NeedSystemVariable = append(bvn.NeedSystemVariable, other.NeedSystemVariable...)
+ bvn.NeedUserDefinedVariables = append(bvn.NeedUserDefinedVariables, other.NeedUserDefinedVariables...)
+}
+
+//AddFuncResult adds a function bindvar need
+func (bvn *BindVarNeeds) AddFuncResult(name string) {
+ bvn.NeedFunctionResult = append(bvn.NeedFunctionResult, name)
+}
+
+//AddSysVar adds a system variable bindvar need
+func (bvn *BindVarNeeds) AddSysVar(name string) {
+ bvn.NeedSystemVariable = append(bvn.NeedSystemVariable, name)
+}
+
+//AddUserDefVar adds a user defined variable bindvar need
+func (bvn *BindVarNeeds) AddUserDefVar(name string) {
+ bvn.NeedUserDefinedVariables = append(bvn.NeedUserDefinedVariables, name)
+}
+
+//NeedsFuncResult says if a function result needs to be provided
+func (bvn *BindVarNeeds) NeedsFuncResult(name string) bool {
+ return contains(bvn.NeedFunctionResult, name)
+}
+
+//NeedsSysVar says if a function result needs to be provided
+func (bvn *BindVarNeeds) NeedsSysVar(name string) bool {
+ return contains(bvn.NeedSystemVariable, name)
+}
+
+func (bvn *BindVarNeeds) HasRewrites() bool {
+ return len(bvn.NeedFunctionResult) > 0 ||
+ len(bvn.NeedUserDefinedVariables) > 0 ||
+ len(bvn.NeedSystemVariable) > 0
+}
+
+func contains(strings []string, name string) bool {
+ for _, s := range strings {
+ if name == s {
+ return true
+ }
+ }
+ return false
+}

go/vt/sqlparser/expression_rewriting.go

@@ -19,7 +19,8 @@ package sqlparser
 import (
  "strings"
 
- "vitess.io/vitess/go/vt/log"
+ "vitess.io/vitess/go/vt/sysvars"
+
  querypb "vitess.io/vitess/go/vt/proto/query"
  "vitess.io/vitess/go/vt/proto/vtrpc"
  "vitess.io/vitess/go/vt/vterrors"
@@ -33,15 +34,6 @@ func PrepareAST(in Statement, bindVars map[string]*querypb.BindVariable, prefix
  return RewriteAST(in)
 }
 
-// BindVarNeeds represents the bind vars that need to be provided as the result of expression rewriting.
-type BindVarNeeds struct {
- NeedLastInsertID bool
- NeedDatabase bool
- NeedFoundRows bool
- NeedRowCount bool
- NeedUserDefinedVariables []string
-}
-
 // RewriteAST rewrites the whole AST, replacing function calls and adding column aliases to queries
 func RewriteAST(in Statement) (*RewriteASTResult, error) {
  er := newExpressionRewriter()
@@ -56,21 +48,8 @@ func RewriteAST(in Statement) (*RewriteASTResult, error) {
  }
 
  r := &RewriteASTResult{
- AST: out,
- }
- for k := range er.bindVars {
- switch k {
- case LastInsertIDName:
- r.NeedLastInsertID = true
- case DBVarName:
- r.NeedDatabase = true
- case FoundRowsName:
- r.NeedFoundRows = true
- case RowCountName:
- r.NeedRowCount = true
- default:
- r.NeedUserDefinedVariables = append(r.NeedUserDefinedVariables, k)
- }
+ AST: out,
+ BindVarNeeds: er.bindVars,
  }
  return r, nil
 }
@@ -96,18 +75,18 @@ func shouldRewriteDatabaseFunc(in Statement) bool {
 
 // RewriteASTResult contains the rewritten ast and meta information about it
 type RewriteASTResult struct {
- BindVarNeeds
+ *BindVarNeeds
  AST Statement // The rewritten AST
 }
 
 type expressionRewriter struct {
- bindVars map[string]struct{}
+ bindVars *BindVarNeeds
  shouldRewriteDatabaseFunc bool
  err error
 }
 
 func newExpressionRewriter() *expressionRewriter {
- return &expressionRewriter{bindVars: make(map[string]struct{})}
+ return &expressionRewriter{bindVars: &BindVarNeeds{}}
 }
 
 const (
@@ -127,6 +106,18 @@ const (
  UserDefinedVariableName = "__vtudv"
 )
 
+func (er *expressionRewriter) rewriteAliasedExpr(cursor *Cursor, node *AliasedExpr) (*BindVarNeeds, error) {
+ inner := newExpressionRewriter()
+ inner.shouldRewriteDatabaseFunc = er.shouldRewriteDatabaseFunc
+ tmp := Rewrite(node.Expr, inner.goingDown, nil)
+ newExpr, ok := tmp.(Expr)
+ if !ok {
+ return nil, vterrors.Errorf(vtrpc.Code_INTERNAL, "failed to rewrite AST. function expected to return Expr returned a %s", String(tmp))
+ }
+ node.Expr = newExpr
+ return inner.bindVars, nil
+}
+
 func (er *expressionRewriter) goingDown(cursor *Cursor) bool {
  switch node := cursor.Node().(type) {
  // select last_insert_id() -> select :__lastInsertId as `last_insert_id()`
@@ -136,35 +127,50 @@ func (er *expressionRewriter) goingDown(cursor *Cursor) bool {
  if ok && aliasedExpr.As.IsEmpty() {
  buf := NewTrackedBuffer(nil)
  aliasedExpr.Expr.Format(buf)
- inner := newExpressionRewriter()
- inner.shouldRewriteDatabaseFunc = er.shouldRewriteDatabaseFunc
- tmp := Rewrite(aliasedExpr.Expr, inner.goingDown, nil)
- newExpr, ok := tmp.(Expr)
- if !ok {
- log.Errorf("failed to rewrite AST. function expected to return Expr returned a %s", String(tmp))
+ innerBindVarNeeds, err := er.rewriteAliasedExpr(cursor, aliasedExpr)
+ if err != nil {
+ er.err = err
  return false
  }
- aliasedExpr.Expr = newExpr
- if inner.didAnythingChange() {
+ if innerBindVarNeeds.HasRewrites() {
  aliasedExpr.As = NewColIdent(buf.String())
  }
- for k := range inner.bindVars {
- er.needBindVarFor(k)
- }
+ er.bindVars.MergeWith(innerBindVarNeeds)
  }
  }
  case *FuncExpr:
  er.funcRewrite(cursor, node)
  case *ColName:
- if node.Name.at == SingleAt {
- udv := strings.ToLower(node.Name.CompliantName())
- cursor.Replace(bindVarExpression(UserDefinedVariableName + udv))
- er.needBindVarFor(udv)
+ switch node.Name.at {
+ case SingleAt:
+ er.udvRewrite(cursor, node)
+ case DoubleAt:
+ er.sysVarRewrite(cursor, node)
  }
  }
  return true
 }
 
+func (er *expressionRewriter) sysVarRewrite(cursor *Cursor, node *ColName) {
+ lowered := node.Name.Lowered()
+ switch lowered {
+ case sysvars.Autocommit.Name,
+ sysvars.ClientFoundRows.Name,
+ sysvars.SkipQueryPlanCache.Name,
+ sysvars.SQLSelectLimit.Name,
+ sysvars.TransactionMode.Name,
+ sysvars.Workload.Name:
+ cursor.Replace(bindVarExpression("__vt" + lowered))
+ er.bindVars.AddSysVar(lowered)
+ }
+}
+
+func (er *expressionRewriter) udvRewrite(cursor *Cursor, node *ColName) {
+ udv := strings.ToLower(node.Name.CompliantName())
+ cursor.Replace(bindVarExpression(UserDefinedVariableName + udv))
+ er.bindVars.AddUserDefVar(udv)
+}
+
 func (er *expressionRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
  switch {
  // last_insert_id() -> :__lastInsertId
@@ -173,7 +179,7 @@ func (er *expressionRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
  er.err = vterrors.New(vtrpc.Code_UNIMPLEMENTED, "Argument to LAST_INSERT_ID() not supported")
  } else {
  cursor.Replace(bindVarExpression(LastInsertIDName))
- er.needBindVarFor(LastInsertIDName)
+ er.bindVars.AddFuncResult(LastInsertIDName)
  }
  // database() -> :__vtdbname
  case er.shouldRewriteDatabaseFunc &&
@@ -183,38 +189,27 @@ func (er *expressionRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
  er.err = vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "Syntax error. %s() takes no arguments", node.Name.String())
  } else {
  cursor.Replace(bindVarExpression(DBVarName))
- er.needBindVarFor(DBVarName)
+ er.bindVars.AddFuncResult(DBVarName)
  }
  // found_rows() -> :__vtfrows
  case node.Name.EqualString("found_rows"):
  if len(node.Exprs) > 0 {
  er.err = vterrors.New(vtrpc.Code_INVALID_ARGUMENT, "Arguments to FOUND_ROWS() not supported")
  } else {
  cursor.Replace(bindVarExpression(FoundRowsName))
- er.needBindVarFor(FoundRowsName)
+ er.bindVars.AddFuncResult(FoundRowsName)
  }
  // row_count() -> :__vtrcount
  case node.Name.EqualString("row_count"):
  if len(node.Exprs) > 0 {
  er.err = vterrors.New(vtrpc.Code_INVALID_ARGUMENT, "Arguments to ROW_COUNT() not supported")
  } else {
  cursor.Replace(bindVarExpression(RowCountName))
- er.needBindVarFor(RowCountName)
+ er.bindVars.AddFuncResult(RowCountName)
  }
  }
 }
 
-// instead of creating new objects, we'll reuse this one
-var token = struct{}{}
-
-func (er *expressionRewriter) needBindVarFor(name string) {
- er.bindVars[name] = token
-}
-
-func (er *expressionRewriter) didAnythingChange() bool {
- return len(er.bindVars) > 0
-}
-
 func bindVarExpression(name string) Expr {
  return NewArgument([]byte(":" + name))
 }