Delegate Column Availability Checks to MySQL for Single-Route Queries (…

…#17077)

Signed-off-by: Harshit Gangal <[email protected]>
Signed-off-by: Andres Taylor <[email protected]>
Co-authored-by: Andres Taylor <[email protected]>

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 (
@@ -71,13 +72,35 @@ 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
+	if err := checkSingleRouteError(ctx, op); err != nil {
+		return nil, err
 	}
 
-	return op, err
+	return op, 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 Operator) error {
+	if ctx.SemTable.NotSingleRouteErr == nil {
+		return nil
+	}
+	routes := 0
+	visitF := func(op Operator, _ semantics.TableSet, _ bool) (Operator, *ApplyResult) {
+		switch op.(type) {
+		case *Route:
+			routes++
+		}
+		return op, NoRewrite
+	}
+
+	// we'll walk the tree and count the number of routes
+	TopDown(op, TableID, visitF, stopAtRoute)
+
+	if routes <= 1 {
+		return nil
+	}
+
+	return ctx.SemTable.NotSingleRouteErr
 }
 
 func PanicHandler(err *error) {

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": "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 20",
+            "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",