From 47d9053e6e57e220531ce74ba8d838430570f572 Mon Sep 17 00:00:00 2001
From: Harshit Gangal <harshit@planetscale.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>
---
 .../vtgate/planbuilder/operators/operator.go  | 31 ++++++++++++++++---
 .../planbuilder/testdata/select_cases.json    | 29 +++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/operator.go b/go/vt/vtgate/planbuilder/operators/operator.go
index a1cfbfd0cc0..32c10e7128f 100644
--- a/go/vt/vtgate/planbuilder/operators/operator.go
+++ b/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
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
index 2412fbaadcf..800ad170bc8 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": "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",