From f20183da065b058315f5d2ed75cfb4859b1c8061 Mon Sep 17 00:00:00 2001
From: "vitess-bot[bot]" <108069721+vitess-bot[bot]@users.noreply.github.com>
Date: Mon, 28 Oct 2024 16:53:35 +0530
Subject: [PATCH] Delegate Column Availability Checks to MySQL for Single-Route
 Queries (#17077)

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
Signed-off-by: Andres Taylor <andres@planetscale.com>
Co-authored-by: Andres Taylor <andres@planetscale.com>
---
 .../planbuilder/operators/plan_query.go       | 33 ++++++++++++++++---
 .../planbuilder/testdata/select_cases.json    | 29 ++++++++++++++++
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/plan_query.go b/go/vt/vtgate/planbuilder/operators/plan_query.go
index d5794a1bcf1..4a9f356d2bd 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 (
@@ -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) {
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
index bf1e1950760..cbe14c48b5a 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.json
+++ b/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",