boost: draft for vtgate boost query checks (#8)

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
 }