boost: draft for vtgate boost query checks

go/vt/sqlparser/analyzer.go

@@ -65,7 +65,7 @@ const (
  StmtShowMigrationLogs
 )
 
-//ASTToStatementType returns a StatementType from an AST stmt
+// ASTToStatementType returns a StatementType from an AST stmt
 func ASTToStatementType(stmt Statement) StatementType {
  switch stmt.(type) {
  case *Select, *Union:
@@ -121,7 +121,7 @@ func ASTToStatementType(stmt Statement) StatementType {
  }
 }
 
-//CanNormalize takes Statement and returns if the statement can be normalized.
+// CanNormalize takes Statement and returns if the statement can be normalized.
 func CanNormalize(stmt Statement) bool {
  switch stmt.(type) {
  case *Select, *Union, *Insert, *Update, *Delete, *Set, *CallProc, *Stream: // TODO: we could merge this logic into ASTrewriter
@@ -140,7 +140,7 @@ func CachePlan(stmt Statement) bool {
  return false
 }
 
-//MustRewriteAST takes Statement and returns true if RewriteAST must run on it for correct execution irrespective of user flags.
+// MustRewriteAST takes Statement and returns true if RewriteAST must run on it for correct execution irrespective of user flags.
 func MustRewriteAST(stmt Statement) bool {
  switch node := stmt.(type) {
  case *Set:
@@ -301,7 +301,7 @@ func IsDML(sql string) bool {
  return false
 }
 
-//IsDMLStatement returns true if the query is an INSERT, UPDATE or DELETE statement.
+// IsDMLStatement returns true if the query is an INSERT, UPDATE or DELETE statement.
 func IsDMLStatement(stmt Statement) bool {
  switch stmt.(type) {
  case *Insert, *Update, *Delete:
@@ -490,7 +490,7 @@ func NewPlanValue(node Expr) (sqltypes.PlanValue, error) {
  return sqltypes.PlanValue{}, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expression is too complex '%v'", String(node))
 }
 
-//IsLockingFunc returns true for all functions that are used to work with mysql advisory locks
+// IsLockingFunc returns true for all functions that are used to work with mysql advisory locks
 func IsLockingFunc(node Expr) bool {
  switch p := node.(type) {
  case *FuncExpr:

go/vt/sqlparser/ast_rewriting.go

@@ -18,6 +18,7 @@ package sqlparser
 
 import (
  "strconv"
+ "vitess.io/vitess/go/vt/vtgate/boost"
 
  querypb "vitess.io/vitess/go/vt/proto/query"
  vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
@@ -31,7 +32,8 @@ import (
 // RewriteASTResult contains the rewritten ast and meta information about it
 type RewriteASTResult struct {
  *BindVarNeeds
- AST Statement // The rewritten AST
+ AST Statement // The rewritten AST
+ Columns boost.Columns
 }
 
 // ReservedVars keeps track of the bind variable names that have already been used
@@ -145,17 +147,21 @@ func NewReservedVars(prefix string, known BindVars) *ReservedVars {
 
 // PrepareAST will normalize the query
 func PrepareAST(in Statement, reservedVars *ReservedVars, bindVars map[string]*querypb.BindVariable, parameterize bool, keyspace string) (*RewriteASTResult, error) {
+ var boostColumns boost.Columns
+ var err error
+
  if parameterize {
- err := Normalize(in, reservedVars, bindVars)
+ boostColumns, err = Normalize(in, reservedVars, bindVars)
  if err != nil {
  return nil, err
  }
  }
- return RewriteAST(in, keyspace)
+
+ return RewriteAST(in, keyspace, boostColumns)
 }
 
 // RewriteAST rewrites the whole AST, replacing function calls and adding column aliases to queries
-func RewriteAST(in Statement, keyspace string) (*RewriteASTResult, error) {
+func RewriteAST(in Statement, keyspace string, columns map[string]string) (*RewriteASTResult, error) {
  er := newExpressionRewriter(keyspace)
  er.shouldRewriteDatabaseFunc = shouldRewriteDatabaseFunc(in)
  setRewriter := &setNormalizer{}
@@ -172,6 +178,7 @@ func RewriteAST(in Statement, keyspace string) (*RewriteASTResult, error) {
  r := &RewriteASTResult{
  AST: out,
  BindVarNeeds: er.bindVars,
+ Columns: columns,
  }
  return r, nil
 }

go/vt/sqlparser/normalizer.go

@@ -18,8 +18,8 @@ package sqlparser
 
 import (
  "vitess.io/vitess/go/sqltypes"
-
  querypb "vitess.io/vitess/go/vt/proto/query"
+ "vitess.io/vitess/go/vt/vtgate/boost"
 )
 
 // BindVars is a set of reserved bind variables from a SQL statement
@@ -32,24 +32,28 @@ type BindVars map[string]struct{}
 // Within Select constructs, bind vars are deduped. This allows
 // us to identify vindex equality. Otherwise, every value is
 // treated as distinct.
-func Normalize(stmt Statement, reserved *ReservedVars, bindVars map[string]*querypb.BindVariable) error {
+
+func Normalize(stmt Statement, reserved *ReservedVars, bindVars map[string]*querypb.BindVariable) (boost.Columns, error) {
  nz := newNormalizer(reserved, bindVars)
  _ = Rewrite(stmt, nz.WalkStatement, nil)
- return nz.err
+ return nz.columns, nz.err
 }
 
 type normalizer struct {
- bindVars map[string]*querypb.BindVariable
- reserved *ReservedVars
- vals map[string]string
- err error
+ bindVars map[string]*querypb.BindVariable
+ reserved *ReservedVars
+ vals map[string]string
+ columns boost.Columns
+ currentColumn string
+ err error
 }
 
 func newNormalizer(reserved *ReservedVars, bindVars map[string]*querypb.BindVariable) *normalizer {
  return &normalizer{
  bindVars: bindVars,
  reserved: reserved,
  vals: make(map[string]string),
+ columns: make(map[string]string),
  }
 }
 
@@ -86,7 +90,10 @@ func (nz *normalizer) WalkSelect(cursor *Cursor) bool {
  nz.convertLiteralDedup(node, cursor)
  case *ComparisonExpr:
  nz.convertComparison(node)
- case *ColName, TableName:
+ case *ColName:
+ nz.currentColumn = node.Name.String()
+ return false
+ case *TableName:
  // Common node types that never contain Literals or ListArgs but create a lot of object
  // allocations.
  return false
@@ -134,6 +141,11 @@ func (nz *normalizer) convertLiteralDedup(node *Literal, cursor *Cursor) {
  nz.bindVars[bvname] = bval
  }
 
+ // <BOOST> Store the column to bind var mapping.
+ if nz.currentColumn != "" && !ok {
+ nz.columns[nz.currentColumn] = bvname
+ }
+
  // Modify the AST node to a bindvar.
  cursor.Replace(NewArgument(bvname))
 }
@@ -148,6 +160,11 @@ func (nz *normalizer) convertLiteral(node *Literal, cursor *Cursor) {
  bvname := nz.reserved.nextUnusedVar()
  nz.bindVars[bvname] = bval
 
+ // <BOOST> Store the column to bind var mapping.
+ if nz.currentColumn != "" {
+ nz.columns[nz.currentColumn] = bvname
+ }
+
  cursor.Replace(NewArgument(bvname))
 }
 

go/vt/sqlparser/redact_query.go

@@ -28,7 +28,7 @@ func RedactSQLQuery(sql string) (string, error) {
  return "", err
  }
 
- err = Normalize(stmt, NewReservedVars("redacted", reservedVars), bv)
+ _, err = Normalize(stmt, NewReservedVars("redacted", reservedVars), bv)
  if err != nil {
  return "", err
  }

go/vt/sqlparser/utils.go

@@ -40,7 +40,7 @@ func QueryMatchesTemplates(query string, queryTemplates []string) (match bool, e
  if err != nil {
  return "", err
  }
- err = Normalize(stmt, NewReservedVars("", reservedVars), bv)
+ _, err = Normalize(stmt, NewReservedVars("", reservedVars), bv)
  if err != nil {
  return "", err
  }

go/vt/vtgate/boost/boost.go

@@ -0,0 +1,8 @@
+package boost
+
+type Columns map[string]string
+
+type PlanConfig struct {
+ IsBoosted bool
+ BoostColumns Columns
+}

go/vt/vtgate/engine/primitive.go

@@ -21,6 +21,7 @@ import (
  "encoding/json"
  "sync/atomic"
  "time"
+ "vitess.io/vitess/go/vt/vtgate/boost"
 
  "golang.org/x/sync/errgroup"
 
@@ -161,11 +162,12 @@ type (
  // each node does its part by combining the results of the
  // sub-nodes.
  Plan struct {
- Type sqlparser.StatementType // The type of query we have
- Original string // Original is the original query.
- Instructions Primitive // Instructions contains the instructions needed to fulfil the query.
- BindVarNeeds *sqlparser.BindVarNeeds // Stores BindVars needed to be provided as part of expression rewriting
- Warnings []*querypb.QueryWarning // Warnings that need to be yielded every time this query runs
+ Type sqlparser.StatementType // The type of query we have
+ Original string // Original is the original query.
+ Instructions Primitive // Instructions contains the instructions needed to fulfil the query.
+ BindVarNeeds *sqlparser.BindVarNeeds // Stores BindVars needed to be provided as part of expression rewriting
+ Warnings []*querypb.QueryWarning // Warnings that need to be yielded every time this query runs
+ BoostPlanConfig *boost.PlanConfig
 
  ExecCount uint64 // Count of times this plan was executed
  ExecTime uint64 // Total execution time
@@ -249,7 +251,7 @@ func Exists(m Match, p Primitive) bool {
  return Find(m, p) != nil
 }
 
-//MarshalJSON serializes the plan into a JSON representation.
+// MarshalJSON serializes the plan into a JSON representation.
 func (p *Plan) MarshalJSON() ([]byte, error) {
  var instructions *PrimitiveDescription
  if p.Instructions != nil {

go/vt/vtgate/executor.go

@@ -28,6 +28,7 @@ import (
  "strings"
  "sync"
  "time"
+ "vitess.io/vitess/go/vt/vtgate/boost"
 
  "vitess.io/vitess/go/acl"
  "vitess.io/vitess/go/cache"
@@ -398,7 +399,7 @@ func (e *Executor) handleCommit(ctx context.Context, safeSession *SafeSession, l
  return &sqltypes.Result{}, err
 }
 
-//Commit commits the existing transactions
+// Commit commits the existing transactions
 func (e *Executor) Commit(ctx context.Context, safeSession *SafeSession) error {
  return e.txConn.Commit(ctx, safeSession)
 }
@@ -1201,6 +1202,7 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
  }
  ignoreMaxMemoryRows := sqlparser.IgnoreMaxMaxMemoryRowsDirective(stmt)
  vcursor.SetIgnoreMaxMemoryRows(ignoreMaxMemoryRows)
+ var boostPlanConfig *boost.PlanConfig
 
  // Normalize if possible and retry.
  if (e.normalize && sqlparser.CanNormalize(stmt)) || sqlparser.MustRewriteAST(stmt) {
@@ -1212,6 +1214,7 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
  statement = result.AST
  bindVarNeeds = result.BindVarNeeds
  query = sqlparser.String(statement)
+ boostPlanConfig = configForBoost(result.Columns, "<TODO>")
  }
 
  if logStats != nil {
@@ -1224,7 +1227,7 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
  return plan.(*engine.Plan), nil
  }
 
- plan, err := planbuilder.BuildFromStmt(query, statement, reservedVars, vcursor, bindVarNeeds, *enableOnlineDDL, *enableDirectDDL)
+ plan, err := planbuilder.BuildFromStmt(query, statement, reservedVars, vcursor, bindVarNeeds, *enableOnlineDDL, *enableDirectDDL, boostPlanConfig)
  if err != nil {
  return nil, err
  }
@@ -1235,9 +1238,41 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
  if !skipQueryPlanCache && !sqlparser.SkipQueryPlanCacheDirective(statement) && sqlparser.CachePlan(statement) {
  e.plans.Set(planKey, plan)
  }
+
  return plan, nil
 }
 
+// TODO
+func configForBoost(columns boost.Columns, table string) *boost.PlanConfig {
+ configColumns := map[string]string{
+ "user_id": "1337",
+ }
+
+ //todo compare sets for ordering
+ if !keysMatch(columns, configColumns) {
+ return &boost.PlanConfig{}
+ }
+
+ return &boost.PlanConfig{
+ IsBoosted: true,
+ BoostColumns: columns,
+ }
+}
+
+func keysMatch(map1, map2 map[string]string) bool {
+ if len(map1) != len(map2) {
+ return false
+ }
+
+ for k := range map1 {
+ if _, exists := map2[k]; !exists {
+ return false
+ }
+ }
+
+ return true
+}
+
 // skipQueryPlanCache extracts SkipQueryPlanCache from session
 func skipQueryPlanCache(safeSession *SafeSession) bool {
  if safeSession == nil || safeSession.Options == nil {

go/vt/vtgate/plan_execute.go

@@ -120,7 +120,14 @@ func (e *Executor) newExecute(ctx context.Context, safeSession *SafeSession, sql
  e.executePlan(ctx, plan, vcursor, bindVars, execStart))
  }
 
- return e.executePlan(ctx, plan, vcursor, bindVars, execStart)(logStats, safeSession)
+ // Check if boosted and hit Redis
+ // plan.BoostPlanConfig.IsBoosted == true
+
+ statementTypeResult, sqlResult, err := e.executePlan(ctx, plan, vcursor, bindVars, execStart)(logStats, safeSession)
+
+ // Maybe store in Redis here if boosted, but cache miss
+
+ return statementTypeResult, sqlResult, err
 }
 
 func (e *Executor) startTxIfNecessary(ctx context.Context, safeSession *SafeSession) error {

go/vt/vtgate/planbuilder/builder.go

@@ -19,6 +19,7 @@ package planbuilder
 import (
  "errors"
  "sort"
+ "vitess.io/vitess/go/vt/vtgate/boost"
 
  "vitess.io/vitess/go/sqltypes"
  querypb "vitess.io/vitess/go/vt/proto/query"
@@ -93,29 +94,30 @@ func TestBuilder(query string, vschema ContextVSchema) (*engine.Plan, error) {
  if err != nil {
  return nil, err
  }
- result, err := sqlparser.RewriteAST(stmt, "")
+ result, err := sqlparser.RewriteAST(stmt, "", make(map[string]string))
  if err != nil {
  return nil, err
  }
 
  reservedVars := sqlparser.NewReservedVars("vtg", reserved)
- return BuildFromStmt(query, result.AST, reservedVars, vschema, result.BindVarNeeds, true, true)
+ return BuildFromStmt(query, result.AST, reservedVars, vschema, result.BindVarNeeds, true, true, nil)
 }
 
 // ErrPlanNotSupported is an error for plan building not supported
 var ErrPlanNotSupported = errors.New("plan building not supported")
 
 // BuildFromStmt builds a plan based on the AST provided.
-func BuildFromStmt(query string, stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema ContextVSchema, bindVarNeeds *sqlparser.BindVarNeeds, enableOnlineDDL, enableDirectDDL bool) (*engine.Plan, error) {
+func BuildFromStmt(query string, stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema ContextVSchema, bindVarNeeds *sqlparser.BindVarNeeds, enableOnlineDDL, enableDirectDDL bool, config *boost.PlanConfig) (*engine.Plan, error) {
  instruction, err := createInstructionFor(query, stmt, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
  if err != nil {
  return nil, err
  }
  plan := &engine.Plan{
- Type: sqlparser.ASTToStatementType(stmt),
- Original: query,
- Instructions: instruction,
- BindVarNeeds: bindVarNeeds,
+ Type: sqlparser.ASTToStatementType(stmt),
+ Original: query,
+ Instructions: instruction,
+ BindVarNeeds: bindVarNeeds,
+ BoostPlanConfig: config,
  }
  return plan, nil
 }