From 0122d504822f98fd66c426c020faeb59b1f230a8 Mon Sep 17 00:00:00 2001
From: Manan Gupta <35839558+GuptaManan100@users.noreply.github.com>
Date: Wed, 17 Jul 2024 11:41:56 +0530
Subject: [PATCH] [release-20.0]  Fix panic in schema tracker in presence of
 keyspace routing rules (#16383) (#16405)

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 go/vt/vtgate/vschema_manager.go      |  8 +++---
 go/vt/vtgate/vschema_manager_test.go | 43 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/go/vt/vtgate/vschema_manager.go b/go/vt/vtgate/vschema_manager.go
index dbac5589ce8..8afe707bd9e 100644
--- a/go/vt/vtgate/vschema_manager.go
+++ b/go/vt/vtgate/vschema_manager.go
@@ -229,9 +229,9 @@ func (vm *VSchemaManager) updateTableInfo(vschema *vindexes.VSchema, ks *vindexe
 	// Now that we have ensured that all the tables are created, we can start populating the foreign keys
 	// in the tables.
 	for tblName, tblInfo := range m {
-		rTbl, err := vschema.FindRoutedTable(ksName, tblName, topodatapb.TabletType_PRIMARY)
-		if err != nil {
-			log.Errorf("error finding routed table %s: %v", tblName, err)
+		rTbl := ks.Tables[tblName]
+		if rTbl == nil {
+			log.Errorf("unable to find table %s in %s", tblName, ksName)
 			continue
 		}
 		for _, fkDef := range tblInfo.ForeignKeys {
@@ -240,7 +240,7 @@ func (vm *VSchemaManager) updateTableInfo(vschema *vindexes.VSchema, ks *vindexe
 				continue
 			}
 			parentTbl, err := vschema.FindRoutedTable(ksName, fkDef.ReferenceDefinition.ReferencedTable.Name.String(), topodatapb.TabletType_PRIMARY)
-			if err != nil {
+			if err != nil || parentTbl == nil {
 				log.Errorf("error finding parent table %s: %v", fkDef.ReferenceDefinition.ReferencedTable.Name.String(), err)
 				continue
 			}
diff --git a/go/vt/vtgate/vschema_manager_test.go b/go/vt/vtgate/vschema_manager_test.go
index 32f83f0021a..8dfb889df0d 100644
--- a/go/vt/vtgate/vschema_manager_test.go
+++ b/go/vt/vtgate/vschema_manager_test.go
@@ -336,6 +336,49 @@ func TestVSchemaUpdate(t *testing.T) {
 	}
 }
 
+// TestKeyspaceRoutingRules tests that the vschema manager doens't panic in the presence of keyspace routing rules.
+func TestKeyspaceRoutingRules(t *testing.T) {
+	cols1 := []vindexes.Column{{
+		Name: sqlparser.NewIdentifierCI("id"),
+		Type: querypb.Type_INT64,
+	}}
+	// Create a vschema manager with a fake vschema that returns a table with a column and a primary key.
+	vm := &VSchemaManager{}
+	vm.schema = &fakeSchema{t: map[string]*vindexes.TableInfo{
+		"t1": {
+			Columns: cols1,
+			Indexes: []*sqlparser.IndexDefinition{
+				{
+					Info: &sqlparser.IndexInfo{Type: sqlparser.IndexTypePrimary},
+					Columns: []*sqlparser.IndexColumn{
+						{
+							Column: sqlparser.NewIdentifierCI("id"),
+						},
+					},
+				},
+			},
+		},
+	}}
+	// Define a vschema that has a keyspace routing rule.
+	vs := &vindexes.VSchema{
+		Keyspaces: map[string]*vindexes.KeyspaceSchema{
+			"ks": {
+				Tables:   map[string]*vindexes.Table{},
+				Keyspace: &vindexes.Keyspace{Name: "ks", Sharded: true},
+			},
+			"ks2": {
+				Tables:   map[string]*vindexes.Table{},
+				Keyspace: &vindexes.Keyspace{Name: "ks2", Sharded: true},
+			},
+		},
+		KeyspaceRoutingRules: map[string]string{
+			"ks": "ks2",
+		},
+	}
+	// Ensure that updating the vschema manager from the vschema doesn't cause a panic.
+	vm.updateFromSchema(vs)
+}
+
 func TestRebuildVSchema(t *testing.T) {
 	cols1 := []vindexes.Column{{
 		Name: sqlparser.NewIdentifierCI("id"),