From e26e54a4aaed7f9d7e4c80d37f9ac1851b18338e Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Sat, 26 Oct 2024 01:22:48 +0530 Subject: [PATCH] fix: check for operator type before route to know a complex plan and fail accordingly Signed-off-by: Harshit Gangal --- .../planbuilder/operators/plan_query.go | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/plan_query.go b/go/vt/vtgate/planbuilder/operators/plan_query.go index baa9a7883e2..bf93683677c 100644 --- a/go/vt/vtgate/planbuilder/operators/plan_query.go +++ b/go/vt/vtgate/planbuilder/operators/plan_query.go @@ -43,6 +43,7 @@ import ( "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" + "vitess.io/vitess/go/vt/vtgate/semantics" ) type ( @@ -73,13 +74,28 @@ func PlanQuery(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (resu checkValid(op) op = planQuery(ctx, op) - _, isRoute := op.(*Route) - if !isRoute && ctx.SemTable.NotSingleRouteErr != nil { - // If we got here, we don't have a single shard plan - return nil, ctx.SemTable.NotSingleRouteErr + err = shouldErrorOnComplexPlan(ctx, op) + return op, err +} + +func shouldErrorOnComplexPlan(ctx *plancontext.PlanningContext, op Operator) error { + complexPlan := false + visitF := func(op Operator, _ semantics.TableSet, _ bool) (Operator, *ApplyResult) { + switch op.(type) { + case *Limit, *Route: + default: + complexPlan = true + } + return op, NoRewrite } - return op, err + TopDown(op, TableID, visitF, stopAtRoute) + + if complexPlan && ctx.SemTable.NotSingleRouteErr != nil { + // If we got here, we don't have a single shard plan + return ctx.SemTable.NotSingleRouteErr + } + return nil } func PanicHandler(err *error) {