refactor: avoid scary global state

Signed-off-by: Andres Taylor <andres@planetscale.com>
vitessio · May 16, 2024 · 6e08d76 · 6e08d76
1 parent 490bfa7
commit 6e08d76
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/go/vt/vtgate/planbuilder/operators/apply_join.go b/go/vt/vtgate/planbuilder/operators/apply_join.go
@@ -163,9 +163,7 @@ func (aj *ApplyJoin) AddJoinPredicate(ctx *plancontext.PlanningContext, expr sql
 		col := breakExpressionInLHSandRHSForApplyJoin(ctx, pred, TableID(aj.LHS))
 		aj.JoinPredicates.add(col)
 		ctx.AddJoinPredicates(pred, col.RHSExpr)
-		ctx.JoinPredInProgress = pred
 		rhs = rhs.AddPredicate(ctx, col.RHSExpr)
-		ctx.JoinPredInProgress = nil
 	}
 	aj.RHS = rhs
 }

diff --git a/go/vt/vtgate/planbuilder/plancontext/planning_context.go b/go/vt/vtgate/planbuilder/plancontext/planning_context.go
@@ -28,9 +28,6 @@ type PlanningContext struct {
 	SemTable     *semantics.SemTable
 	VSchema      VSchema
 
-	// JoinPredInProgress is used to track the current join predicate being processed.
-	JoinPredInProgress sqlparser.Expr
-
 	// joinPredicates maps each original join predicate (key) to a slice of
 	// variations of the RHS predicates (value). This map is used to handle
 	// different scenarios in join planning, where the RHS predicates are
@@ -194,8 +191,13 @@ func (ctx *PlanningContext) execOnJoinPredicateEqual(joinPred sqlparser.Expr, fn
 
 func (ctx *PlanningContext) RewriteDerivedTableExpression(expr sqlparser.Expr, tableInfo semantics.TableInfo) sqlparser.Expr {
 	modifiedExpr := semantics.RewriteDerivedTableExpression(expr, tableInfo)
-	if ctx.JoinPredInProgress != nil {
-		ctx.AddJoinPredicates(ctx.JoinPredInProgress, modifiedExpr)
+	for key, exprs := range ctx.joinPredicates {
+		for _, rhsExpr := range exprs {
+			if ctx.SemTable.EqualsExpr(expr, rhsExpr) {
+				ctx.joinPredicates[key] = append(ctx.joinPredicates[key], modifiedExpr)
+				return modifiedExpr
+			}
+		}
 	}
 	return modifiedExpr
 }