refactoring

Signed-off-by: Andres Taylor <[email protected]>
vitessio · Aug 28, 2024 · aa788bd · aa788bd
1 parent b0c5192
commit aa788bd
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 48 deletions.
diff --git a/go/vt/vtctl/workflow/stream_migrator.go b/go/vt/vtctl/workflow/stream_migrator.go
@@ -1130,15 +1130,15 @@ func (sm *StreamMigrator) templatizeRule(ctx context.Context, rule *binlogdatapb
 	case rule.Filter == vreplication.ExcludeStr:
 		return StreamTypeUnknown, fmt.Errorf("unexpected rule in vreplication: %v", rule)
 	default:
-		if err := sm.templatizeKeyRange(ctx, rule); err != nil {
+		if err := sm.templatizeKeyRange(rule); err != nil {
 			return StreamTypeUnknown, err
 		}
 
 		return StreamTypeSharded, nil
 	}
 }
 
-func (sm *StreamMigrator) templatizeKeyRange(ctx context.Context, rule *binlogdatapb.Rule) error {
+func (sm *StreamMigrator) templatizeKeyRange(rule *binlogdatapb.Rule) error {
 	statement, err := sm.parser.Parse(rule.Filter)
 	if err != nil {
 		return err
@@ -1149,12 +1149,7 @@ func (sm *StreamMigrator) templatizeKeyRange(ctx context.Context, rule *binlogda
 		return fmt.Errorf("unexpected query: %v", rule.Filter)
 	}
 
-	var expr sqlparser.Expr
-	if sel.Where != nil {
-		expr = sel.Where.Expr
-	}
-
-	exprs := sqlparser.SplitAndExpression(nil, expr)
+	exprs := sqlparser.SplitAndExpression(nil, sel.GetWherePredicate())
 	for _, subexpr := range exprs {
 		funcExpr, ok := subexpr.(*sqlparser.FuncExpr)
 		if !ok || !funcExpr.Name.EqualString("in_keyrange") {

diff --git a/go/vt/vtctl/workflow/vexec/query_planner.go b/go/vt/vtctl/workflow/vexec/query_planner.go
@@ -181,7 +181,7 @@ func (planner *VReplicationQueryPlanner) planDelete(del *sqlparser.Delete) (*Fix
 		)
 	}
 
-	del.Where = addDefaultWheres(planner, del.Where)
+	addDefaultWheres(planner, del)
 
 	buf := sqlparser.NewTrackedBuffer(nil)
 	buf.Myprintf("%v", del)
@@ -194,7 +194,7 @@ func (planner *VReplicationQueryPlanner) planDelete(del *sqlparser.Delete) (*Fix
 }
 
 func (planner *VReplicationQueryPlanner) planSelect(sel *sqlparser.Select) (*FixedQueryPlan, error) {
-	sel.Where = addDefaultWheres(planner, sel.Where)
+	addDefaultWheres(planner, sel)
 
 	buf := sqlparser.NewTrackedBuffer(nil)
 	buf.Myprintf("%v", sel)
@@ -230,7 +230,7 @@ func (planner *VReplicationQueryPlanner) planUpdate(upd *sqlparser.Update) (*Fix
 		}
 	}
 
-	upd.Where = addDefaultWheres(planner, upd.Where)
+	addDefaultWheres(planner, upd)
 
 	buf := sqlparser.NewTrackedBuffer(nil)
 	buf.Myprintf("%v", upd)
@@ -289,8 +289,7 @@ func (planner *VReplicationLogQueryPlanner) QueryParams() QueryParams {
 }
 
 func (planner *VReplicationLogQueryPlanner) planSelect(sel *sqlparser.Select) (QueryPlan, error) {
-	where := sel.Where
-	cols := extractWhereComparisonColumns(where)
+	cols := extractWhereComparisonColumns(sel.GetWherePredicate())
 	hasVReplIDCol := false
 
 	for _, col := range cols {
@@ -313,10 +312,6 @@ func (planner *VReplicationLogQueryPlanner) planSelect(sel *sqlparser.Select) (Q
 	// streamIDs.
 	queriesByTarget := make(map[string]*sqlparser.ParsedQuery, len(planner.tabletStreamIDs))
 	for target, streamIDs := range planner.tabletStreamIDs {
-		targetWhere := &sqlparser.Where{
-			Type: sqlparser.WhereClause,
-		}
-
 		var expr sqlparser.Expr
 		switch len(streamIDs) {
 		case 0: // WHERE vreplication_log.vrepl_id IN () => WHERE 1 != 1
@@ -349,15 +344,7 @@ func (planner *VReplicationLogQueryPlanner) planSelect(sel *sqlparser.Select) (Q
 				Right: tuple,
 			}
 		}
-
-		switch where {
-		case nil:
-			targetWhere.Expr = expr
-		default:
-			targetWhere.Expr = sqlparser.CreateAndExpr(expr, where.Expr)
-		}
-
-		sel.Where = targetWhere
+		sel.AddWhere(expr)
 
 		buf := sqlparser.NewTrackedBuffer(nil)
 		buf.Myprintf("%v", sel)
@@ -371,8 +358,8 @@ func (planner *VReplicationLogQueryPlanner) planSelect(sel *sqlparser.Select) (Q
 	}, nil
 }
 
-func addDefaultWheres(planner QueryPlanner, where *sqlparser.Where) *sqlparser.Where {
-	cols := extractWhereComparisonColumns(where)
+func addDefaultWheres(planner QueryPlanner, stmt sqlparser.WhereAble) {
+	cols := extractWhereComparisonColumns(stmt.GetWherePredicate())
 
 	params := planner.QueryParams()
 	hasDBNameCol := false
@@ -387,8 +374,6 @@ func addDefaultWheres(planner QueryPlanner, where *sqlparser.Where) *sqlparser.W
 		}
 	}
 
-	newWhere := where
-
 	if !hasDBNameCol {
 		expr := &sqlparser.ComparisonExpr{
 			Left: &sqlparser.ColName{
@@ -398,15 +383,7 @@ func addDefaultWheres(planner QueryPlanner, where *sqlparser.Where) *sqlparser.W
 			Right:    sqlparser.NewStrLiteral(params.DBName),
 		}
 
-		switch newWhere {
-		case nil:
-			newWhere = &sqlparser.Where{
-				Type: sqlparser.WhereClause,
-				Expr: expr,
-			}
-		default:
-			newWhere.Expr = sqlparser.CreateAndExpr(newWhere.Expr, expr)
-		}
+		stmt.AddWhere(expr)
 	}
 
 	if !hasWorkflowCol && params.Workflow != "" {
@@ -417,23 +394,20 @@ func addDefaultWheres(planner QueryPlanner, where *sqlparser.Where) *sqlparser.W
 			Operator: sqlparser.EqualOp,
 			Right:    sqlparser.NewStrLiteral(params.Workflow),
 		}
-
-		newWhere.Expr = sqlparser.CreateAndExpr(newWhere.Expr, expr)
+		stmt.AddWhere(expr)
 	}
-
-	return newWhere
 }
 
 // extractWhereComparisonColumns extracts the column names used in AND-ed
 // comparison expressions in a where clause, given the following assumptions:
 // - (1) The column name is always the left-hand side of the comparison.
 // - (2) There are no compound expressions within the where clause involving OR.
-func extractWhereComparisonColumns(where *sqlparser.Where) []string {
+func extractWhereComparisonColumns(where sqlparser.Expr) []string {
 	if where == nil {
 		return nil
 	}
 
-	exprs := sqlparser.SplitAndExpression(nil, where.Expr)
+	exprs := sqlparser.SplitAndExpression(nil, where)
 	cols := make([]string, 0, len(exprs))
 
 	for _, expr := range exprs {

diff --git a/go/vt/vtgate/planbuilder/operators/ast_to_op.go b/go/vt/vtgate/planbuilder/operators/ast_to_op.go
@@ -50,8 +50,8 @@ func translateQueryToOp(ctx *plancontext.PlanningContext, selStmt sqlparser.Stat
 func createOperatorFromSelect(ctx *plancontext.PlanningContext, sel *sqlparser.Select) Operator {
 	op := crossJoin(ctx, sel.From)
 
-	if sel.Where != nil {
-		op = addWherePredicates(ctx, sel.Where.Expr, op)
+	if expr := sel.GetWherePredicate(); expr != nil {
+		op = addWherePredicates(ctx, expr, op)
 	}
 
 	if sel.Comments != nil || sel.Lock != sqlparser.NoLock {

diff --git a/go/vt/vtgate/planbuilder/operators/phases.go b/go/vt/vtgate/planbuilder/operators/phases.go
@@ -193,7 +193,7 @@ func createDMLWithInput(ctx *plancontext.PlanningContext, op, src Operator, in *
 
 	if in.OwnedVindexQuery != nil {
 		in.OwnedVindexQuery.From = sqlparser.TableExprs{targetQT.Alias}
-		in.OwnedVindexQuery.Where = sqlparser.NewWhere(sqlparser.WhereClause, compExpr)
+		in.OwnedVindexQuery.AddWhere(compExpr)
 		in.OwnedVindexQuery.OrderBy = nil
 		in.OwnedVindexQuery.Limit = nil
 	}