boost: move boost config to plan itself

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,17 +32,18 @@ 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
-	columns       map[string]string
+	columns       boost.Columns
 	currentColumn string
 	err           error
 }

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"
@@ -1203,6 +1204,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) {
@@ -1214,6 +1216,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, "x")
 	}
 
 	if logStats != nil {
@@ -1226,24 +1229,7 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
 		return plan.(*engine.Plan), nil
 	}
 
-	isBoosted := true
-
-	if isBoosted {
-		plan, err := planbuilder.BuildBoost(query, statement)
-		if err != nil {
-			return nil, err
-		}
-
-		plan.Warnings = vcursor.warnings
-		vcursor.warnings = nil
-
-		if !skipQueryPlanCache && !sqlparser.SkipQueryPlanCacheDirective(statement) && sqlparser.CachePlan(statement) {
-			e.plans.Set(planKey, plan)
-		}
-		return 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
 	}
@@ -1254,9 +1240,40 @@ 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
 }
 
+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/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,40 +94,41 @@ 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
 }
 
-func BuildBoost(query string, stmt sqlparser.Statement) (*engine.Plan, error) {
-	return &engine.Plan{
-		Type:         sqlparser.StmtBoost,
-		Instructions: &engine.EmptyPrimitive{},
-		BindVarNeeds: &sqlparser.BindVarNeeds{},
-	}, nil
-}
+//func BuildBoost(query string, stmt sqlparser.Statement) (*engine.Plan, error) {
+//	return &engine.Plan{
+//		Type:         sqlparser.StmtBoost,
+//		Instructions: &engine.EmptyPrimitive{},
+//		BindVarNeeds: &sqlparser.BindVarNeeds{},
+//	}, nil
+//}
 
 func getConfiguredPlanner(vschema ContextVSchema) (selectPlanner, error) {
 	switch vschema.Planner() {