Merge branch 'release-18.0' of https://github.com/vitessio/vitess into release-18.0-github

Signed-off-by: Arthur Schreiber <[email protected]>

Author: arthurschreiber

go/vt/vtgate/planbuilder/operators/SQL_builder.go

@@ -564,14 +564,20 @@ func buildProjection(op *Projection, qb *queryBuilder) error {
 }
 
 func buildApplyJoin(op *ApplyJoin, qb *queryBuilder) error {
-	predicates := slice.Map(op.JoinPredicates, func(jc JoinColumn) sqlparser.Expr {
+	predicates, err := slice.MapWithError(op.JoinPredicates, func(jc JoinColumn) (sqlparser.Expr, error) {
 		// since we are adding these join predicates, we need to mark to broken up version (RHSExpr) of it as done
-		_ = qb.ctx.SkipJoinPredicates(jc.Original.Expr)
-		return jc.Original.Expr
+		err := qb.ctx.SkipJoinPredicates(jc.Original.Expr)
+		if err != nil {
+			return nil, err
+		}
+		return jc.Original.Expr, nil
 	})
+	if err != nil {
+		return err
+	}
 
 	pred := sqlparser.AndExpressions(predicates...)
-	err := buildQuery(op.LHS, qb)
+	err = buildQuery(op.LHS, qb)
 	if err != nil {
 		return err
 	}

go/vt/vtgate/planbuilder/operators/apply_join.go

@@ -153,15 +153,21 @@ func (aj *ApplyJoin) AddJoinPredicate(ctx *plancontext.PlanningContext, expr sql
 	if expr == nil {
 		return nil
 	}
-	col, err := BreakExpressionInLHSandRHS(ctx, expr, TableID(aj.LHS))
-	if err != nil {
-		return err
-	}
-	aj.JoinPredicates = append(aj.JoinPredicates, col)
-	rhs, err := aj.RHS.AddPredicate(ctx, col.RHSExpr)
-	if err != nil {
-		return err
+	rhs := aj.RHS
+	predicates := sqlparser.SplitAndExpression(nil, expr)
+	for _, pred := range predicates {
+		col, err := BreakExpressionInLHSandRHS(ctx, pred, TableID(aj.LHS))
+		if err != nil {
+			return err
+		}
+		aj.JoinPredicates = append(aj.JoinPredicates, col)
+		ctx.AddJoinPredicates(pred, col.RHSExpr)
+		rhs, err = rhs.AddPredicate(ctx, col.RHSExpr)
+		if err != nil {
+			return err
+		}
 	}
+
 	aj.RHS = rhs
 
 	return nil

go/vt/vtgate/planbuilder/operators/operator.go

@@ -44,6 +44,7 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/ops"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/rewrite"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+	"vitess.io/vitess/go/vt/vtgate/semantics"
 )
 
 type (
@@ -81,10 +82,8 @@ func PlanQuery(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (ops.
 		return nil, err
 	}
 
-	_, 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
+	if err := checkSingleRouteError(ctx, op); err != nil {
+		return nil, err
 	}
 
 	return op, err
@@ -161,3 +160,27 @@ func transformColumnsToSelectExprs(ctx *plancontext.PlanningContext, op ops.Oper
 	})
 	return selExprs, nil
 }
+
+// checkSingleRouteError checks if the query has a NotSingleRouteErr and more than one route, and returns an error if it does
+func checkSingleRouteError(ctx *plancontext.PlanningContext, op ops.Operator) error {
+	if ctx.SemTable.NotSingleRouteErr == nil {
+		return nil
+	}
+	routes := 0
+	visitF := func(op ops.Operator, _ semantics.TableSet, _ bool) (ops.Operator, *rewrite.ApplyResult, error) {
+		switch op.(type) {
+		case *Route:
+			routes++
+		}
+		return op, rewrite.SameTree, nil
+	}
+
+	// we'll walk the tree and count the number of routes
+	_, _ = rewrite.TopDown(op, TableID, visitF, stopAtRoute)
+
+	if routes <= 1 {
+		return nil
+	}
+
+	return ctx.SemTable.NotSingleRouteErr
+}

go/vt/vtgate/planbuilder/testdata/from_cases.json

@@ -511,7 +511,7 @@
           "Sharded": true
         },
         "FieldQuery": "select 42 from `user` as u, user_extra as ue, music as m where 1 != 1",
-        "Query": "select 42 from `user` as u, user_extra as ue, music as m where u.id = ue.user_id and u.id = :m_user_id and (u.foo or :m_foo or ue.foo) and m.user_id = u.id and (u.foo or m.foo or ue.foo)",
+        "Query": "select 42 from `user` as u, user_extra as ue, music as m where u.id = ue.user_id and m.user_id = u.id and (u.foo or m.foo or ue.foo)",
         "Table": "`user`, music, user_extra"
       },
       "TablesUsed": [
@@ -4193,5 +4193,27 @@
         "zlookup_unique.t1"
       ]
     }
+  },
+  {
+    "comment": "lots of joins that can be merged should not confuse the planner",
+    "query": "SELECT `music`.id FROM `music` INNER JOIN music `music2` ON `music2`.`id` = `music`.`id` INNER JOIN music `music3` ON `music3`.`id` = `music2`.`id` INNER JOIN music `music4` ON `music4`.`id` = `music3`.`id` WHERE music.user_index = music4.user_index",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT `music`.id FROM `music` INNER JOIN music `music2` ON `music2`.`id` = `music`.`id` INNER JOIN music `music3` ON `music3`.`id` = `music2`.`id` INNER JOIN music `music4` ON `music4`.`id` = `music3`.`id` WHERE music.user_index = music4.user_index",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select music.id from music, music as music2, music as music3, music as music4 where 1 != 1",
+        "Query": "select music.id from music, music as music2, music as music3, music as music4 where music2.id = music.id and music4.id = music3.id and music3.id = music2.id and music.user_index = music4.user_index",
+        "Table": "music"
+      },
+      "TablesUsed": [
+        "user.music"
+      ]
+    }
   }
 ]

go/vt/vtgate/planbuilder/testdata/select_cases.json

@@ -65,6 +65,35 @@
       ]
     }
   },
+  {
+    "comment": "join on sharding column with limit - should be a simple scatter query and limit on top with non resolved columns",
+    "query": "select * from user u join user_metadata um on u.id = um.user_id where foo=41 limit 20",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select * from user u join user_metadata um on u.id = um.user_id where foo=41 limit 20",
+      "Instructions": {
+        "OperatorType": "Limit",
+        "Count": "INT64(20)",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "Scatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select * from `user` as u, user_metadata as um where 1 != 1",
+            "Query": "select * from `user` as u, user_metadata as um where foo = 41 and u.id = um.user_id limit :__upper_limit",
+            "Table": "`user`, user_metadata"
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user",
+        "user.user_metadata"
+      ]
+    }
+  },
   {
     "comment": "select with timeout directive sets QueryTimeout in the route",
     "query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ * from user",