From 5c08da6a5db5db5eb872608aae462f3519b38388 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Taylor?= <andres@planetscale.com>
Date: Wed, 30 Oct 2024 06:54:16 +0100
Subject: [PATCH] Bugfix for Panic on Joined Queries with Non-Authoritative
 Tables in Vitess 19.0 (#17103)

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../planbuilder/testdata/from_cases.json      | 25 +++++++++++++++++
 go/vt/vtgate/semantics/analyzer_test.go       |  1 +
 go/vt/vtgate/semantics/early_rewriter.go      | 10 ++++++-
 go/vt/vtgate/semantics/early_rewriter_test.go | 27 +++++++++++++++++++
 4 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/go/vt/vtgate/planbuilder/testdata/from_cases.json b/go/vt/vtgate/planbuilder/testdata/from_cases.json
index 18f9f376810..7ed42771579 100644
--- a/go/vt/vtgate/planbuilder/testdata/from_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/from_cases.json
@@ -1069,6 +1069,31 @@
       ]
     }
   },
+  {
+    "comment": "Ambiguous column is not a problem when we can merge and push down",
+    "query": "select foobar from user join music on user.id = music.user_id order by foobar",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select foobar from user join music on user.id = music.user_id order by foobar",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select foobar, weight_string(foobar) from `user`, music where 1 != 1",
+        "OrderBy": "(0|1) ASC",
+        "Query": "select foobar, weight_string(foobar) from `user`, music where `user`.id = music.user_id order by foobar asc",
+        "ResultColumns": 1,
+        "Table": "`user`, music"
+      },
+      "TablesUsed": [
+        "user.music",
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "index hints, make sure they are not stripped.",
     "query": "select user.col from user use index(a)",
diff --git a/go/vt/vtgate/semantics/analyzer_test.go b/go/vt/vtgate/semantics/analyzer_test.go
index e02d3c62917..a4271cbe859 100644
--- a/go/vt/vtgate/semantics/analyzer_test.go
+++ b/go/vt/vtgate/semantics/analyzer_test.go
@@ -1401,6 +1401,7 @@ func fakeSchemaInfo() *FakeSI {
 	si := &FakeSI{
 		Tables: map[string]*vindexes.Table{
 			"t":  {Name: sqlparser.NewIdentifierCS("t"), Keyspace: unsharded},
+			"t3": {Name: sqlparser.NewIdentifierCS("t3"), Keyspace: unsharded},
 			"t1": {Name: sqlparser.NewIdentifierCS("t1"), Columns: cols1, ColumnListAuthoritative: true, Keyspace: ks2},
 			"t2": {Name: sqlparser.NewIdentifierCS("t2"), Columns: cols2, ColumnListAuthoritative: true, Keyspace: ks3},
 		},
diff --git a/go/vt/vtgate/semantics/early_rewriter.go b/go/vt/vtgate/semantics/early_rewriter.go
index 8345c5fb8eb..2c731c4bfc8 100644
--- a/go/vt/vtgate/semantics/early_rewriter.go
+++ b/go/vt/vtgate/semantics/early_rewriter.go
@@ -668,7 +668,15 @@ func (r *earlyRewriter) fillInQualifiers(cursor *sqlparser.CopyOnWriteCursor) {
 	if !found {
 		panic("uh oh")
 	}
-	tbl := r.tables.Tables[ts.TableOffset()]
+	offset := ts.TableOffset()
+	if offset < 0 {
+		// this is a column that is not coming from a table - it's an alias introduced in a SELECT expression
+		// Example: select (1+1) as foo from bar order by foo
+		// we don't want to add a qualifier to foo here
+		cursor.Replace(sqlparser.NewColName(col.Name.String()))
+		return
+	}
+	tbl := r.tables.Tables[offset]
 	tblName, err := tbl.Name()
 	if err != nil {
 		panic(err)
diff --git a/go/vt/vtgate/semantics/early_rewriter_test.go b/go/vt/vtgate/semantics/early_rewriter_test.go
index e9576b5a3b7..adf26017d87 100644
--- a/go/vt/vtgate/semantics/early_rewriter_test.go
+++ b/go/vt/vtgate/semantics/early_rewriter_test.go
@@ -739,6 +739,33 @@ func TestOrderByColumnName(t *testing.T) {
 	}
 }
 
+func TestOrderByNotAllColumnsAuthoritative(t *testing.T) {
+	schemaInfo := fakeSchemaInfo()
+	cDB := "db"
+	tcases := []struct {
+		sql    string
+		expSQL string
+		deps   TableSet
+	}{{
+		sql:    "select foo from t3 join t order by foo",
+		expSQL: "select foo from t3 join t order by foo asc",
+		deps:   None,
+	}}
+	for _, tcase := range tcases {
+		t.Run(tcase.sql, func(t *testing.T) {
+			ast, err := sqlparser.NewTestParser().Parse(tcase.sql)
+			require.NoError(t, err)
+			selectStatement := ast.(sqlparser.SelectStatement)
+			semTable, err := Analyze(selectStatement, cDB, schemaInfo)
+
+			require.NoError(t, err)
+			assert.Equal(t, tcase.expSQL, sqlparser.String(selectStatement))
+			orderByExpr := selectStatement.GetOrderBy()[0].Expr
+			assert.Equal(t, tcase.deps, semTable.RecursiveDeps(orderByExpr))
+		})
+	}
+}
+
 func TestSemTableDependenciesAfterExpandStar(t *testing.T) {
 	schemaInfo := &FakeSI{Tables: map[string]*vindexes.Table{
 		"t1": {