From 964b18123e4a6005c343535d3b76f5ec8fe39fb7 Mon Sep 17 00:00:00 2001
From: "vitess-bot[bot]" <108069721+vitess-bot[bot]@users.noreply.github.com>
Date: Tue, 19 Mar 2024 11:33:07 +0530
Subject: [PATCH] fix: remove keyspace from column during query builder
 (#15514)

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
---
 .../planbuilder/operators/SQL_builder.go      |   9 +-
 .../planbuilder/testdata/filter_cases.json    | 104 ++++++++++++++++++
 2 files changed, 105 insertions(+), 8 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/SQL_builder.go b/go/vt/vtgate/planbuilder/operators/SQL_builder.go
index 998f849ba3c..1f4dbd68ce2 100644
--- a/go/vt/vtgate/planbuilder/operators/SQL_builder.go
+++ b/go/vt/vtgate/planbuilder/operators/SQL_builder.go
@@ -622,15 +622,8 @@ func buildDerivedSelect(op *Horizon, qb *queryBuilder, sel *sqlparser.Select) {
 
 func buildHorizon(op *Horizon, qb *queryBuilder) {
 	buildQuery(op.Source, qb)
-
 	stripDownQuery(op.Query, qb.asSelectStatement())
-
-	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
-		if aliasedExpr, ok := node.(sqlparser.SelectExpr); ok {
-			removeKeyspaceFromSelectExpr(aliasedExpr)
-		}
-		return true, nil
-	}, qb.stmt)
+	sqlparser.RemoveKeyspaceInCol(qb.stmt)
 }
 
 func mergeHaving(h1, h2 *sqlparser.Where) *sqlparser.Where {
diff --git a/go/vt/vtgate/planbuilder/testdata/filter_cases.json b/go/vt/vtgate/planbuilder/testdata/filter_cases.json
index d144da2441d..aee0bac3365 100644
--- a/go/vt/vtgate/planbuilder/testdata/filter_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/filter_cases.json
@@ -4585,5 +4585,109 @@
         "user.multicol_tbl"
       ]
     }
+  },
+  {
+    "comment": "order by with filter removing the keyspace from order by",
+    "query": "select col from user.user where id = 1 order by user.user.user_id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select col from user.user where id = 1 order by user.user.user_id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select col from `user` where 1 != 1",
+        "Query": "select col from `user` where id = 1 order by `user`.user_id asc",
+        "Table": "`user`",
+        "Values": [
+          "1"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "group by with filter removing the keyspace from order by",
+    "query": "select col from user.user where id = 1 group by user.user.user_id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select col from user.user where id = 1 group by user.user.user_id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select col from `user` where 1 != 1 group by `user`.user_id",
+        "Query": "select col from `user` where id = 1 group by `user`.user_id",
+        "Table": "`user`",
+        "Values": [
+          "1"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "order with authoritative table - removing keyspace from group by",
+    "query": "select * from user.authoritative where user_id = 5 order by user_id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select * from user.authoritative where user_id = 5 order by user_id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select user_id, col1, col2 from authoritative where 1 != 1",
+        "Query": "select user_id, col1, col2 from authoritative where user_id = 5 order by authoritative.user_id asc",
+        "Table": "authoritative",
+        "Values": [
+          "5"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
+  },
+  {
+    "comment": "group by and having with authoritative table - removing keyspace from having",
+    "query": "select * from user.authoritative where user_id = 5 group by user_id having count(user_id) = 6 order by user_id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select * from user.authoritative where user_id = 5 group by user_id having count(user_id) = 6 order by user_id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select user_id, col1, col2 from authoritative where 1 != 1 group by user_id",
+        "Query": "select user_id, col1, col2 from authoritative where user_id = 5 group by user_id having count(user_id) = 6 order by authoritative.user_id asc",
+        "Table": "authoritative",
+        "Values": [
+          "5"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
   }
 ]